Extract task cancellation as utility function

Author: Cito

src/graphql/execution/execute.py

@@ -4,9 +4,7 @@
 
 from asyncio import (
     CancelledError,
-    create_task,
     ensure_future,
-    gather,
     shield,
     wait_for,
 )
@@ -52,6 +50,7 @@
     RefMap,
     Undefined,
     async_reduce,
+    gather_with_cancel,
     inspect,
     is_iterable,
 )
@@ -466,21 +465,9 @@ async def get_results() -> dict[str, Any]:
                 field = awaitable_fields[0]
                 results[field] = await results[field]
             else:
-                tasks = [
-                    create_task(results[field])  # type: ignore[arg-type]
-                    for field in awaitable_fields
-                ]
-
-                try:
-                    awaited_results = await gather(*tasks)
-                except Exception:
-                    # Cancel unfinished tasks before raising the exception
-                    for task in tasks:
-                        if not task.done():
-                            task.cancel()
-                    await gather(*tasks, return_exceptions=True)
-                    raise
-
+                awaited_results = await gather_with_cancel(
+                    *(results[field] for field in awaitable_fields)
+                )
                 results.update(zip(awaitable_fields, awaited_results))
 
             return results
@@ -911,20 +898,9 @@ async def complete_async_iterator_value(
             index = awaitable_indices[0]
             completed_results[index] = await completed_results[index]
         else:
-            tasks = [
-                create_task(completed_results[index]) for index in awaitable_indices
-            ]
-
-            try:
-                awaited_results = await gather(*tasks)
-            except Exception:
-                # Cancel unfinished tasks before raising the exception
-                for task in tasks:
-                    if not task.done():
-                        task.cancel()
-                await gather(*tasks, return_exceptions=True)
-                raise
-
+            awaited_results = await gather_with_cancel(
+                *(completed_results[index] for index in awaitable_indices)
+            )
             for index, sub_result in zip(awaitable_indices, awaited_results):
                 completed_results[index] = sub_result
         return completed_results
@@ -1023,20 +999,9 @@ async def get_completed_results() -> list[Any]:
                 index = awaitable_indices[0]
                 completed_results[index] = await completed_results[index]
             else:
-                tasks = [
-                    create_task(completed_results[index]) for index in awaitable_indices
-                ]
-
-                try:
-                    awaited_results = await gather(*tasks)
-                except Exception:
-                    # Cancel unfinished tasks before raising the exception
-                    for task in tasks:
-                        if not task.done():
-                            task.cancel()
-                    await gather(*tasks, return_exceptions=True)
-                    raise
-
+                awaited_results = await gather_with_cancel(
+                    *(completed_results[index] for index in awaitable_indices)
+                )
                 for index, sub_result in zip(awaitable_indices, awaited_results):
                     completed_results[index] = sub_result
             return completed_results
@@ -2123,21 +2088,7 @@ def default_type_resolver(
     if awaitable_is_type_of_results:
         # noinspection PyShadowingNames
         async def get_type() -> str | None:
-            tasks = [
-                create_task(result)  # type: ignore[arg-type]
-                for result in awaitable_is_type_of_results
-            ]
-
-            try:
-                is_type_of_results = await gather(*tasks)
-            except Exception:
-                # Cancel unfinished tasks before raising the exception
-                for task in tasks:
-                    if not task.done():
-                        task.cancel()
-                await gather(*tasks, return_exceptions=True)
-                raise
-
+            is_type_of_results = await gather_with_cancel(*awaitable_is_type_of_results)
             for is_type_of_result, type_ in zip(is_type_of_results, awaitable_types):
                 if is_type_of_result:
                     return type_.name

src/graphql/pyutils/__init__.py

@@ -9,6 +9,7 @@
 """
 
 from .async_reduce import async_reduce
+from .gather_with_cancel import gather_with_cancel
 from .convert_case import camel_to_snake, snake_to_camel
 from .cached_property import cached_property
 from .description import (
@@ -52,6 +53,7 @@
     "cached_property",
     "camel_to_snake",
     "did_you_mean",
+    "gather_with_cancel",
     "group_by",
     "identity_func",
     "inspect",

src/graphql/pyutils/gather_with_cancel.py

@@ -0,0 +1,36 @@
+"""Run awaitables concurrently with cancellation support."""
+
+from __future__ import annotations
+
+from asyncio import Task, create_task, gather
+from typing import Any, Awaitable
+
+__all__ = ["gather_with_cancel"]
+
+
+async def gather_with_cancel(*awaitables: Awaitable[Any]) -> list[Any]:
+    """Run awaitable objects in the sequence concurrently.
+
+    The first raised exception is immediately propagated to the task that awaits
+    on this function and all pending awaitables in the sequence will be cancelled.
+
+    This is different from the default behavior or `asyncio.gather` which waits
+    for all tasks to complete even if one of them raises an exception. It is also
+    different from `asyncio.gather` with `return_exceptions` set, which does not
+    cancel the other tasks when one of them raises an exception.
+    """
+    try:
+        tasks: list[Task[Any]] = [
+            aw if isinstance(aw, Task) else create_task(aw)  # type: ignore[arg-type]
+            for aw in awaitables
+        ]
+    except TypeError:
+        return await gather(*awaitables)  # type: ignore[arg-type]
+    try:
+        return await gather(*tasks)
+    except Exception:
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        await gather(*tasks, return_exceptions=True)
+        raise

tests/pyutils/test_gather_with_cancel.py

@@ -0,0 +1,133 @@
+from asyncio import Event, create_task, gather, sleep, wait_for
+from typing import Callable
+
+import pytest
+
+from graphql.pyutils import gather_with_cancel, is_awaitable
+
+
+class Controller:
+    def reset(self, wait=False):
+        self.event = Event()
+        if not wait:
+            self.event.set()
+        self.returned = []
+
+
+controller = Controller()
+
+
+async def coroutine(value: int) -> int:
+    """Simple coroutine that returns a value."""
+    if value > 2:
+        raise RuntimeError("Oops")
+    await controller.event.wait()
+    controller.returned.append(value)
+    return value
+
+
+class CustomAwaitable:
+    """Custom awaitable that return a value."""
+
+    def __init__(self, value: int):
+        self.value = value
+        self.coroutine = coroutine(value)
+
+    def __await__(self):
+        return self.coroutine.__await__()
+
+
+awaitable_factories: dict[str, Callable] = {
+    "coroutine": coroutine,
+    "task": lambda value: create_task(coroutine(value)),
+    "custom": lambda value: CustomAwaitable(value),
+}
+
+with_all_types_of_awaitables = pytest.mark.parametrize(
+    "type_of_awaitable", awaitable_factories
+)
+
+
+def describe_gather_with_cancel():
+    @with_all_types_of_awaitables
+    @pytest.mark.asyncio
+    async def gathers_all_values(type_of_awaitable: str):
+        return  # !!!s
+        factory = awaitable_factories[type_of_awaitable]
+        values = list(range(3))
+
+        controller.reset()
+        aws = [factory(i) for i in values]
+
+        assert await gather(*aws) == values
+        assert controller.returned == values
+
+        controller.reset()
+        aws = [factory(i) for i in values]
+
+        result = gather_with_cancel(*aws)
+        assert is_awaitable(result)
+
+        awaited = await wait_for(result, 1)
+        assert awaited == values
+
+    @with_all_types_of_awaitables
+    @pytest.mark.asyncio
+    async def raises_on_exception(type_of_awaitable: str):
+        return  # !!!
+        factory = awaitable_factories[type_of_awaitable]
+        values = list(range(4))
+
+        controller.reset()
+        aws = [factory(i) for i in values]
+
+        with pytest.raises(RuntimeError, match="Oops"):
+            await gather(*aws)
+        assert controller.returned == values[:-1]
+
+        controller.reset()
+        aws = [factory(i) for i in values]
+
+        result = gather_with_cancel(*aws)
+        assert is_awaitable(result)
+
+        with pytest.raises(RuntimeError, match="Oops"):
+            await wait_for(result, 1)
+        assert controller.returned == values[:-1]
+
+    @with_all_types_of_awaitables
+    @pytest.mark.asyncio
+    async def cancels_on_exception(type_of_awaitable: str):
+        factory = awaitable_factories[type_of_awaitable]
+        values = list(range(4))
+
+        controller.reset(wait=True)
+        aws = [factory(i) for i in values]
+
+        with pytest.raises(RuntimeError, match="Oops"):
+            await gather(*aws)
+        assert not controller.returned
+
+        # check that the standard gather continues to produce results
+        controller.event.set()
+        await sleep(0)
+        assert controller.returned == values[:-1]
+
+        controller.reset(wait=True)
+        aws = [factory(i) for i in values]
+
+        result = gather_with_cancel(*aws)
+        assert is_awaitable(result)
+
+        with pytest.raises(RuntimeError, match="Oops"):
+            await wait_for(result, 1)
+        assert not controller.returned
+
+        # check that gather_with_cancel stops producing results
+        controller.event.set()
+        await sleep(0)
+        if type_of_awaitable == "custom":
+            # Cancellation of custom awaitables is not supported
+            assert controller.returned == values[:-1]
+        else:
+            assert not controller.returned