Cancel remaining iterator items on exceptions

Author: Cito

src/graphql/execution/execute.py

@@ -911,13 +911,22 @@ async def complete_async_iterator_value(
             index = awaitable_indices[0]
             completed_results[index] = await completed_results[index]
         else:
-            for index, result in zip(
-                awaitable_indices,
-                await gather(
-                    *(completed_results[index] for index in awaitable_indices)
-                ),
-            ):
-                completed_results[index] = result
+            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
+
+            for index, sub_result in zip(awaitable_indices, awaited_results):
+                completed_results[index] = sub_result
         return completed_results
 
     def complete_list_value(

tests/execution/test_parallel.py

@@ -281,3 +281,50 @@ async def resolve_list(*args):
         await barrier.wait()
         await asyncio.sleep(0)
         assert not completed
+
+    @pytest.mark.asyncio
+    async def cancel_async_iterator_on_exception():
+        barrier = Barrier(2)
+        completed = False
+
+        async def succeed(*_args):
+            nonlocal completed
+            await barrier.wait()
+            completed = True  # pragma: no cover
+
+        async def fail(*_args):
+            raise RuntimeError("Oops")
+
+        async def resolve_iterator(*args):
+            yield fail(*args)
+            yield succeed(*args)
+
+        schema = GraphQLSchema(
+            GraphQLObjectType(
+                "Query",
+                {
+                    "foo": GraphQLField(
+                        GraphQLList(GraphQLNonNull(GraphQLBoolean)),
+                        resolve=resolve_iterator,
+                    )
+                },
+            )
+        )
+
+        ast = parse("{foo}")
+
+        awaitable_result = execute(schema, ast)
+        assert isinstance(awaitable_result, Awaitable)
+        result = await asyncio.wait_for(awaitable_result, 1)
+
+        assert result == (
+            {"foo": None},
+            [{"message": "Oops", "locations": [(1, 2)], "path": ["foo", 0]}],
+        )
+
+        assert not completed
+
+        # Unblock succeed() and check that it does not complete
+        await barrier.wait()
+        await asyncio.sleep(0)
+        assert not completed