Skip to content

Fix: Prevent unbounded queue memory spikes and task leaks in streams.concat - #165

Open
THRISHAL12345 wants to merge 2 commits into
google-gemini:mainfrom
THRISHAL12345:fix/streams-concat-bounded-queue
Open

Fix: Prevent unbounded queue memory spikes and task leaks in streams.concat#165
THRISHAL12345 wants to merge 2 commits into
google-gemini:mainfrom
THRISHAL12345:fix/streams-concat-bounded-queue

Conversation

@THRISHAL12345

Copy link
Copy Markdown

Summary of Changes

This pull request fixes a potential Out-Of-Memory (OOM) vulnerability and resource leak in streams.concat when processing heavy multimodal streams (video, audio, PDF documents).

Key Improvements:

  1. Bounded Queue Backpressure in streams.concat (genai_processors/streams.py):
    • Added a queue_maxsize: int = 1 keyword argument to streams.concat(*contents, queue_maxsize: int = 1).
    • Replaced unbounded asyncio.Queue() (maxsize=0) with bounded queues (asyncio.Queue(maxsize=queue_maxsize)).
    • Replaced non-blocking put_nowait(c) calls with enqueue(c, output_queues[idx]), which uses await queue.put(part) to pause upstream producers when queue_maxsize items are buffered ahead.
    • Callers who require unbounded queues can explicitly pass queue_maxsize=0.
  2. Task Lifecycle & Cancellation Safety:
    • Wrapped streams.concat iteration in a try...finally block that cleanly cancels background enqueue tasks (t.cancel()) when iteration completes or terminates early.
    • Updated enqueue so that when a task is cancelled (asyncio.CancelledError), it re-raises immediately without attempting to block on await queue.put(None) into a full queue.
  3. Unit Tests (genai_processors/tests/streams_test.py):
    • Added test_concat_bounded_queue_backpressure: Verifies that bounded queueing correctly prevents downstream streams from running ahead and over-buffering.
    • Added test_concat_early_cancellation: Verifies that early termination (aclose()) cleanly cancels background enqueue tasks.

Verification

  • Ran unit tests: test_concat_bounded_queue_backpressure and test_concat_early_cancellation pass cleanly.
  • Verified existing StreamsTest unit tests continue to pass without regression.

Closes #164

@google-cla

google-cla Bot commented Jul 8, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces backpressure and cancellation safety to the concat stream utility by adding a queue_maxsize parameter and ensuring background tasks are cancelled when the generator exits. It also adds corresponding unit tests for bounded queue backpressure and early cancellation. The review feedback suggests leveraging the project's idiomatic context.context() task group in concat to handle background task exceptions robustly and simplify task cleanup. Additionally, it recommends removing a redundant try...except asyncio.CancelledError block in the enqueue function.

Comment thread genai_processors/streams.py Outdated
Comment thread genai_processors/streams.py Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Prevent unbounded queue memory spikes and task leaks in streams.concat

1 participant