Proper cancellation support for async generators.

Author: Michael Elsdörfer

tests/interop/test_calls.py

@@ -259,3 +259,23 @@ async def dly_asyncio():
         with pytest.raises(RuntimeError) as err:
             await async_gen_to_list(loop.wrap_generator(dly_asyncio))
         assert err.value.args[0] == "I has an owie"
+
+    @pytest.mark.trio
+    async def test_trio_asyncio_generator_with_cancellation(self, loop):
+        async def dly_asyncio(hold, seen):
+            yield 1
+            seen.flag |= 1
+            await hold.wait()
+
+        async def cancel_soon(nursery):
+            await trio.sleep(0.01)
+            nursery.cancel_scope.cancel()
+
+        hold = asyncio.Event(loop=loop)
+        seen = Seen()
+
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(async_gen_to_list, loop.wrap_generator(dly_asyncio, hold, seen))
+            nursery.start_soon(cancel_soon, nursery)
+        assert nursery.cancel_scope.cancel_called
+        assert seen.flag == 1

trio_asyncio/util.py

@@ -55,13 +55,17 @@ def abort_cb(raise_cancel_arg):
 async def run_generator(loop, async_generator):
     task = trio.hazmat.current_task()
     raise_cancel = None
+    current_read = None
 
     async def consume_next():
         try:
             item = await async_generator.__anext__()
             result = trio.hazmat.Value(value=item)
         except StopAsyncIteration:
             result = trio.hazmat.Value(value=STOP)
+        except asyncio.CancelledError:
+            # Once we are cancelled, we do not call reschedule() anymore
+            return
         except Exception as e:
             result = trio.hazmat.Error(error=e)
 
@@ -71,16 +75,29 @@ def abort_cb(raise_cancel_arg):
         # Save the cancel-raising function
         nonlocal raise_cancel
         raise_cancel = raise_cancel_arg
-        # XXX: we need to cancel any actice consume_next() call.
-        # Keep waiting
-        return trio.hazmat.Abort.FAILED
+
+        if not current_read:
+            # There is no current read
+            return trio.hazmat.Abort.SUCCEEDED
+        else:
+            # Attempt to cancel the current iterator read, do not
+            # report success until the future was actually cancelled.
+            already_cancelled = current_read.cancel()
+            if already_cancelled:
+                return trio.hazmat.Abort.SUCCEEDED
+            else:
+                # Continue dealing with the cancellation once
+                # future.cancel() goes to the result of
+                # wait_task_rescheduled()
+                return trio.hazmat.Abort.FAILED
 
     try:
         while True:
-            # schedule that we read the next one from the iterator
-            asyncio.ensure_future(consume_next(), loop=loop)
+            # Schedule in asyncio that we read the next item from the iterator
+            current_read = asyncio.ensure_future(consume_next(), loop=loop)
 
             item = await trio.hazmat.wait_task_rescheduled(abort_cb)
+
             if item is STOP:
                 break
             await yield_(item)