Fixed the subscription race condition (#5)

Author: dfee

graphql/subscription/map_async_iterator.py

@@ -1,4 +1,6 @@
-from inspect import isawaitable
+from asyncio import Event, ensure_future, wait
+from concurrent.futures import FIRST_COMPLETED
+from inspect import isasyncgen, isawaitable
 from typing import AsyncIterable, Callable
 
 __all__ = ['MapAsyncIterator']
@@ -19,35 +21,62 @@ def __init__(self, iterable: AsyncIterable, callback: Callable,
         self.iterator = iterable.__aiter__()
         self.callback = callback
         self.reject_callback = reject_callback
-        self.stop = False
+        self._close_event = Event()
+
+    @property
+    def closed(self) -> bool:
+        return self._close_event.is_set()
+
+    @closed.setter
+    def closed(self, value: bool) -> None:
+        if value:
+            self._close_event.set()
+        else:
+            self._close_event.clear()
 
     def __aiter__(self):
         return self
 
     async def __anext__(self):
-        if self.stop:
+        if self.closed:
+            if not isasyncgen(self.iterator):
+                raise StopAsyncIteration
+            result = await self.iterator.__anext__()
+            return self.callback(result)
+
+        _close = ensure_future(self._close_event.wait())
+        _next = ensure_future(self.iterator.__anext__())
+        done, pending = await wait(
+            [_close, _next],
+            return_when=FIRST_COMPLETED,
+        )
+
+        for task in pending:
+            task.cancel()
+
+        if _close.done():
             raise StopAsyncIteration
-        try:
-            value = await self.iterator.__anext__()
-        except Exception as error:
-            if not self.reject_callback or isinstance(error, (
-                    StopAsyncIteration, GeneratorExit)):
-                raise
-            result = self.reject_callback(error)
-        else:
-            result = self.callback(value)
-        if isawaitable(result):
-            result = await result
-        return result
+
+        if _next.done():
+            error = _next.exception()
+            if error:
+                if not self.reject_callback or isinstance(error, (
+                        StopAsyncIteration, GeneratorExit)):
+                    raise error
+                result = self.reject_callback(error)
+            else:
+                result = self.callback(_next.result())
+
+        return (await result) if isawaitable(result) else result
 
     async def athrow(self, type_, value=None, traceback=None):
-        if self.stop:
+        if self.closed:
             return
         athrow = getattr(self.iterator, 'athrow', None)
         if athrow:
             await athrow(type_, value, traceback)
         else:
-            self.stop = True
+            self.closed = True
             if value is None:
                 if traceback is None:
                     raise type_
@@ -57,13 +86,12 @@ async def athrow(self, type_, value=None, traceback=None):
             raise value
 
     async def aclose(self):
-        if self.stop:
+        if self.closed:
             return
         aclose = getattr(self.iterator, 'aclose', None)
         if aclose:
             try:
                 await aclose()
             except RuntimeError:
                 pass
-        else:
-            self.stop = True
+        self.closed = True

tests/subscription/test_map_async_iterator.py

@@ -1,3 +1,4 @@
+from asyncio import Event, ensure_future, sleep
 import sys
 
 from pytest import mark, raises
@@ -247,3 +248,27 @@ async def __anext__(self):
             # Throw error
             with raises(ValueError):
                 await doubles.athrow(ValueError, None, tb)
+
+    @mark.asyncio
+    async def stops_async_iteration_on_close():
+        async def source():
+            yield 1
+            await Event().wait()  # Block forever
+            yield 2
+            yield 3
+
+        doubles = MapAsyncIterator(source(), lambda x: x * 2)
+
+        result = await anext(doubles)
+        assert result == 2
+
+        # Block at event.wait()
+        fut = ensure_future(anext(doubles))
+        await sleep(.01)
+        assert not fut.done()
+
+        # Trigger cancellation and watch StopAsyncIteration propogate
+        await doubles.aclose()
+        await sleep(.01)
+        assert fut.done()
+        assert isinstance(fut.exception(), StopAsyncIteration)

tox.ini

@@ -27,4 +27,4 @@ deps =
     pytest-describe
 commands =
     python -m pip install -U pip
-    pytest
+    pytest {posargs}