Rework map_async to handle failures better

`map` stops the flow of items in the stream when the function raises but
`map_async` is outside of the direct line of return so it fails weirdly
during an exception. To address that, I added the idea about stopping
the stream or not. This way, if the stream does not deliberately invoke
`stop` during an exception, the stream continues to process inputs after
an exception.

Since the `map_async` now conceives of stopping or not, I added a boolean
in the node state to control the loop inside the worker task.

In the case of an exception during mapping, `map_async` will now release
the references held on the metadata for the offending input.

I added an example that shows off the failure modes of `map` and
`map_async` that plainly demonstrates that exceptions can leave the
stream in a weird state.

Author: dwskoog

examples/map_failure_modes.py

@@ -0,0 +1,45 @@
+import asyncio
+from itertools import count
+from streamz import Stream
+
+
+async def flaky_async(x, from_where):
+    return flaky_sync(x, from_where)
+
+
+def flaky_sync(x, from_where):
+    if x % 5 == 4:
+        raise ValueError(f"I flaked out on {from_where}")
+    return x
+
+
+def make_counter(name):
+    return Stream.from_iterable(count(), asynchronous=True, stream_name=name)
+
+
+async def main():
+    async_non_stop_source = make_counter("async not stopping")
+    s_async = async_non_stop_source.map_async(flaky_async, async_non_stop_source)
+    s_async.rate_limit("500ms").sink(print, async_non_stop_source.name)
+
+    sync_source = make_counter("sync")
+    s_sync = sync_source.map(flaky_sync, sync_source)
+    s_sync.rate_limit("500ms").sink(print, sync_source.name)
+
+    async_stopping_source = make_counter("async stopping")
+    s_async = async_stopping_source.map_async(flaky_async, async_stopping_source, stop_on_exception=True)
+    s_async.rate_limit("500ms").sink(print, async_stopping_source.name)
+
+    async_non_stop_source.start()
+    sync_source.start()
+    async_stopping_source.start()
+    print(f"{async_non_stop_source.started=}, {sync_source.started=}, {async_stopping_source.started=}")
+    await asyncio.sleep(3)
+    print(f"{async_non_stop_source.stopped=}, {sync_source.stopped=}, {async_stopping_source.stopped=}")
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        pass

streamz/core.py

@@ -730,6 +730,8 @@ class map_async(Stream):
         The arguments to pass to the function.
     parallelism:
         The maximum number of parallel Tasks for evaluating func, default value is 1
+    stop_on_exception:
+         If the mapped func raises an exception, should the stream stop or not. Default value is False.
     **kwargs:
         Keyword arguments to pass to func
 
@@ -749,16 +751,23 @@ class map_async(Stream):
     6
     8
     """
-    def __init__(self, upstream, func, *args, parallelism=1, **kwargs):
+    def __init__(self, upstream, func, *args, parallelism=1, stop_on_exception=False, **kwargs):
         self.func = func
         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())
 
+    def stop(self):
+        if self.running:
+            self.running = False
+            super().stop()
+
     def update(self, x, who=None, metadata=None):
         return self._create_task(self._insert_job(x, metadata))
 
@@ -768,19 +777,20 @@ def _create_task(self, coro):
         return self.loop.asyncio_loop.create_task(coro)
 
     async def work_callback(self):
-        while True:
+        while self.running:
+            task, metadata = await self.work_queue.get()
+            self.work_queue.task_done()
             try:
-                task, metadata = await self.work_queue.get()
-                self.work_queue.task_done()
                 result = await task
             except Exception as e:
                 logger.exception(e)
-                raise
+                if self.stop_on_exception:
+                    self.stop()
             else:
                 results = self._emit(result, metadata=metadata)
                 if results:
                     await asyncio.gather(*results)
-                self._release_refs(metadata)
+            self._release_refs(metadata)
 
     async def _wait_for_work_slot(self):
         while self.work_queue.full():