diff --git a/streamz/core.py b/streamz/core.py index 6d49d5a8..7f306d81 100644 --- a/streamz/core.py +++ b/streamz/core.py @@ -1356,7 +1356,8 @@ def cb(self): L, self._buffer = self._buffer, [] metadata, self.metadata_buffer = self.metadata_buffer, [] m = [m for ml in metadata for m in ml] - self.last = self._emit(L, m) + # Both this callback and subsequent updates await the same work. + self.last = gen.convert_yielded(self._emit(L, m)) self._release_refs(m) yield self.last yield gen.sleep(self.interval) @@ -1475,7 +1476,8 @@ def cb(self): metadata_result, self._metadata_buffer = list(self._metadata_buffer.values()), {} # TODO: figure out why metadata_result is handled differently here... m = [m for ml in metadata_result for m in ml] - self.last = self._emit(result, m) + # Both this callback and subsequent updates await the same work. + self.last = gen.convert_yielded(self._emit(result, m)) self._release_refs(m) yield self.last yield gen.sleep(self.interval) diff --git a/streamz/tests/test_core.py b/streamz/tests/test_core.py index 67261255..dac928a5 100644 --- a/streamz/tests/test_core.py +++ b/streamz/tests/test_core.py @@ -548,6 +548,84 @@ def read_from_q(): assert stop - start > 0.2 +@pytest.mark.asyncio +@pytest.mark.parametrize("window_type, options", [ + ("timed_window", {}), + ("timed_window_unique", {"keep": "first"}), + ("timed_window_unique", {"keep": "last"}), +]) +@pytest.mark.parametrize("sink_type", ["coroutine", "task", "tornado"]) +async def test_timed_window_shared_backpressure(monkeypatch, window_type, options, sink_type): + started = asyncio.Queue() + intervals = asyncio.Queue() + release = asyncio.Event() + completed = [] + synchronous = [] + + def wait_for_interval(interval): + assert interval == 60 + future = asyncio.get_running_loop().create_future() + intervals.put_nowait(future) + return future + + monkeypatch.setattr(gen, "sleep", wait_for_interval) + + async def write(batch): + started.put_nowait(tuple(batch)) + await release.wait() + completed.append(tuple(batch)) + + if sink_type == "task": + def sink(batch): + return asyncio.create_task(write(batch)) + elif sink_type == "tornado": + @gen.coroutine + def sink(batch): + yield write(batch) + else: + sink = write + + source = Stream(asynchronous=True) + window = getattr(source, window_type)(60, **options) + asynchronous_sink = window.sink(sink) + synchronous_sink = window.sink(lambda batch: synchronous.append(tuple(batch))) + try: + await source.emit(0) + assert await asyncio.wait_for(started.get(), 5) == (0,) + + # Several producers share the in-flight window's backpressure. + pending = [source.emit(1), source.emit(2)] + await asyncio.sleep(0) + assert not any(future.done() for future in pending) + assert completed == [] + release.set() + await asyncio.wait_for(asyncio.gather(*pending), 5) + interval = await asyncio.wait_for(intervals.get(), 5) + + # Subsequent inputs must also be able to await the completed window. + await source.emit(3) + await source.emit(4) + assert completed == [(0,)] + assert synchronous == [(0,)] + + release.clear() + interval.set_result(None) + assert await asyncio.wait_for(started.get(), 5) == (1, 2, 3, 4) + pending = source.emit(5) + await asyncio.sleep(0) + assert not pending.done() + release.set() + await asyncio.wait_for(pending, 5) + await asyncio.wait_for(intervals.get(), 5) + assert completed == [(0,), (1, 2, 3, 4)] + assert synchronous == completed + finally: + release.set() + asynchronous_sink.destroy() + synchronous_sink.destroy() + window.destroy() + + def test_sink_to_file(): with tmpfile() as fn: source = Stream()