fix(incremental): await async incremental cleanup

Asyncio realization of the JS queue/computation cancellation rework:
stopping the incremental publisher cancels all pending incremental work
and awaits its settlement, stream sources are closed lazily and only on
abnormal stop, aborting rejects pending subsequent results with the abort
reason after the cleanup has settled, and concurrent consumer cancellation
is realized by cancelling the task that awaits the pending anext().

Replicates graphql/graphql-js@3ea5c10.
           graphql/graphql-js@87e2929
           graphql/graphql-js@ce68f97
           graphql/graphql-js@06e5bd0
           graphql/graphql-js@41a671c

Author: Cito

src/graphql/execution/execute.py

@@ -4,6 +4,7 @@
 
 from asyncio import (
     FIRST_COMPLETED,
+    Future,
     TimeoutError,  # only needed for Python < 3.11  # noqa: A004
     ensure_future,
     gather,
@@ -217,6 +218,7 @@ class Executor(IncrementalPublisherContext):
     abort_signal: AbortSignal | None
     errors: list[GraphQLError] | None
     cancellable_streams: set[CancellableStreamRecord] | None
+    pending_incremental_futures: set[Future[Any]]
     middleware_manager: MiddlewareManager | None
     error_propagation: bool
 
@@ -271,6 +273,7 @@ def __init__(  # noqa: PLR0913
         self.is_async_iterable = is_async_iterable or default_is_async_iterable
         self.errors = None
         self.cancellable_streams = None
+        self.pending_incremental_futures = set()
         self._relevant_sub_fields: dict[tuple, CollectedFields] = {}
         self._stream_usages: RefMap[FieldDetailsList, StreamUsage] = RefMap()
         self._execution_plans: RefMap[GroupedFieldSet, ExecutionPlan] = RefMap()
@@ -469,6 +472,11 @@ async def await_result() -> (
                         resolved = await self.with_abort_signal(graphql_wrapped_result)
                     except GraphQLError as error:
                         return ExecutionResult(None, with_error(self.errors, error))
+                    except Exception:
+                        # cancel incremental work started early and close the
+                        # stream sources before re-raising, e.g. the abort reason
+                        await self.cancel_incremental_work()
+                        raise
                     return self.build_data_response(
                         resolved.result, resolved.increments
                     )
@@ -975,6 +983,50 @@ def abort_error(self) -> Exception:
         msg = f"Unexpected error value: {inspect(reason)}"
         return TypeError(msg)
 
+    def box_incremental_result(
+        self, result: AwaitableOrValue[T]
+    ) -> BoxedAwaitableOrValue[T]:
+        """Box a possibly awaitable incremental result.
+
+        A pending result is registered so that it can be cancelled when the
+        incremental execution is stopped before it has settled.
+        """
+        boxed = BoxedAwaitableOrValue(result)
+        future = boxed.pending_future
+        if future is not None:
+            futures = self.pending_incremental_futures
+            futures.add(future)
+            future.add_done_callback(futures.discard)
+        return boxed
+
+    async def cancel_incremental_work(self) -> None:
+        """Cancel all pending incremental work and close the stream sources.
+
+        Cancels the still pending incremental execution tasks first and waits for
+        their cancellation to settle, so that no early execution continues and no
+        iteration is pending on the stream sources any more, then triggers and
+        awaits the early return of all remaining cancellable streams.
+        """
+        futures = self.pending_incremental_futures
+        if futures:
+            pending = list(futures)
+            for future in pending:
+                future.cancel()
+            await gather(*pending, return_exceptions=True)
+        cancellable_streams = self.cancellable_streams
+        if cancellable_streams:
+            early_returns = [
+                early_return
+                for early_return in (
+                    stream_record.early_return()
+                    for stream_record in cancellable_streams
+                )
+                if default_is_awaitable(early_return)
+            ]
+            cancellable_streams.clear()
+            if early_returns:
+                await gather(*early_returns, return_exceptions=True)
+
     def cancellable_iterable(self, iterable: AsyncIterable[T]) -> AsyncIterable[T]:
         """Wrap an async iterable so pending iteration is cancelled on abort.
 
@@ -1135,7 +1187,7 @@ async def complete_async_iterator_value(
                         )
                     else:
                         stream_record = CancellableStreamRecord(
-                            early_return(),
+                            early_return,
                             stream_item_queue,
                             path,
                             stream_usage.label,
@@ -1840,17 +1892,17 @@ async def execute_async(
                             return await result
                         return result  # type: ignore
 
-                    pending_group.result = BoxedAwaitableOrValue(execute_async())
+                    pending_group.result = self.box_incremental_result(execute_async())
                 else:
-                    pending_group.result = BoxedAwaitableOrValue(executor())
+                    pending_group.result = self.box_incremental_result(executor())
             else:
 
                 def execute_sync(
                     executor: Callable[
                         [], AwaitableOrValue[CompletedExecutionGroup]
                     ] = executor,
                 ) -> BoxedAwaitableOrValue[CompletedExecutionGroup]:
-                    return BoxedAwaitableOrValue(executor())
+                    return self.box_incremental_result(executor())
 
                 pending_group.result = execute_sync
 
@@ -1934,7 +1986,7 @@ def first_executor() -> StreamItemResult:
             initial_path = stream_path.add_key(initial_index)
 
             first_stream_item: BoxedAwaitableOrValue[StreamItemResult] = (
-                BoxedAwaitableOrValue(
+                self.box_incremental_result(
                     complete_stream_item(
                         initial_path,
                         initial_item,
@@ -1972,9 +2024,9 @@ def current_executor(
                     )
 
                 current_stream_item = (
-                    BoxedAwaitableOrValue(current_executor())
+                    self.box_incremental_result(current_executor())
                     if enable_early_execution
-                    else lambda executor=current_executor: BoxedAwaitableOrValue(
+                    else lambda executor=current_executor: self.box_incremental_result(
                         executor()
                     )
                 )
@@ -1992,9 +2044,9 @@ def current_executor(
             async def await_first_stream_item() -> StreamItemResult:
                 return first_executor()
 
-            append_stream_item(BoxedAwaitableOrValue(await_first_stream_item()))
+            append_stream_item(self.box_incremental_result(await_first_stream_item()))
         else:
-            append_stream_item(lambda: BoxedAwaitableOrValue(first_executor()))
+            append_stream_item(lambda: self.box_incremental_result(first_executor()))
 
         return stream_item_queue
 
@@ -2022,9 +2074,9 @@ def executor() -> AwaitableOrValue[StreamItemResult]:
 
         stream_item_queue: list[StreamItemRecord] = []
         stream_item_queue.append(
-            BoxedAwaitableOrValue(executor())
+            self.box_incremental_result(executor())
             if self.enable_early_execution
-            else lambda: BoxedAwaitableOrValue(executor())
+            else lambda: self.box_incremental_result(executor())
         )
 
         return stream_item_queue
@@ -2076,9 +2128,9 @@ def executor() -> AwaitableOrValue[StreamItemResult]:
             )
 
         stream_item_queue.append(
-            BoxedAwaitableOrValue(executor())
+            self.box_incremental_result(executor())
             if self.enable_early_execution
-            else lambda: BoxedAwaitableOrValue(executor())
+            else lambda: self.box_incremental_result(executor())
         )
 
         if self.is_awaitable(result):

src/graphql/execution/incremental_graph.py

@@ -6,6 +6,7 @@
     Future,
     Task,
     ensure_future,
+    gather,
     get_running_loop,
     isfuture,
     sleep,
@@ -108,7 +109,8 @@ def next_completed_batch(
     def abort(self) -> None:
         """Abort the incremental graph execution."""
         for resolve in self._next_queue:
-            resolve.set_result(None)  # pragma: no cover
+            if not resolve.cancelled():  # pragma: no cover
+                resolve.set_result(None)
 
     def has_next(self) -> bool:
         """Check if there are more results to process."""
@@ -161,10 +163,20 @@ def remove_stream(self, stream_record: StreamRecord) -> None:
         """Remove a stream record as no longer pending."""
         del self._root_nodes[stream_record]
 
-    def stop_incremental_data(self) -> None:
-        """Stop the delivery of incremental data."""
+    async def stop_incremental_data(self) -> None:
+        """Stop the delivery and execution of incremental data.
+
+        Cancels all pending requests for the next completed batch and all still
+        running incremental execution tasks, waiting until their cancellation
+        has settled.
+        """
         for future in self._next_queue:
-            future.cancel()  # pragma: no cover
+            future.cancel()
+        tasks = list(self._tasks)
+        if tasks:
+            for task in tasks:
+                task.cancel()
+            await gather(*tasks, return_exceptions=True)
 
     def _add_incremental_data_records(
         self,
@@ -321,12 +333,14 @@ async def _on_stream_items(self, stream_record: StreamRecord) -> None:
     def _enqueue(self, completed: IncrementalDataRecordResult) -> None:
         """Enqueue completed incremental data record result."""
         self._completed_queue.append(completed)
-        try:
-            future = self._next_queue.pop(0)
-        except IndexError:
-            pass
-        else:
+        next_queue = self._next_queue
+        while next_queue:
+            future = next_queue.pop(0)
+            if future.cancelled():  # pragma: no cover
+                # defensive guard against a race with a stopping consumer
+                continue
             future.set_result(self.current_completed_batch())
+            break
 
     def _add_task(self, awaitable: Awaitable[Any]) -> None:
         """Add the given task to the tasks set for later execution."""

src/graphql/execution/incremental_publisher.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from asyncio import gather, sleep
+from asyncio import FIRST_COMPLETED, ensure_future, sleep, wait
 from contextlib import suppress
 from typing import (
     TYPE_CHECKING,
@@ -12,6 +12,7 @@
     cast,
 )
 
+from ..pyutils import is_awaitable
 from .incremental_graph import IncrementalGraph
 from .types import (
     CompletedResult,
@@ -30,6 +31,7 @@
     from collections.abc import AsyncGenerator, Iterable, Sequence
 
     from ..error import GraphQLError
+    from ..pyutils import AbortSignal
     from .types import (
         CancellableStreamRecord,
         CompletedExecutionGroup,
@@ -54,8 +56,17 @@
 class IncrementalPublisherContext(Protocol):
     """The context for incremental publishing."""
 
+    abort_signal: AbortSignal | None
     cancellable_streams: set[CancellableStreamRecord] | None
 
+    def abort_error(self) -> Exception:
+        """Return the exception to raise when execution has been aborted."""
+        ...  # pragma: no cover
+
+    async def cancel_incremental_work(self) -> None:
+        """Cancel all pending incremental work and close the stream sources."""
+        ...  # pragma: no cover
+
 
 class SubsequentIncrementalExecutionResultContext(NamedTuple):
     """The context for subsequent incremental execution results."""
@@ -134,9 +145,13 @@ async def _subscribe(
         incremental_graph = self._incremental_graph
         check_has_next = incremental_graph.has_next
         handle_completed_incremental_data = self._handle_completed_incremental_data
+        abort_signal = self._context.abort_signal
 
         try:
             while True:
+                if abort_signal is not None and abort_signal.aborted:
+                    raise self._context.abort_error()
+
                 batch: Iterable[IncrementalDataRecordResult] | None = (
                     incremental_graph.current_completed_batch()
                 )
@@ -160,21 +175,29 @@ async def _subscribe(
 
                         if not has_next:
                             return
-                    batch = await incremental_graph.next_completed_batch()
+
+                    next_batch = incremental_graph.next_completed_batch()
+                    if abort_signal is None:
+                        batch = await next_batch
+                    else:
+                        # reject the pending request when the operation is aborted
+                        abort = ensure_future(abort_signal.wait())
+                        try:
+                            await wait({next_batch, abort}, return_when=FIRST_COMPLETED)
+                        finally:
+                            if not abort.done():
+                                abort.cancel()
+                        if abort_signal.aborted:
+                            next_batch.cancel()
+                            raise self._context.abort_error()
+                        batch = next_batch.result()
         finally:
             await self._stop_async_iterators()
 
     async def _stop_async_iterators(self) -> None:
-        """Finish all async iterators."""
-        self._incremental_graph.stop_incremental_data()
-        cancellable_streams = self._context.cancellable_streams
-        if cancellable_streams is None:
-            return
-        early_returns = [
-            stream_record.early_return for stream_record in cancellable_streams
-        ]
-        if early_returns:
-            await gather(*early_returns, return_exceptions=True)
+        """Stop the incremental execution and finish all async iterators."""
+        await self._incremental_graph.stop_incremental_data()
+        await self._context.cancel_incremental_work()
 
     async def _handle_completed_incremental_data(
         self,
@@ -267,7 +290,9 @@ async def _handle_completed_stream_items(
                 if cancellable_streams:  # pragma: no branch
                     cancellable_streams.discard(stream_record)
                 with suppress(Exception):
-                    await stream_record.early_return
+                    early_return = stream_record.early_return()
+                    if is_awaitable(early_return):  # pragma: no branch
+                        await early_return
         elif stream_items_result.result is None:
             context.completed.append(CompletedResult(id_))
             incremental_graph.remove_stream(stream_record)

src/graphql/execution/types.py

@@ -891,13 +891,13 @@ def __repr__(self) -> str:
 class CancellableStreamRecord(StreamRecord):
     """Cancellable stream record"""
 
-    early_return: Awaitable[None]
+    early_return: Callable[[], Awaitable[None]]
 
     __slots__ = ("early_return",)
 
     def __init__(
         self,
-        early_return: Awaitable[None],
+        early_return: Callable[[], Awaitable[None]],
         stream_item_queue: list[StreamItemRecord],
         path: Path,
         label: str | None = None,

src/graphql/pyutils/boxed_awaitable_or_value.py

@@ -43,6 +43,14 @@ def value(self) -> T:
             self._value = value = value.result()
         return value  # type: ignore
 
+    @property
+    def pending_future(self) -> Future[T] | None:
+        """Get the still pending Future, or None if the value is already settled."""
+        value = self._value
+        if isfuture(value) and not value.done():
+            return value
+        return None
+
     def _update_value(self, value: Future[T]) -> None:
         """Update the boxed value when the Awaitable is done."""
         with suppress(CancelledError):