Interrupt in-flight TiKV coprocessor RPC when a killed/disconnected SELECT is still executing
- Dominant language
- Go
- Stars
- 40.5k
- Forks
- 6.2k
- PR merge metrics
- PR metrics pending
Description
## Enhancement
### Background
When a client connection is closed (or a query is killed) while a `SELECT` is
still executing, TiDB should stop the query **and** promptly abort the in-flight
TiKV coprocessor RPC. Two gaps exist today:
1. The passive connection-liveness probe
(`shouldInstallConnectionAliveDuringExecute`) was only installed for
autocommit non-txn `INSERT`/`UPDATE`/`DELETE`, so long-running `SELECT`s were
not covered.
2. When a kill signal is raised through the passive path, only the polled
`Killed` flag is set. That flag is observed only at coprocessor task
boundaries (`finished()`) and backoff points (`bo.CheckKilled()`). A request
already blocked inside a single `SendReqCtx` RPC does not observe the kill
until `CoprReqTimeout` (default 60s) elapses.
### Proposal
This tracks two related optimizations:
**1. Extend the liveness probe to cover `SELECT`.**
Add `*ast.SelectStmt` to `shouldInstallConnectionAliveDuringExecute` so the
liveness probe is installed for autocommit non-txn SELECTs (it already unwraps
`*ast.ExecuteStmt` to the prepared inner statement). Long-running read queries
whose client has disconnected can now be detected and killed via the existing
passive checkpoint path, instead of running to completion.
**2. Cancel the execution context on kill to abort the in-flight RPC.**
Register a context-cancel hook on `SQLKiller`
(`ConnCancel atomic.Pointer[func()]`), wired to the dispatch context's cancel
(`cc.cancelDispatch`). When `sendKillSignal` wins the CAS that sets the kill
signal, it also invokes this cancel. Cancelling the dispatch context aborts the
in-flight gRPC/coprocessor RPC via a stream reset instead of waiting for
`CoprReqTimeout`. This makes the passive disconnect path behave like the
existing `KILL` path (`killQuery` already calls `cancelDispatch()`
unconditionally). `SQLKiller.Reset()` clears `ConnCancel` so no stale cancel
leaks across statements on the same connection.
This keeps the design of #68685 (no always-on hot-path monitor goroutine): the
probe is a passive, throttled checkpoint check, and short statements pay no
goroutine/ticker/channel cost.
### Compatibility / correctness notes
- All kill sources funnel through `sendKillSignal`'s
`CompareAndSwapUint32(&Signal, 0, reason)`; first writer wins, so
`MaxExecTimeExceeded` set by the expensive-query background checker is not
overwritten by a later passive `QueryInterrupted`. The passive probe only
raises `QueryInterrupted` when the connection is confirmed dead, never on
unknown state.
- `context.CancelFunc` is idempotent, so the redundant cancel on the
`KILL`/`max_execution_time` path is harmless.
### Known limitation
If a coprocessor request is stuck inside TiKV and never returns, the execution
goroutine is blocked in `SendReqCtx` and never reaches any checkpoint
(`HandleSignal` / `finished()`), so nothing on the execution thread can trigger
the cancel. A blocked thread cannot self-detect; the cancel must come from
another execution context. `CoprReqTimeout` (default 60s) remains the hard
backstop for that case. A scoped, lazily-started low-frequency watcher (active
only while a long-running probe-installed query runs) is a candidate follow-up.
### Affected code
- `pkg/server/conn.go`: `shouldInstallConnectionAliveDuringExecute`,
`setSQLKillerConnectionAlive`.
- `pkg/util/sqlkiller/sqlkiller.go`: new `ConnCancel` field, cancel invocation
in `sendKillSignal`, cleanup in `Reset`.
Contributor guide
Research direction
Start in pkg/server/conn.go by tracing shouldInstallConnectionAliveDuringExecute and setSQLKillerConnectionAlive, then read pkg/util/sqlkiller/sqlkiller.go around sendKillSignal and Reset. Done means autocommit non-transaction SELECTs use the liveness probe and a kill cancels the dispatch context so an in-flight coprocessor RPC aborts promptly without stale cancellation across statements.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- go
- Domain
- databases, distributed-systems
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100