Merged
Conversation
This PR ensures that active database queries are stopped immediately when a client connection is lost. Key changes: - Added a connection-level context to 'clientConn' that is cancelled when the session ends. - Updated 'queryContext()' to derive from the connection context, ensuring all queries (Query, Exec, Cursor operations) stop when the client disconnects. - Updated 'FlightExecutor' to correctly propagate cancellation to worker processes via gRPC context cancellation. - Hardened constructors and destructors to prevent nil-pointer panics in unit tests.
The previous commits set up a context chain (clientConn.ctx → queryContext → FlightExecutor.mergedContext) but nothing actually detected TCP disconnects while a query was in-flight. The message loop is single-threaded, so it only noticed disconnects after the current query completed. This commit adds a Peek-based disconnect monitor goroutine that runs during query execution. It uses bufio.Reader.Peek(1) with short read deadlines to detect TCP FIN/RST without consuming bytes from the stream. When a disconnect is detected, the connection context is cancelled, propagating through the context chain to cancel gRPC calls on Flight SQL workers. Key changes: - startDisconnectMonitor() in conn.go polls the connection during queries - RunMessageLoop() now defers cc.cancel() for proper cleanup in control plane - FlightExecutor context merging extracted into mergedContext() helper - Tests for disconnect detection, stop-before-disconnect, and data preservation Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Cursors span multiple message loop iterations (FETCH, FETCH, CLOSE). The disconnect monitor uses bufio.Reader.Peek which would race with readMessage between FETCH calls. Use queryContextForCursor() which skips the monitor — disconnect detection still works via c.ctx cancellation when the connection handler exits. Also add a comment at the COPY dispatch site noting it must stay above queryContext() since handleCopyIn reads from c.reader directly. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
When pipelined data is in the bufio.Reader buffer, Peek(1) returns instantly without hitting the network, causing a tight spin loop. Add a 50ms throttle on the success path to prevent CPU waste. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Previously, if a client disconnected while a query was running, the worker would continue executing the query until completion. This was especially problematic for long-running analytical queries.
This PR implements end-to-end query cancellation: