Make map_async restartable

Author: dwskoog

streamz/core.py

@@ -1,12 +1,13 @@
 import asyncio
+import concurrent.futures
 from collections import deque, defaultdict
 from datetime import timedelta
 from itertools import chain
 import functools
 import logging
 import threading
 from time import time
-from typing import Any, Callable, Hashable, Union
+from typing import Any, Callable, Coroutine, Hashable, Tuple, Union, overload
 import weakref
 
 import toolz
@@ -756,28 +757,54 @@ def __init__(self, upstream, func, *args, parallelism=1, stop_on_exception=False
         stream_name = kwargs.pop('stream_name', None)
         self.kwargs = kwargs
         self.args = args
-        self.running = True
         self.stop_on_exception = stop_on_exception
         self.work_queue = asyncio.Queue(maxsize=parallelism)
 
         Stream.__init__(self, upstream, stream_name=stream_name, ensure_io_loop=True)
-        self.work_task = self._create_task(self.work_callback())
+        self.work_task = None
+
+    def _create_work_task(self) -> Tuple[asyncio.Event, asyncio.Task[None]]:
+        stop_work = asyncio.Event()
+        work_task = self._create_task(self.work_callback(stop_work))
+        return stop_work, work_task
+
+    def start(self):
+        if self.work_task:
+            stop_work, _ = self.work_task
+            stop_work.set()
+        self.work_task = self._create_work_task()
+        super().start()
 
     def stop(self):
-        if self.running:
-            self.running = False
-            super().stop()
+        stop_work, _ = self.work_task
+        stop_work.set()
+        self.work_task = None
+        super().stop()
 
     def update(self, x, who=None, metadata=None):
+        if not self.work_task:
+            self.work_task = self._create_work_task()
         return self._create_task(self._insert_job(x, metadata))
 
+    @overload
+    def _create_task(self, coro: asyncio.Future) -> asyncio.Future:
+        ...
+
+    @overload
+    def _create_task(self, coro: concurrent.futures.Future) -> concurrent.futures.Future:
+        ...
+
+    @overload
+    def _create_task(self, coro: Coroutine) -> asyncio.Task:
+        ...
+
     def _create_task(self, coro):
         if gen.is_future(coro):
             return coro
         return self.loop.asyncio_loop.create_task(coro)
 
-    async def work_callback(self):
-        while self.running:
+    async def work_callback(self, stop_work: asyncio.Event):
+        while not stop_work.is_set():
             task, metadata = await self.work_queue.get()
             self.work_queue.task_done()
             try:

streamz/tests/test_core.py

@@ -151,6 +151,28 @@ def fail_func():
     assert (time() - start) == pytest.approx(0.1, abs=4e-3)
 
 
+@pytest.mark.asyncio
+async def test_map_async_restart():
+    async def flake_out(x):
+        if x == 2:
+            raise RuntimeError("I fail on 2.")
+        if x > 4:
+            raise RuntimeError("I fail on > 4.")
+        return x
+
+    source = Stream.from_iterable(itertools.count())
+    mapped = source.map_async(flake_out, stop_on_exception=True)
+    results = mapped.sink_to_list()
+    source.start()
+
+    await await_for(lambda: results == [0, 1], 1)
+    await await_for(lambda: not mapped.work_task, 1)
+
+    source.start()
+
+    await await_for(lambda: results == [0, 1, 3, 4], 1)
+
+
 @pytest.mark.asyncio
 async def test_map_async():
     @gen.coroutine