feat(crdb): experimental crdb query cancellation#3176
Conversation
github-actions
Bot
added
area/cli
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
ecordell
force-pushed
the
crdb-inband-cancel
branch
from
0d2ddf3 to
442de98
Compare
If the fallback behavior when the cancelpool isn't available is the old connection cancellation behavior, couldn't you still hit the degenerate condition? |
tstirrat15
left a comment
tstirrat15
left a comment
There was a problem hiding this comment.
See comments, otherwise LGTM
| // HandleCancel implements ctxwatch.Handler. | ||
| func (h *CancelQueryContextWatcherHandler) HandleCancel(_ context.Context) { | ||
| h.cancelFinished = make(chan struct{}) | ||
| h.cancelSucceeded.Store(false) |
There was a problem hiding this comment.
Is the idea here that there are potentially multiple requests using this same handler? I'm wondering why this isn't in the constructor
| go func() { | ||
| defer close(h.cancelFinished) | ||
| if err := h.canceler.CancelSessionQueries(h.pgConn); err != nil { | ||
| log.Warn().Err(err).Msg("in-band query cancellation failed; connection will be discarded") |
There was a problem hiding this comment.
At the risk of Yet Another Metric, successful and failed cancellations seem like they'd be good things to count
There was a problem hiding this comment.
Nvm I see you beat me to it
| config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { | ||
| if afterConnect != nil { | ||
| if err := afterConnect(ctx, conn); err != nil { | ||
| return err | ||
| } | ||
| } | ||
| return c.RegisterSession(ctx, conn) | ||
| } |
There was a problem hiding this comment.
It feels like there ought to be a nicer way to do this but I know that's not on this code.
| config.ConnConfig.BuildContextWatcherHandler = func(pgConn *pgconn.PgConn) ctxwatch.Handler { | ||
| return NewCancelQueryContextWatcherHandler(pgConn, c, defaultCancelDeadlineDelay) | ||
| } |
There was a problem hiding this comment.
Do we need to do the same wrapping behavior here?
| func beginFuncExec(ctx context.Context, tx pgx.Tx, fn func(pgx.Tx) error) (err error) { | ||
| defer func() { | ||
| rollbackErr := tx.Rollback(ctx) | ||
| rollbackErr := tx.Rollback(context.WithoutCancel(ctx)) |
There was a problem hiding this comment.
Was this a latent bug or is this necessitated by the new approach?
There was a problem hiding this comment.
latent bug! in cases where we would need to rollback due to a context cancel the rollback would always fail too.
ecordell
changed the title
2 participants
When a request context is canceled mid-query on the CRDB datastore, pgx's default behavior is to destroy the connection. Under load this drains the write pool and triggers the metastable death spiral described in #2576. A previous attempt to fix this using pgx's built-in
CancelRequestContextWatcherHandler(pgwire cancel protocol) was reverted in #2434 because CRDB applies pgwire cancels asynchronously; a late-arriving cancel could kill the next query on the connection.This PR implements cancellation using CockroachDB's
CANCEL QUERIESstatement instead. A newpool.Cancelerowns a small dedicated connection pool and a registry mapping each pooled write connection to its CRDBsession_id(captured viaSHOW session_idat connect). A customctxwatch.Handlerfires on context cancellation and issuesCANCEL QUERIES IF EXISTS (SELECT query_id FROM [SHOW CLUSTER STATEMENTS] WHERE session_id = '...')on a sibling connection, then blocks until that statement completes before releasing the original connection back to the pool. This sequencing eliminates the wrong-query race by construction: no cancellation can be in flight when the next query starts on the same session.The feature is gated behind
--datastore-experimental-crdb-query-cancellation(default off). When enabled, set--write-conn-acquisition-timeout=0; with connections no longer being destroyed on cancel, pool exhaustion under load is far less likely and the acquisition-timeout backpressure is no longer needed.