feat(crdb): experimental crdb query cancellation #3176
There are no files selected for viewing
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,121 @@ | ||
| //go:build datastore | ||
|
|
||
| package crdb | ||
|
|
||
| import ( | ||
| "context" | ||
| "math/rand/v2" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/jackc/pgx/v5" | ||
| "github.com/jackc/pgx/v5/pgconn" | ||
| "github.com/prometheus/client_golang/prometheus/testutil" | ||
| "github.com/stretchr/testify/require" | ||
|
|
||
| "github.com/authzed/spicedb/internal/datastore/crdb/pool" | ||
| testdatastore "github.com/authzed/spicedb/internal/testserver/datastore" | ||
| "github.com/authzed/spicedb/pkg/datastore" | ||
| ) | ||
|
|
||
| // newCancellationTestDatastore builds an unwrapped *crdbDatastore with | ||
| // in-band query cancellation enabled, against a dockerized CRDB. | ||
| func newCancellationTestDatastore(t *testing.T) *crdbDatastore { | ||
| t.Helper() | ||
| b := testdatastore.RunCRDBForTesting(t, "", crdbTestVersion()) | ||
|
|
||
| var cds *crdbDatastore | ||
| b.NewDatastore(t, func(engine, uri string) datastore.Datastore { | ||
| ds, err := newCRDBDatastore(t.Context(), uri, | ||
| GCWindow(veryLargeGCWindow), | ||
| RevisionQuantization(0), | ||
| WithQueryCancellation(true), | ||
| WithAcquireTimeout(5*time.Second), | ||
| ) | ||
| require.NoError(t, err) | ||
| t.Cleanup(func() { _ = ds.Close() }) | ||
| cds = ds.(*crdbDatastore) | ||
| return ds | ||
| }) | ||
| require.NotNil(t, cds) | ||
| return cds | ||
| } | ||
|
|
||
| func scanIgnore(_ context.Context, row pgx.Row) error { | ||
| var ignored any | ||
| return row.Scan(&ignored) | ||
| } | ||
|
|
||
| func TestQueryCancellationSessionIDsRegistered(t *testing.T) { | ||
| cds := newCancellationTestDatastore(t) | ||
| ctx := t.Context() | ||
|
|
||
| // Force at least one write connection to exist. | ||
| require.NoError(t, cds.writePool.QueryRowFunc(ctx, scanIgnore, "SELECT 1")) | ||
|
|
||
| found := 0 | ||
| cds.writePool.Range(func(conn *pgx.Conn, _ uint32) { | ||
| if _, ok := cds.canceler.SessionID(conn.PgConn()); ok { | ||
| found++ | ||
| } | ||
| }) | ||
| require.Positive(t, found, "write pool connections must have registered session IDs") | ||
| } | ||
|
|
||
| func TestQueryCancellationRollsBackWriteAndKeepsConnection(t *testing.T) { | ||
| cds := newCancellationTestDatastore(t) | ||
| ctx := t.Context() | ||
|
|
||
| require.NoError(t, cds.writePool.ExecFunc(ctx, | ||
| func(_ context.Context, _ pgconn.CommandTag, err error) error { return err }, | ||
| "CREATE TABLE IF NOT EXISTS cancellation_test (id INT PRIMARY KEY)")) | ||
|
|
||
| baselineNewConns := cds.writePool.Stat().NewConnsCount() | ||
|
|
||
| cancelCtx, cancel := context.WithTimeout(ctx, 200*time.Millisecond) | ||
| defer cancel() | ||
| err := cds.writePool.BeginTxFunc(cancelCtx, pgx.TxOptions{}, func(tx pgx.Tx) error { | ||
| if _, err := tx.Exec(cancelCtx, "INSERT INTO cancellation_test (id) VALUES (1)"); err != nil { | ||
| return err | ||
| } | ||
| _, err := tx.Exec(cancelCtx, "SELECT pg_sleep(10)") | ||
| return err | ||
| }) | ||
| require.ErrorIs(t, err, context.DeadlineExceeded) | ||
|
|
||
| // The transaction must have been rolled back on a live connection. | ||
| var count int | ||
| require.NoError(t, cds.writePool.QueryRowFunc(ctx, func(_ context.Context, row pgx.Row) error { | ||
| return row.Scan(&count) | ||
| }, "SELECT COUNT(*) FROM cancellation_test")) | ||
| require.Zero(t, count, "canceled transaction must be rolled back") | ||
| require.Equal(t, baselineNewConns, cds.writePool.Stat().NewConnsCount(), | ||
| "write cancellation must not destroy and replace the connection") | ||
| } | ||
|
|
||
| // TestQueryCancellationStressNoStaleErrors races cancellation against query | ||
| // completion many times and verifies that no cancellation ever bleeds into a | ||
| // subsequent query on the same pool — the failure mode that forced the revert | ||
| // of pgwire-protocol cancellation in #2434. | ||
| func TestQueryCancellationStressNoStaleErrors(t *testing.T) { | ||
| cds := newCancellationTestDatastore(t) | ||
| ctx := t.Context() | ||
|
|
||
| tripwireBefore := testutil.ToFloat64(pool.UnexpectedCancellationErrors) | ||
|
|
||
| for i := 0; i < 300; i++ { | ||
| // Jitter the deadline around the query duration (~10ms) so the | ||
| // cancellation races completion in both directions. | ||
| timeout := time.Duration(5+rand.IntN(15)) * time.Millisecond //nolint:gosec // jitter, not crypto | ||
| cancelCtx, cancel := context.WithTimeout(ctx, timeout) | ||
| _ = cds.writePool.QueryRowFunc(cancelCtx, scanIgnore, "SELECT pg_sleep(0.01)") | ||
| cancel() | ||
|
|
||
| err := cds.writePool.QueryRowFunc(ctx, scanIgnore, "SELECT 1") | ||
| require.NoError(t, err, | ||
| "iteration %d: follow-up query failed — possible stale cancellation", i) | ||
| } | ||
|
|
||
| require.InDelta(t, tripwireBefore, testutil.ToFloat64(pool.UnexpectedCancellationErrors), 1e-9, | ||
| "no query may observe 57014 without its own context being canceled") | ||
| } |
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,101 @@ | ||
| package pool | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync/atomic" | ||
| "time" | ||
|
|
||
| "github.com/jackc/pgx/v5/pgconn" | ||
|
|
||
| log "github.com/authzed/spicedb/internal/logging" | ||
| ) | ||
|
|
||
| // deadlineSetter is the subset of net.Conn the handler needs; it exists so | ||
| // the handler can be unit tested without a live connection. | ||
| type deadlineSetter interface { | ||
| SetDeadline(t time.Time) error | ||
| } | ||
|
|
||
| // sessionCanceler cancels all queries running on a tracked session. It is | ||
| // implemented by *Canceler. | ||
| type sessionCanceler interface { | ||
| CancelSessionQueries(pgConn *pgconn.PgConn) error | ||
| } | ||
|
|
||
| // CancelQueryContextWatcherHandler responds to context cancellation during an | ||
| // in-flight pgconn operation by issuing CockroachDB's in-band CANCEL QUERIES | ||
| // statement from a sibling connection (via the Canceler), so the query stops | ||
| // promptly without destroying this connection. | ||
| // | ||
| // Sequencing guarantee: HandleUnwatchAfterCancel blocks until the CANCEL | ||
| // statement round trip completes. pgconn invokes it before the canceled | ||
| // operation returns to the caller, and the pool cannot release the connection | ||
| // before the operation returns. A cancellation therefore can never still be | ||
| // in flight when a later query runs on this session — unlike the pgwire | ||
| // cancel protocol, which CockroachDB applies asynchronously after closing the | ||
| // cancel connection. | ||
| // | ||
| // If the CANCEL statement fails or times out, the connection is poisoned with | ||
| // an immediate deadline: its next use errors out, the pool destroys it, and | ||
| // the pool-level retry logic transparently moves the work to another | ||
| // connection. This degrades to the legacy close-the-connection behavior. | ||
| type CancelQueryContextWatcherHandler struct { | ||
| conn deadlineSetter | ||
| pgConn *pgconn.PgConn | ||
| canceler sessionCanceler | ||
| deadlineDelay time.Duration | ||
|
|
||
| cancelFinished chan struct{} | ||
| cancelSucceeded atomic.Bool | ||
| } | ||
|
|
||
| // NewCancelQueryContextWatcherHandler builds a handler for a connection. | ||
| // deadlineDelay is the fallback net.Conn deadline applied when a context is | ||
| // canceled: if in-band cancellation does not resolve the in-flight operation | ||
| // before it elapses, the pending read errors out and the connection is | ||
| // destroyed. It must comfortably exceed the canceler's statement timeout. | ||
| func NewCancelQueryContextWatcherHandler(pgConn *pgconn.PgConn, canceler sessionCanceler, deadlineDelay time.Duration) *CancelQueryContextWatcherHandler { | ||
| return &CancelQueryContextWatcherHandler{ | ||
| conn: pgConn.Conn(), | ||
| pgConn: pgConn, | ||
| canceler: canceler, | ||
| deadlineDelay: deadlineDelay, | ||
| } | ||
| } | ||
|
|
||
| // HandleCancel implements ctxwatch.Handler. | ||
| func (h *CancelQueryContextWatcherHandler) HandleCancel(_ context.Context) { | ||
| h.cancelFinished = make(chan struct{}) | ||
| h.cancelSucceeded.Store(false) | ||
|
|
||
| // Fallback: if in-band cancellation fails to resolve the in-flight | ||
| // operation, this deadline unblocks the pending read so the connection | ||
| // errors out and is destroyed. | ||
| _ = h.conn.SetDeadline(time.Now().Add(h.deadlineDelay)) | ||
|
|
||
| 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") | ||
|
Contributor
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
Contributor
There was a problem hiding this comment. Nvm I see you beat me to it |
||
| return | ||
| } | ||
| h.cancelSucceeded.Store(true) | ||
| }() | ||
| } | ||
|
|
||
| // HandleUnwatchAfterCancel implements ctxwatch.Handler. | ||
| func (h *CancelQueryContextWatcherHandler) HandleUnwatchAfterCancel() { | ||
| // Block until the CANCEL statement completes. Releasing the connection | ||
| // any earlier would allow a still-in-flight CANCEL to hit whatever query | ||
| // runs next on this session. | ||
| <-h.cancelFinished | ||
|
|
||
| if h.cancelSucceeded.Load() { | ||
| _ = h.conn.SetDeadline(time.Time{}) | ||
| return | ||
| } | ||
| // The CANCEL statement failed or timed out, so a cancellation may still | ||
| // be in flight for this session. Poison the connection so it is | ||
| // destroyed instead of reused. | ||
| _ = h.conn.SetDeadline(time.Now()) | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is the idea here that there are potentially multiple requests using this same handler? I'm wondering why this isn't in the constructor