Repository navigation
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a single-consumer threading contract to ResumableStreamIterator and AbstractResultSet, adding robust synchronization and state management to safely handle concurrent closures, lazy stream initialization, and asynchronous polling. Additionally, the internal buffer was migrated from LinkedList to ArrayDeque. The review feedback highlights two critical issues: first, startFailure is not cleared during resetStream(), which breaks the retry mechanism for startup failures; second, concurrent closures can cause uncaught exceptions to propagate from readNextFromStreamIntoBuffer when polling the underlying stream.
…umableStreamIterator Asynchronous result sets interact with ResumableStreamIterator from multiple threads: the caller thread initiating or closing the stream, gRPC callback threads notifying that messages or errors have arrived, and consumer threads polling for readiness and reading rows. Previously, stream lifecycle state and the retry buffer were not synchronized across these interactions, which caused several issues: 1. Duplicate or prematurely consumed streams during startup: When a gRPC callback fired while the initial stream was still being started and prefetched, the consumer thread could observe a partially initialized stream or start a second stream concurrently. Conversely, if a retryable error occurred immediately and the consumer thread reset the stream to retry, the finishing startup call could overwrite the reset state with the failed stream. 2. Uncancelled streams and deadlocks on concurrent close: Closing a result set while a stream was being created or retried could leave the newly created stream open, close it multiple times, or hold an internal lock while invoking external close callbacks. In addition, transaction metadata flags (whether the stream included an inlined BeginTransaction or was the last statement) were lost when the active stream reference was cleared on close or during retry backoff, preventing callers from properly cleaning up transactions and sessions. 3. Premature readiness signals for unresumable chunks: ResumableStreamIterator buffers chunks that do not carry a resume token so that a transient failure can be retried transparently without yielding duplicate data. However, non-blocking readiness checks delegated directly to the underlying gRPC stream and reported data as available as soon as a single tokenless chunk arrived. This caused asynchronous consumers to wake up and block on network I/O waiting for a resume token or buffer limit. How this change fixes these problems: - Coordinates stream startup so only one thread starts a stream at a time while concurrent callers wait for startup to complete, and tracks an in-flight stream separately from the published stream. A stream is only exposed to consumers and readiness checks after startup and initial prefetching have completed, while still allowing concurrent close calls to cancel an in-flight stream immediately and tearing down streams outside the lock. - Preserves stream transaction flags across stream resets, startup failures, and concurrent closes so upper layers always know whether an inlined transaction or final statement was in progress, and ensures the tracing span is always ended exactly once. - Drains available non-error chunks into the internal retry buffer during non-blocking readiness checks and only reports data as ready to poll when the buffer can actually be flushed (a resume token has been seen, the buffer is full, or the stream has reached the end), or when an error or closure needs to be surfaced.
a264732 to
d6a4474
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces thread-safety improvements and documents a single-consumer threading contract for ResumableStreamIterator and AbstractResultSet. Key changes include replacing LinkedList with ArrayDeque for buffering, introducing fine-grained synchronization via streamLock and buffer locks, and adding volatile fields to safely manage stream state transitions, lazy initialization, and concurrent cancellation. Additionally, JSpecify nullness annotations have been integrated. There are no review comments to assess, and I have no further feedback to provide on these changes.
Asynchronous result sets interact with ResumableStreamIterator from multiple threads: the caller thread initiating or closing the stream, gRPC callback threads notifying that messages or errors have arrived, and consumer threads polling for readiness and reading rows. Previously, stream lifecycle state and the retry buffer were not synchronized across these interactions, which caused several issues:
Duplicate or prematurely consumed streams during startup: When a gRPC callback fired while the initial stream was still being started and prefetched, the consumer thread could observe a partially initialized stream or start a second stream concurrently. Conversely, if a retryable error occurred immediately and the consumer thread reset the stream to retry, the finishing startup call could overwrite the reset state with the failed stream.
Uncancelled streams and deadlocks on concurrent close: Closing a result set while a stream was being created or retried could leave the newly created stream open, close it multiple times, or hold an internal lock while invoking external close callbacks. In addition, transaction metadata flags (whether the stream included an inlined BeginTransaction or was the last statement) were lost when the active stream reference was cleared on close or during retry backoff, preventing callers from properly cleaning up transactions and sessions.
Premature readiness signals for unresumable chunks: ResumableStreamIterator buffers chunks that do not carry a resume token so that a transient failure can be retried transparently without yielding duplicate data. However, non-blocking readiness checks delegated directly to the underlying gRPC stream and reported data as available as soon as a single tokenless chunk arrived. This caused asynchronous consumers to wake up and block on network I/O waiting for a resume token or buffer limit.
How this change fixes these problems: