From b4938b84d86e80cec8b416b389a3d8f4f92dfb4e Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Tue, 7 Jul 2026 16:59:30 +0200 Subject: [PATCH] fix: more fixes ... --- .../providers/in_memory_stream_service.py | 1 + .../streaming/stream/stream_reader.py | 350 +++++++++++------- .../streaming/test_stream_multiplexer.py | 49 +-- 3 files changed, 244 insertions(+), 156 deletions(-) diff --git a/private_gpt/components/streaming/providers/in_memory_stream_service.py b/private_gpt/components/streaming/providers/in_memory_stream_service.py index 138b3fa4..574ba869 100644 --- a/private_gpt/components/streaming/providers/in_memory_stream_service.py +++ b/private_gpt/components/streaming/providers/in_memory_stream_service.py @@ -175,6 +175,7 @@ class InMemoryStreamService(StreamService): self._waiters[correlation_id].add(waiter) if waiter is not None: + assert block_ms is not None try: await asyncio.wait_for(waiter.wait(), timeout=block_ms / 1000) except TimeoutError: diff --git a/private_gpt/components/streaming/stream/stream_reader.py b/private_gpt/components/streaming/stream/stream_reader.py index e86e5475..54c1b68b 100644 --- a/private_gpt/components/streaming/stream/stream_reader.py +++ b/private_gpt/components/streaming/stream/stream_reader.py @@ -3,7 +3,7 @@ import contextlib import logging from collections import defaultdict from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from injector import inject, singleton from pydantic import BaseModel @@ -20,7 +20,9 @@ logger = logging.getLogger(__name__) DEFAULT_BLOCK_MS: Final[int] = 1000 BATCH_SIZE: Final[int] = 5 -STATUS_CHECK_INTERVAL: Final[float] = 5.0 +# How often (seconds) to poll for terminal status when no events arrive. +# Halves Redis metadata round-trips vs. checking every block cycle (1 s). +STATUS_CHECK_INTERVAL: Final[float] = 2.0 TERMINAL_STATUSES: Final[set[StreamStatus]] = { StreamStatus.COMPLETED, StreamStatus.CANCELLED, @@ -28,39 +30,118 @@ TERMINAL_STATUSES: Final[set[StreamStatus]] = { } -class StreamConsumer: - """A single consumer receiving events for one correlation id. +class StreamBroadcast: + """Append-only batch log that broadcasts to N concurrent readers. - Events are buffered in an unbounded ``asyncio.Queue``. A ``None`` entry - is the sentinel that signals the consumer the stream is finished. + A single producer calls ``push`` / ``close`` (both synchronous); each + subscriber holds an integer *cursor* into ``_batches`` and parks on a + ``asyncio.Future`` until new data arrives. One ``push`` resolves *all* + pending futures in a single pass — no per-subscriber queues. + + Safety: all operations are synchronous except ``read_from``, so no lock is + needed (asyncio cooperative threading guarantees no interleaving between + non-await statements). """ - def __init__(self, correlation_id: str, event_handler: EventHandler): - self.correlation_id = correlation_id - self.event_handler = event_handler - self.queue: asyncio.Queue[BaseModel | None] = asyncio.Queue() - self.ref_count = 1 + def __init__(self) -> None: + self._batches: list[list[BaseModel] | None] = [] + self._waiters: list[asyncio.Future[None]] = [] + + def _notify(self) -> None: + waiters, self._waiters = self._waiters, [] + for fut in waiters: + if not fut.done(): + fut.set_result(None) + + def push(self, batch: list[BaseModel]) -> None: + self._batches.append(batch) + self._notify() def close(self) -> None: - """Signal the consumer that no more events will arrive.""" - self.queue.put_nowait(None) + """Append the terminal sentinel and wake all parked readers.""" + self._batches.append(None) + self._notify() + + def next_cursor(self) -> int: + """Cursor for a subscriber joining right now (no backfill).""" + return len(self._batches) + + @property + def is_closed(self) -> bool: + return bool(self._batches) and self._batches[-1] is None + + async def read_from( + self, + cursor: int, + stop_event: asyncio.Event | None = None, + ) -> AsyncGenerator[BaseModel, None]: + """Yield events starting at *cursor*, blocking until new ones arrive.""" + while True: + # Drain already-available batches. Safe without a lock: _batches is + # append-only and list operations are atomic between await points. + while cursor < len(self._batches): + batch = self._batches[cursor] + cursor += 1 + if batch is None: + return + for event in batch: + yield event + + if stop_event and stop_event.is_set(): + return + + # Park until the producer pushes more data. + fut: asyncio.Future[None] = asyncio.get_event_loop().create_future() + self._waiters.append(fut) + try: + await asyncio.wait_for(fut, timeout=1.0) + except TimeoutError: + pass # Re-check stop_event and drain any new batches + finally: + # Clean up if push() hasn't already swapped us out. + if fut in self._waiters: + self._waiters.remove(fut) + + +class StreamConsumer: + """A subscriber on a ``StreamBroadcast`` with an individual read cursor.""" + + def __init__( + self, + correlation_id: str, + event_handler: EventHandler, + broadcast: StreamBroadcast, + cursor: int, + ) -> None: + self.correlation_id = correlation_id + self.event_handler = event_handler + self.broadcast = broadcast + self.cursor = cursor # Index into broadcast._batches class StreamState: - """Read position and cache for a single multiplexed stream.""" + """Read position and broadcast log for a single multiplexed stream.""" - def __init__(self, last_id: str = "0"): + def __init__(self, last_id: str = "0") -> None: self.last_id = last_id self.cached_events: list[BaseModel] | None = None + # 0.0 → check on the first idle cycle (ancient enough to pass the gate) self.last_status_check: float = 0.0 + self.broadcast = StreamBroadcast() @singleton class StreamReader: - """Reads events from the backend (Redis or in-memory) and deserializes them.""" + """Reads events from the backend (Redis or in-memory) and deserializes them. + + On Redis, ``read_events`` with ``block_ms > 0`` issues ``XREAD BLOCK`` which + parks the connection — the event loop stays free. On the in-memory backend + the provider uses an ``asyncio.Event`` to avoid busy-waiting. + ``block_ms=None`` is always non-blocking (used for catch-up / drain reads). + """ @inject - def __init__(self, stream_component: StreamComponent): + def __init__(self, stream_component: StreamComponent) -> None: self.stream_service = stream_component.stream async def read_events( @@ -107,9 +188,9 @@ class StreamReader: logger.error(f"Error checking cached status for {correlation_id}: {e}") try: - metadata: ( - StreamMetadata | None - ) = await self.stream_service.get_stream_metadata(correlation_id) + metadata: StreamMetadata | None = ( + await self.stream_service.get_stream_metadata(correlation_id) + ) if metadata and metadata.status in TERMINAL_STATUSES: return True except Exception as e: @@ -140,7 +221,7 @@ class StreamReader: async def produce_events(current_last_id: str) -> None: try: cached_events: list[BaseModel] | None = None - last_status_check: float = 0.0 + last_status_check: float = 0.0 # 0 → check on first idle cycle while not stop_event.is_set(): events, current_last_id = await self.read_events( @@ -205,24 +286,22 @@ class StreamReader: class StreamMultiplexer: """Multiplexes many stream consumers onto a single worker task. - Each correlation id has one ``StreamState`` (read position) shared by all - consumers on that stream. The worker loops over every active stream, - issues a blocking ``read_events`` and dispatches the batch to every - current consumer. + Each correlation id has one ``StreamState`` containing a ``StreamBroadcast`` + shared by all subscribers. The worker calls ``_process_stream`` once per + active stream per cycle: one Redis read, one ``broadcast.push`` — no + per-consumer queues. Key invariants --------------- - * Dispatch and ``last_id`` advance happen atomically under the lock, and - the consumer list is re-fetched immediately before dispatch. This means - a consumer that joins while the worker is blocked inside ``read_events`` - still receives the batch about to be dispatched. - * ``add_consumer`` accepts a ``last_id`` and initialises the stream state - from it, so a reconnecting consumer never receives a replay from ``"0"``. - * ``stop`` cancels the worker, clears all state and closes every consumer; - a subsequent ``start`` begins from a clean slate. + * ``broadcast.push`` is synchronous and resolves all parked futures in one + pass, so all N subscribers wake atomically after a single call. + * State is updated under the lock; ``broadcast.push`` happens outside it so + the lock is held for as short a time as possible. + * ``close`` on the broadcast terminates all ``read_from`` generators, making + ``stop`` and terminal-status cleanup self-contained. """ - def __init__(self, stream_reader: StreamReader): + def __init__(self, stream_reader: StreamReader) -> None: self.stream_reader = stream_reader self.consumers: dict[str, list[StreamConsumer]] = defaultdict(list) self.states: dict[str, StreamState] = {} @@ -237,7 +316,7 @@ class StreamMultiplexer: logger.info("Stream multiplexer started") async def stop(self) -> None: - """Cancel the worker, clear all state, and close every consumer.""" + """Cancel the worker, clear all state, and close every broadcast.""" self.shutdown_event.set() task = self.worker_task self.worker_task = None @@ -247,11 +326,12 @@ class StreamMultiplexer: await task async with self._lock: - all_consumers = [c for cl in self.consumers.values() for c in cl] + states_to_close = list(self.states.values()) self.consumers.clear() self.states.clear() - for consumer in all_consumers: - consumer.close() + + for state in states_to_close: + state.broadcast.close() async def add_consumer( self, @@ -259,47 +339,40 @@ class StreamMultiplexer: event_handler: EventHandler, last_id: str = "0", ) -> StreamConsumer: - """Register a consumer for *correlation_id*. + """Register a new subscriber for *correlation_id*. - If a consumer with the same handler type already exists its ref-count - is incremented and the existing consumer is returned (reconnect case). - Otherwise a new consumer is created. When this is the first consumer - for the stream, the read position is initialised to *last_id*. + Each call returns a distinct ``StreamConsumer`` with its own cursor so + all concurrent subscribers receive every event independently (true + fan-out). When this is the first consumer for the stream the read + position is initialised to *last_id*. """ async with self._lock: - for consumer in self.consumers.get(correlation_id, []): - if type(consumer.event_handler) is type(event_handler): - consumer.ref_count += 1 - return consumer - - consumer = StreamConsumer(correlation_id, event_handler) - self.consumers[correlation_id].append(consumer) - if correlation_id not in self.states: self.states[correlation_id] = StreamState(last_id=last_id) + state = self.states[correlation_id] + cursor = state.broadcast.next_cursor() + consumer = StreamConsumer( + correlation_id, event_handler, state.broadcast, cursor + ) + self.consumers[correlation_id].append(consumer) return consumer async def remove_consumer(self, consumer: StreamConsumer) -> None: - should_close = False + broadcast_to_close: StreamBroadcast | None = None async with self._lock: correlation_id = consumer.correlation_id - consumer.ref_count -= 1 + cl = self.consumers.get(correlation_id) + if cl and consumer in cl: + cl.remove(consumer) + if correlation_id in self.consumers and not self.consumers[correlation_id]: + del self.consumers[correlation_id] + state = self.states.pop(correlation_id, None) + if state: + broadcast_to_close = state.broadcast - if consumer.ref_count <= 0: - cl = self.consumers.get(correlation_id) - if cl and consumer in cl: - cl.remove(consumer) - if ( - correlation_id in self.consumers - and not self.consumers[correlation_id] - ): - del self.consumers[correlation_id] - self.states.pop(correlation_id, None) - should_close = True - - if should_close: - consumer.close() + if broadcast_to_close is not None: + broadcast_to_close.close() async def _run_worker(self) -> None: try: @@ -317,9 +390,12 @@ class StreamMultiplexer: except Exception as e: logger.error(f"Worker error: {e}", exc_info=True) finally: - for cl in list(self.consumers.values()): - for consumer in cl: - consumer.close() + async with self._lock: + states_to_close = list(self.states.values()) + self.consumers.clear() + self.states.clear() + for state in states_to_close: + state.broadcast.close() async def _process_stream(self, correlation_id: str) -> None: try: @@ -330,9 +406,10 @@ class StreamMultiplexer: consumer_list = list(self.consumers.get(correlation_id, [])) if not consumer_list: return + event_handler = consumer_list[0].event_handler events, new_last_id = await self.stream_reader.read_events( - event_handler=consumer_list[0].event_handler, + event_handler=event_handler, correlation_id=correlation_id, last_id=state.last_id, block_ms=DEFAULT_BLOCK_MS, @@ -349,11 +426,11 @@ class StreamMultiplexer: state.last_status_check = current_time if await self.stream_reader.check_terminal_status( correlation_id, - consumer_list[0].event_handler, + event_handler, state.cached_events, ): drain_events, drain_last_id = await self.stream_reader.read_events( - event_handler=consumer_list[0].event_handler, + event_handler=event_handler, correlation_id=correlation_id, last_id=state.last_id, block_ms=None, @@ -380,43 +457,45 @@ class StreamMultiplexer: new_last_id: str, current_time: float, ) -> None: - """Dispatch *events* to every current consumer and advance ``last_id``. + """Advance stream state then broadcast *events* to all subscribers. - The consumer list is re-fetched under the lock so that consumers that - joined or left during the preceding blocking ``read_events`` are - reflected. Dispatch and ``last_id`` advance are atomic: no consumer - can join between the two and miss the batch. + State is updated under the lock; ``broadcast.push`` is synchronous and + happens outside so the lock is held for the minimum time. """ async with self._lock: if correlation_id not in self.states: return - consumer_list = list(self.consumers.get(correlation_id, [])) - if not consumer_list: + if not self.consumers.get(correlation_id): return - for consumer in consumer_list: - for event in events: - consumer.queue.put_nowait(event) state.last_id = new_last_id state.cached_events = events state.last_status_check = current_time + state.broadcast.push(events) + async def _close_all_consumers(self, correlation_id: str) -> None: async with self._lock: - consumer_list = list(self.consumers.get(correlation_id, [])) + state = self.states.pop(correlation_id, None) self.consumers.pop(correlation_id, None) - self.states.pop(correlation_id, None) - for consumer in consumer_list: - consumer.close() + if state: + state.broadcast.close() @singleton class AdaptiveStreamReader: - """Switches between the direct reader and the multiplexer based on load. + """Routes stream consumers between direct and multiplexed paths. - When the number of concurrently active streams reaches - ``multiplexing_threshold`` the multiplexed path is used; otherwise the - direct ``StreamReader.stream_events`` path is used. Both paths must - produce identical results for the same stream and ``last_id``. + Direct path: one XREAD BLOCK per consumer → one Redis connection held. + Multiplexed path: one shared worker for all active streams on this process. + + When concurrent direct-path readers reach ``multiplexing_threshold`` new + consumers are routed to the multiplexer instead of opening another Redis + connection. When the multiplexer drains to zero active streams it is + stopped, so load-shedding is fully reversible — consumers that arrive when + load is back below the threshold go direct again. + + ``multiplexing_threshold=None`` disables multiplexing entirely. + ``enable_multiplexing=False`` is a hard kill-switch (e.g. for tests). """ @inject @@ -424,13 +503,22 @@ class AdaptiveStreamReader: self, settings: Settings, stream_reader: StreamReader, - ): + ) -> None: self.stream_reader = stream_reader self.multiplexer: StreamMultiplexer | None = None - self._active_count = 0 self._lock = asyncio.Lock() self.enable_multiplexing = True - self.multiplexing_threshold = settings.chat.multiplexing_threshold or None + self.multiplexing_threshold: int | None = settings.chat.multiplexing_threshold + # Number of consumers currently on the direct path. + # Each one holds an XREAD BLOCK connection; this is what we throttle. + self._direct_count = 0 + + def _should_multiplex(self) -> bool: + return ( + self.enable_multiplexing + and self.multiplexing_threshold is not None + and self._direct_count >= self.multiplexing_threshold + ) async def read_events( self, @@ -457,16 +545,14 @@ class AdaptiveStreamReader: stop_event: asyncio.Event | None = None, ) -> AsyncGenerator[BaseModel, None]: async def gen() -> AsyncGenerator[BaseModel, None]: + # Decide path under the lock so _direct_count is consistent. async with self._lock: - self._active_count += 1 - should_multiplex = ( - self.enable_multiplexing - and self.multiplexing_threshold is not None - and self._active_count >= self.multiplexing_threshold - ) + use_mux = self._should_multiplex() + if not use_mux: + self._direct_count += 1 try: - if should_multiplex: + if use_mux: async for event in await self._stream_multiplexed( event_handler, correlation_id, last_id, stop_event ): @@ -477,8 +563,9 @@ class AdaptiveStreamReader: ): yield event finally: - async with self._lock: - self._active_count -= 1 + if not use_mux: + async with self._lock: + self._direct_count -= 1 return gen() @@ -490,41 +577,38 @@ class AdaptiveStreamReader: stop_event: asyncio.Event | None, ) -> AsyncGenerator[BaseModel, None]: async def gen() -> AsyncGenerator[BaseModel, None]: - if self.multiplexer is None: - self.multiplexer = StreamMultiplexer(self.stream_reader) - await self.multiplexer.start() - logger.info("Multiplexed streaming activated") + # Start the multiplexer on first use; capture the reference so the + # finally block targets the right instance even if it is replaced. + async with self._lock: + if self.multiplexer is None: + self.multiplexer = StreamMultiplexer(self.stream_reader) + await self.multiplexer.start() + logger.info("Multiplexer started (threshold reached)") + mux = self.multiplexer - consumer = await self.multiplexer.add_consumer( + consumer = await mux.add_consumer( correlation_id, event_handler, last_id=last_id ) - try: - events_processed = 0 - while True: - if stop_event and stop_event.is_set(): - break - - try: - event: Any | None = await asyncio.wait_for( - consumer.queue.get(), timeout=1.0 - ) - if event is None: - break - yield event - - events_processed += 1 - if events_processed % BATCH_SIZE == 0: - await asyncio.sleep(0) - - except TimeoutError: - continue + async for event in consumer.broadcast.read_from( + cursor=consumer.cursor, + stop_event=stop_event, + ): + yield event finally: - await self.multiplexer.remove_consumer(consumer) + await mux.remove_consumer(consumer) + # Stop the multiplexer when it has no more active streams so the + # next consumer below the threshold can go direct again. + async with self._lock: + if self.multiplexer is mux and not mux.states: + await mux.stop() + self.multiplexer = None + logger.info("Multiplexer stopped (no active streams)") return gen() async def shutdown(self) -> None: - if self.multiplexer: - await self.multiplexer.stop() - self.multiplexer = None + async with self._lock: + mux, self.multiplexer = self.multiplexer, None + if mux: + await mux.stop() diff --git a/tests/components/streaming/test_stream_multiplexer.py b/tests/components/streaming/test_stream_multiplexer.py index 50d02a8d..3b76df45 100644 --- a/tests/components/streaming/test_stream_multiplexer.py +++ b/tests/components/streaming/test_stream_multiplexer.py @@ -38,6 +38,7 @@ from private_gpt.components.streaming.providers.redis_stream_service import ( from private_gpt.components.streaming.providers.stream_service import StreamService from private_gpt.components.streaming.stream.stream_reader import ( DEFAULT_BLOCK_MS, + StreamConsumer, StreamMultiplexer, StreamReader, ) @@ -76,16 +77,18 @@ class MockStreamComponent: self.stream = service -async def collect_until_closed( - queue: asyncio.Queue[BaseModel | None], +async def collect_from_consumer( + consumer: StreamConsumer, timeout: float = 10.0, ) -> list[BaseModel]: + """Collect all events from a consumer's broadcast until it is closed.""" events: list[BaseModel] = [] - while True: - event = await asyncio.wait_for(queue.get(), timeout=timeout) - if event is None: - break - events.append(event) + + async def drain() -> None: + async for event in consumer.broadcast.read_from(cursor=consumer.cursor): + events.append(event) + + await asyncio.wait_for(drain(), timeout=timeout) return events @@ -113,7 +116,7 @@ async def run_multiplexed( last_id: str = "0", ) -> list[BaseModel]: consumer = await multiplexer.add_consumer(cid, handler, last_id=last_id) - return await collect_until_closed(consumer.queue, timeout=30.0) + return await collect_from_consumer(consumer, timeout=30.0) async def push_n_and_complete( @@ -396,9 +399,9 @@ async def test_late_joiner_receives_inflight_batch( await service.update_stream_status(cid, StreamStatus.COMPLETED) try: - recv1 = await collect_until_closed(c1.queue, timeout=10.0) + recv1 = await collect_from_consumer(c1, timeout=10.0) c2 = late_box["c2"] - recv2 = await collect_until_closed(c2.queue, timeout=10.0) + recv2 = await collect_from_consumer(c2, timeout=10.0) finally: await multiplexer.stop() stream_reader.read_events = original_read # type: ignore[method-assign] @@ -450,7 +453,7 @@ async def test_concurrent_remove_preserves_delivery( await service.update_stream_status(cid, StreamStatus.COMPLETED) try: - recv1 = await collect_until_closed(c1.queue, timeout=5.0) + recv1 = await collect_from_consumer(c1, timeout=5.0) finally: await multiplexer.stop() stream_reader.read_events = original_read # type: ignore[method-assign] @@ -475,8 +478,8 @@ async def test_multiple_consumers_all_receive( c1 = await multiplexer.add_consumer(cid, handler) c2 = await multiplexer.add_consumer(cid, handler2) try: - recv1 = await collect_until_closed(c1.queue, timeout=10.0) - recv2 = await collect_until_closed(c2.queue, timeout=10.0) + recv1 = await collect_from_consumer(c1, timeout=10.0) + recv2 = await collect_from_consumer(c2, timeout=10.0) finally: await multiplexer.stop() @@ -505,8 +508,6 @@ async def test_events_in_terminal_gap_drained( correlation_id: str, event_handler: Any, cached_events: Any, - last_flush_time: Any, - cache_flush_interval: Any, ) -> bool: await service.push_event(correlation_id, json.dumps({"n": 1})) await service.push_event(correlation_id, json.dumps({"n": 2})) @@ -518,7 +519,7 @@ async def test_events_in_terminal_gap_drained( await multiplexer.start() consumer = await multiplexer.add_consumer(cid, handler) try: - received = await collect_until_closed(consumer.queue, timeout=10.0) + received = await collect_from_consumer(consumer, timeout=10.0) finally: await multiplexer.stop() stream_reader.check_terminal_status = original_check # type: ignore[method-assign] @@ -542,17 +543,19 @@ async def test_stop_clears_state_and_closes_consumers( await multiplexer.start() consumer = await multiplexer.add_consumer(cid, handler) - await asyncio.wait_for(consumer.queue.get(), timeout=5.0) + + # Wait until the worker has pushed at least one batch to the broadcast. + async def get_one() -> None: + async for _ in consumer.broadcast.read_from(cursor=consumer.cursor): + break + + await asyncio.wait_for(get_one(), timeout=5.0) await multiplexer.stop() assert not multiplexer.consumers, "consumers dict must be empty after stop()" assert not multiplexer.states, "states dict must be empty after stop()" - - remaining: list[Any] = [] - while not consumer.queue.empty(): - remaining.append(consumer.queue.get_nowait()) - assert None in remaining, "consumer queue must have a None sentinel" + assert consumer.broadcast.is_closed, "broadcast must be closed after stop()" async def test_restart_after_stop( @@ -577,7 +580,7 @@ async def test_restart_after_stop( await multiplexer.start() try: consumer = await multiplexer.add_consumer(cid2, handler) - received = await collect_until_closed(consumer.queue, timeout=10.0) + received = await collect_from_consumer(consumer, timeout=10.0) finally: await multiplexer.stop()