perf(spanner): support PartialResultSet.last with background stream draining

[olavloite]

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.

Benchmark Results

To verify that handling the last PartialResultSet and marking the stream as completed does not introduce any performance regression, a 15-minute scheduled-steady-point-select benchmark was executed on GCE (n2-standard-4, sidecar enabled) and compared against the 7-day nightly baseline for the Python client.

• Workload: Steady Point-Select (100 TPS, 15m duration, sidecar enabled)
• Target PR Run: 84,477 queries executed, 0 errors
• Baseline: 7-Day Nightly Average (5,014,740 samples)

Latency Comparison

Metric	PR Branch (Full 15m)	PR Branch (Warm Steady-State)	7-Day Nightly Baseline	Delta (vs Baseline)
Mean	4.563 ms	4.542 ms	5.388 ms	-15.3%
P50	4.401 ms	4.410 ms	5.079 ms	-13.4%
P90	5.723 ms	5.731 ms	6.850 ms	-16.5%
P99	7.573 ms	7.527 ms	10.542 ms	-28.2%

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.

[gemini-code-assist]

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.

[olavloite]

/gemini review

[gemini-code-assist]

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.

[olavloite]

/gemini review

[gemini-code-assist]

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.

[olavloite]

/gemini review

[gemini-code-assist]

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.

[olavloite]

/gemini review

[gemini-code-assist]

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.

[olavloite]

/gemini review

[gemini-code-assist]

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.

[olavloite]

/gemini review

[gemini-code-assist]

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.