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,15 @@
 
 from __future__ import annotations
 
-from asyncio import ensure_future, gather, shield, wait_for
+from asyncio import (
+    FIRST_EXCEPTION,
+    create_task,
+    ensure_future,
+    gather,
+    shield,
+    wait,
+    wait_for,
+)
 from contextlib import suppress
 from copy import copy
 from typing import (
@@ -459,12 +467,16 @@ 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)
+                for task in pending:
+                    task.cancel()
+
+                results.update((tasks[task], task.result()) for task in done)
+
             return results
 
         return get_results()

tests/execution/test_parallel.py

@@ -11,6 +11,7 @@
     GraphQLInt,
     GraphQLInterfaceType,
     GraphQLList,
+    GraphQLNonNull,
     GraphQLObjectType,
     GraphQLSchema,
     GraphQLString,
@@ -193,3 +194,39 @@ 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 = asyncio.Event()
+
+        async def succeed(*_args):
+            await barrier.wait()
+            completed.set()
+
+        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.is_set()