perf(spanner): support PartialResultSet.last with background stream draining - #18320
perf(spanner): support PartialResultSet.last with background stream draining#18320olavloite wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces background stream draining for both synchronous and asynchronous gRPC streams in the Google Cloud Spanner client, allowing callers to return immediately upon receiving the last result set while the stream terminates cleanly in the background. Review feedback highlights performance improvements in both sync and async snapshot implementations by replacing inefficient O(N^2) pop(0) operations on item_buffer with O(N) iteration and clearing. Additionally, the reviewer noted a style guide violation in the unit tests where asyncio.create_task was mocked globally rather than at its local module import path.
7aece4d to
c4d0666
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements background stream draining for both synchronous and asynchronous gRPC streams in Google Cloud Spanner. When a partial result set indicates it is the last item, the stream is drained to EOF in the background (using an asyncio task for async streams, and a bounded thread pool for synchronous streams), allowing callers to return immediately. Additionally, iterators are properly cancelled in finally blocks to handle early termination. Feedback on the changes highlights a potential memory leak in _BoundedStreamDrainer where registering a bound method with os.register_at_fork creates a strong reference cycle, which should be resolved using a weak reference.
c4d0666 to
604b3a2
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces background stream draining for both synchronous and asynchronous gRPC streams in Google Cloud Spanner. When a PartialResultSet with last=True is encountered, the stream is drained to EOF in the background (using a thread pool for synchronous streams and asyncio.create_task for asynchronous streams), allowing callers to return immediately while ensuring clean stream termination. Feedback on the changes highlights a potential memory leak in unit tests due to registering the fork handler inside the _BoundedStreamDrainer constructor, and suggests moving the os.register_at_fork registration to the module level next to the global singleton instance.
604b3a2 to
fc3ecbd
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces background stream draining to EOF when PartialResultSet.last is True, allowing callers to return immediately while trailing gRPC metadata is consumed. This is implemented for both synchronous streams (using a bounded background thread pool _BoundedStreamDrainer) and asynchronous streams (using background asyncio tasks). The review feedback recommends replacing hardcoded asyncio.sleep calls in the async unit tests with deterministic polling of _PENDING_DRAIN_TASKS to prevent flaky tests in busy CI environments.
fc3ecbd to
8b4d58e
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements background stream draining for both synchronous and asynchronous gRPC streams in Google Cloud Spanner, allowing callers to return immediately when PartialResultSet.last is True while trailing metadata is consumed cleanly. The feedback suggests wrapping the thread pool initialization and queue operations in a try-except block to handle potential thread creation limits in restricted environments by falling back to inline draining. Additionally, a redundant local import of asyncio in the test suite should be removed.
8b4d58e to
56db3f3
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces background stream draining to EOF for both synchronous and asynchronous Spanner gRPC streams, allowing callers to return immediately upon receiving the last result set without blocking on trailing metadata. Feedback on the changes highlights a potential thread leak in the synchronous _BoundedStreamDrainer._ensure_started method if thread creation fails, suggesting a robust try-except block to safely manage the _started state.
…raining When Cloud Spanner finishes transmitting query results, it marks `last = True` on the final PartialResultSet chunk. Previously, the client blocked synchronously waiting for gRPC trailers and EOF frames over the wire before returning the final rows. Additionally, abandoning streams early caused gRPC's C-core finalizer to mark them as CANCELLED upon garbage collection. This change enables immediate return upon observing `last = True` and offloads trailing metadata consumption to the background so streams complete cleanly with status OK: - In sync mode, completed streams are handed off to `_BoundedStreamDrainer`, which uses a bounded queue and daemon worker threads to drain to EOF. If the queue is full or the interpreter is shutting down, it falls back to inline draining. Process fork safety is ensured via `os.register_at_fork`. - In async mode, trailing frames are drained via a background `asyncio.create_task` with strong reference retention to prevent premature task garbage collection and clean cancellation handling. - Transaction precommit tokens, query stats, and metadata present on the final chunk are captured before handing the stream off to background draining.
56db3f3 to
874a894
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces background stream draining for both synchronous and asynchronous gRPC streams in Google Cloud Spanner. When a result set indicates it is the last chunk, the stream is drained to EOF in the background (using an async task for async streams, and a bounded thread pool worker queue for sync streams) to allow the caller to return immediately while ensuring the stream terminates cleanly. Additionally, robust cancellation handling has been added to clean up active iterators upon early termination, accompanied by comprehensive unit tests. There are no review comments, so I have no feedback to provide on the review itself.
1 participant
When Cloud Spanner finishes transmitting query results, it marks
last = Trueon the final PartialResultSet chunk. Previously, the client blocked synchronously waiting for gRPC trailers and EOF frames over the wire before returning the final rows. Additionally, abandoning streams early caused gRPC's C-core finalizer to mark them as CANCELLED upon garbage collection.This change enables immediate return upon observing
last = Trueand offloads trailing metadata consumption to the background so streams complete cleanly with status OK:_BoundedStreamDrainer, which uses a bounded queue and daemon worker threads to drain to EOF. If the queue is full or the interpreter is shutting down, it falls back to inline draining. Process fork safety is ensured viaos.register_at_fork.asyncio.create_taskwith strong reference retention to prevent premature task garbage collection and clean cancellation handling.Benchmark Results
To verify that handling the last
PartialResultSetand marking the stream as completed does not introduce any performance regression, a 15-minutescheduled-steady-point-selectbenchmark was executed on GCE (n2-standard-4, sidecar enabled) and compared against the 7-day nightly baseline for the Python client.Latency Comparison
Takeaway
The change introduces zero performance overhead. All 84,477 queries completed with 0 errors, and latency remained well within the expected healthy baseline across all percentiles.