Cancel remaining fields on exceptions

gather() returns when the first exception is raised, but does not cancel
any remaining tasks. These continue to run which is inefficient, and can
also cause problems if they access shared resources like database
connections.

Fixes: #236

Author: mgorven

src/graphql/execution/execute.py

@@ -2,7 +2,16 @@
 
 from __future__ import annotations
 
-from asyncio import ensure_future, gather, shield, wait_for
+from asyncio import (
+    FIRST_EXCEPTION,
+    CancelledError,
+    create_task,
+    ensure_future,
+    gather,
+    shield,
+    wait,
+    wait_for,
+)
 from contextlib import suppress
 from copy import copy
 from typing import (
@@ -459,12 +468,18 @@ async def get_results() -> dict[str, Any]:
                 field = awaitable_fields[0]
                 results[field] = await results[field]
             else:
-                results.update(
-                    zip(
-                        awaitable_fields,
-                        await gather(*(results[field] for field in awaitable_fields)),
-                    )
-                )
+                tasks = {}
+                for field in awaitable_fields:
+                    tasks[create_task(results[field])] = field  # type: ignore[arg-type]
+
+                done, pending = await wait(tasks, return_when=FIRST_EXCEPTION)
+                if pending:
+                    for task in pending:
+                        task.cancel()
+                    await wait(pending)
+
+                results.update((tasks[task], task.result()) for task in done)
+
             return results
 
         return get_results()
@@ -538,6 +553,10 @@ async def await_completed() -> Any:
                     try:
                         return await completed
                     except Exception as raw_error:
+                        # Before Python 3.8 CancelledError inherits Exception and
+                        # so gets caught here.
+                        if isinstance(raw_error, CancelledError):
+                            raise
                         self.handle_field_error(
                             raw_error,
                             return_type,

tests/execution/test_parallel.py

@@ -11,6 +11,7 @@
     GraphQLInt,
     GraphQLInterfaceType,
     GraphQLList,
+    GraphQLNonNull,
     GraphQLObjectType,
     GraphQLSchema,
     GraphQLString,
@@ -193,3 +194,40 @@ async def is_type_of_baz(obj, *_args):
             {"foo": [{"foo": "bar", "foobar": 1}, {"foo": "baz", "foobaz": 2}]},
             None,
         )
+
+    @pytest.mark.asyncio
+    async def cancel_on_exception():
+        barrier = Barrier(2)
+        completed = False
+
+        async def succeed(*_args):
+            nonlocal completed
+            await barrier.wait()
+            completed = True
+
+        async def fail(*_args):
+            raise Exception
+
+        schema = GraphQLSchema(
+            GraphQLObjectType(
+                "Query",
+                {
+                    "foo": GraphQLField(GraphQLNonNull(GraphQLBoolean), resolve=fail),
+                    "bar": GraphQLField(GraphQLBoolean, resolve=succeed),
+                },
+            )
+        )
+
+        ast = parse("{foo, bar}")
+
+        awaitable_result = execute(schema, ast)
+        assert isinstance(awaitable_result, Awaitable)
+        result = await asyncio.wait_for(awaitable_result, 1.0)
+
+        assert result.errors
+        assert not result.data
+
+        # Unblock succeed() and check that it does not complete
+        await barrier.wait()
+        await asyncio.sleep(0)
+        assert not completed