diff --git a/CHANGELOG.md b/CHANGELOG.md index 3089eb1ce5..5cb0e36f6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). ## [Unreleased] ### Added - Query Planner: fast serialize/deserialize for query plans (https://github.com/authzed/spicedb/pull/3122) +- CockroachDB: new experimental flag `--datastore-experimental-crdb-query-cancellation` that cancels in-flight queries via CockroachDB's `CANCEL QUERIES` statement when their context is canceled, keeping connections alive instead of destroying them and preventing the pool-drain death spiral under load. Recommended pairing: `--write-conn-acquisition-timeout=0`. ### Changed - Cache: switch to [otter](https://maypok86.github.io/otter/) as the primary cache implementation (https://github.com/authzed/spicedb/pull/3112) diff --git a/docs/spicedb.md b/docs/spicedb.md index 3efc6bd7c5..0bd106008a 100644 --- a/docs/spicedb.md +++ b/docs/spicedb.md @@ -91,6 +91,7 @@ spicedb datastore gc [flags] --datastore-disable-watch-support disable watch support (only enable if you absolutely do not need watch) --datastore-engine string type of datastore to initialize ("cockroachdb", "mysql", "postgres", "spanner") (default "memory") --datastore-experimental-column-optimization enable experimental column optimization (default true) + --datastore-experimental-crdb-query-cancellation enables experimental in-band query cancellation: canceled requests cancel their in-flight queries via CockroachDB's CANCEL QUERIES statement instead of destroying the connection (CockroachDB driver only) --datastore-follower-read-delay-duration duration amount of time to subtract from non-sync revision timestamps to ensure they are sufficiently in the past to enable follower reads (CockroachDB and Spanner drivers only) or read replicas (Postgres and MySQL drivers only) (default 4.8s) --datastore-gc-interval duration amount of time between passes of garbage collection (Postgres driver only) (default 3m0s) --datastore-gc-max-operation-time duration maximum amount of time a garbage collection pass can operate before timing out (Postgres driver only) (default 1m0s) @@ -236,6 +237,7 @@ spicedb datastore repair [flags] --datastore-disable-watch-support disable watch support (only enable if you absolutely do not need watch) --datastore-engine string type of datastore to initialize ("cockroachdb", "mysql", "postgres", "spanner") (default "memory") --datastore-experimental-column-optimization enable experimental column optimization (default true) + --datastore-experimental-crdb-query-cancellation enables experimental in-band query cancellation: canceled requests cancel their in-flight queries via CockroachDB's CANCEL QUERIES statement instead of destroying the connection (CockroachDB driver only) --datastore-follower-read-delay-duration duration amount of time to subtract from non-sync revision timestamps to ensure they are sufficiently in the past to enable follower reads (CockroachDB and Spanner drivers only) or read replicas (Postgres and MySQL drivers only) (default 4.8s) --datastore-gc-interval duration amount of time between passes of garbage collection (Postgres driver only) (default 3m0s) --datastore-gc-max-operation-time duration maximum amount of time a garbage collection pass can operate before timing out (Postgres driver only) (default 1m0s) @@ -414,6 +416,7 @@ spicedb serve [flags] --datastore-disable-watch-support disable watch support (only enable if you absolutely do not need watch) --datastore-engine string type of datastore to initialize ("cockroachdb", "mysql", "postgres", "spanner") (default "memory") --datastore-experimental-column-optimization enable experimental column optimization (default true) + --datastore-experimental-crdb-query-cancellation enables experimental in-band query cancellation: canceled requests cancel their in-flight queries via CockroachDB's CANCEL QUERIES statement instead of destroying the connection (CockroachDB driver only) --datastore-follower-read-delay-duration duration amount of time to subtract from non-sync revision timestamps to ensure they are sufficiently in the past to enable follower reads (CockroachDB and Spanner drivers only) or read replicas (Postgres and MySQL drivers only) (default 4.8s) --datastore-gc-interval duration amount of time between passes of garbage collection (Postgres driver only) (default 3m0s) --datastore-gc-max-operation-time duration maximum amount of time a garbage collection pass can operate before timing out (Postgres driver only) (default 1m0s) diff --git a/internal/datastore/crdb/cancellation_test.go b/internal/datastore/crdb/cancellation_test.go new file mode 100644 index 0000000000..de191e393b --- /dev/null +++ b/internal/datastore/crdb/cancellation_test.go @@ -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") +} diff --git a/internal/datastore/crdb/crdb.go b/internal/datastore/crdb/crdb.go index cbb2f07db8..9ad0196a3e 100644 --- a/internal/datastore/crdb/crdb.go +++ b/internal/datastore/crdb/crdb.go @@ -183,6 +183,7 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas MigrationValidator: common.NewMigrationValidator(headMigration, config.allowedMigrations), dburl: url, acquireTimeout: config.acquireTimeout, + queryCancellationEnabled: config.enableQueryCancellation, watchBufferLength: config.watchBufferLength, watchChangeBufferMaximumSize: config.watchChangeBufferMaximumSize, watchBufferWriteTimeout: config.watchBufferWriteTimeout, @@ -203,6 +204,21 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas // this ctx and cancel is tied to the lifetime of the datastore ds.ctx, ds.cancel = context.WithCancel(context.Background()) + + if config.enableQueryCancellation { + cancelPoolConfig, err := pgxpool.ParseConfig(url) + if err != nil { + ds.cancel() + return nil, common.RedactAndLogSensitiveConnString(ctx, errUnableToInstantiate, err, url) + } + ds.canceler, err = pool.NewCanceler(ds.ctx, cancelPoolConfig) + if err != nil { + ds.cancel() + return nil, common.RedactAndLogSensitiveConnString(ctx, errUnableToInstantiate, err, url) + } + ds.canceler.InstallOn(writePoolConfig) + } + ds.writePool, err = pool.NewRetryPool(ds.ctx, "write", writePoolConfig, healthChecker, config.maxRetries, config.connectRate) if err != nil { ds.cancel() @@ -275,6 +291,8 @@ type crdbDatastore struct { gcWindow time.Duration schema common.SchemaInformation acquireTimeout time.Duration + canceler *pool.Canceler + queryCancellationEnabled bool beginChangefeedQuery string transactionNowQuery string @@ -482,6 +500,9 @@ func (cds *crdbDatastore) Close() error { } cds.readPool.Close() cds.writePool.Close() + if cds.canceler != nil { + cds.canceler.Close() + } for _, collector := range cds.collectors { ok := prometheus.Unregister(collector) if !ok { diff --git a/internal/datastore/crdb/options.go b/internal/datastore/crdb/options.go index 9e745fc213..7f1e133991 100644 --- a/internal/datastore/crdb/options.go +++ b/internal/datastore/crdb/options.go @@ -36,6 +36,7 @@ type crdbOptions struct { includeQueryParametersInTraces bool watchDisabled bool acquireTimeout time.Duration + enableQueryCancellation bool } const ( @@ -67,6 +68,7 @@ const ( defaultColumnOptimizationOption = common.ColumnOptimizationOptionStaticValues defaultIncludeQueryParametersInTraces = false defaultWatchDisabled = false + defaultEnableQueryCancellation = false ) // Option provides the facility to configure how clients within the CRDB @@ -94,6 +96,7 @@ func generateConfig(options []Option) (crdbOptions, error) { includeQueryParametersInTraces: defaultIncludeQueryParametersInTraces, watchDisabled: defaultWatchDisabled, acquireTimeout: defaultAcquireTimeout, + enableQueryCancellation: defaultEnableQueryCancellation, } for _, option := range options { @@ -404,3 +407,13 @@ func WithWatchDisabled(isDisabled bool) Option { func WithAcquireTimeout(timeout time.Duration) Option { return func(po *crdbOptions) { po.acquireTimeout = timeout } } + +// WithQueryCancellation enables in-band cancellation of in-flight queries via +// CockroachDB's CANCEL QUERIES statement when their context is canceled, +// keeping connections alive instead of destroying them (pgx's default +// behavior). When enabled, reader query contexts are no longer severed. +// +// This is experimental and disabled by default. +func WithQueryCancellation(enabled bool) Option { + return func(po *crdbOptions) { po.enableQueryCancellation = enabled } +} diff --git a/internal/datastore/crdb/pool/cancel_handler.go b/internal/datastore/crdb/pool/cancel_handler.go new file mode 100644 index 0000000000..d575e407a1 --- /dev/null +++ b/internal/datastore/crdb/pool/cancel_handler.go @@ -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") + 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()) +} diff --git a/internal/datastore/crdb/pool/cancel_handler_test.go b/internal/datastore/crdb/pool/cancel_handler_test.go new file mode 100644 index 0000000000..7851849ef3 --- /dev/null +++ b/internal/datastore/crdb/pool/cancel_handler_test.go @@ -0,0 +1,98 @@ +package pool + +import ( + "errors" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/require" +) + +type fakeSessionCanceler struct { + delay time.Duration + err error + calls atomic.Int32 +} + +func (f *fakeSessionCanceler) CancelSessionQueries(_ *pgconn.PgConn) error { + f.calls.Add(1) + time.Sleep(f.delay) + return f.err +} + +type fakeDeadlineConn struct { + mu sync.Mutex + deadlines []time.Time // GUARDED_BY(mu) +} + +func (f *fakeDeadlineConn) SetDeadline(t time.Time) error { + f.mu.Lock() + defer f.mu.Unlock() + f.deadlines = append(f.deadlines, t) + return nil +} + +func (f *fakeDeadlineConn) last() time.Time { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.deadlines) == 0 { + return time.Time{} + } + return f.deadlines[len(f.deadlines)-1] +} + +func TestCancelHandlerBlocksUntilCancelCompletesAndClearsDeadline(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + fake := &fakeSessionCanceler{delay: 50 * time.Millisecond} + conn := &fakeDeadlineConn{} + h := &CancelQueryContextWatcherHandler{conn: conn, canceler: fake, deadlineDelay: time.Second} + + h.HandleCancel(t.Context()) + start := time.Now() + h.HandleUnwatchAfterCancel() + + require.GreaterOrEqual(t, time.Since(start), 50*time.Millisecond, + "HandleUnwatchAfterCancel must block until the CANCEL statement completes") + require.Equal(t, int32(1), fake.calls.Load()) + require.True(t, conn.last().IsZero(), "deadline must be cleared after a successful cancel") + }) +} + +func TestCancelHandlerPoisonsConnectionWhenCancelFails(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + fake := &fakeSessionCanceler{err: errors.New("cancel pool exhausted")} + conn := &fakeDeadlineConn{} + h := &CancelQueryContextWatcherHandler{conn: conn, canceler: fake, deadlineDelay: time.Second} + + h.HandleCancel(t.Context()) + h.HandleUnwatchAfterCancel() + + require.False(t, conn.last().IsZero(), "deadline must be poisoned (non-zero) when the cancel fails") + require.False(t, conn.last().After(time.Now()), "poison deadline must already be due") + }) +} + +func TestCancelHandlerResetsStateBetweenCycles(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + fake := &fakeSessionCanceler{} + conn := &fakeDeadlineConn{} + h := &CancelQueryContextWatcherHandler{conn: conn, canceler: fake, deadlineDelay: time.Second} + + // Cycle 1: failure poisons. + fake.err = errors.New("boom") + h.HandleCancel(t.Context()) + h.HandleUnwatchAfterCancel() + require.False(t, conn.last().IsZero()) + + // Cycle 2: success clears. + fake.err = nil + h.HandleCancel(t.Context()) + h.HandleUnwatchAfterCancel() + require.True(t, conn.last().IsZero()) + require.Equal(t, int32(2), fake.calls.Load()) + }) +} diff --git a/internal/datastore/crdb/pool/canceler.go b/internal/datastore/crdb/pool/canceler.go new file mode 100644 index 0000000000..a9f989bd68 --- /dev/null +++ b/internal/datastore/crdb/pool/canceler.go @@ -0,0 +1,186 @@ +package pool + +import ( + "context" + "errors" + "fmt" + "regexp" + "sync" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgconn/ctxwatch" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/prometheus/client_golang/prometheus" +) + +const ( + // defaultCancelTimeout bounds the full round trip of one CANCEL QUERIES + // statement, including acquiring a connection from the cancel pool. + defaultCancelTimeout = 2 * time.Second + + // defaultCancelDeadlineDelay is the fallback net.Conn deadline applied to + // the original connection when its context is canceled. It must + // comfortably exceed defaultCancelTimeout so the in-band cancellation has + // time to resolve the in-flight query before the connection is destroyed. + defaultCancelDeadlineDelay = 5 * time.Second + + // defaultCancelPoolMaxConns bounds concurrent CANCEL QUERIES statements. + // Excess cancellations queue on the pool, applying natural backpressure; + // queued past the statement timeout, they fail safe (the original + // connection is destroyed, the legacy behavior). + defaultCancelPoolMaxConns = 5 +) + +var ( + cancellationDuration = prometheus.NewHistogram(prometheus.HistogramOpts{ + Name: "crdb_query_cancellation_duration_seconds", + Help: "Round-trip time of in-band CANCEL QUERIES statements. This is the window during which the canceled request still holds its connection.", + Buckets: prometheus.ExponentialBuckets(0.0005, 2, 13), + }) + cancellationCounter = prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "crdb_query_cancellations_total", + Help: "Number of in-band CANCEL QUERIES statements issued, by outcome. Outcome 'error' means the connection was discarded instead (legacy behavior).", + }, []string{"outcome"}) +) + +func init() { + prometheus.MustRegister(cancellationDuration, cancellationCounter) +} + +// validSessionID matches CockroachDB session IDs as returned by +// `SHOW session_id` (a hex string). Validated at registration so the ID can +// be safely inlined into the CANCEL QUERIES statement. +var validSessionID = regexp.MustCompile(`^[0-9a-fA-F]+$`) + +// Canceler cancels in-flight queries on tracked CockroachDB sessions using +// the in-band CANCEL QUERIES statement, issued over a small dedicated pool of +// sibling connections. Unlike the pgwire cancel protocol — which CockroachDB +// applies asynchronously after closing the cancel connection — CANCEL QUERIES +// completes only after the target query's context has been canceled on its +// gateway node, and it targets globally-unique query IDs, so it can never +// cancel the wrong query. +type Canceler struct { + pool *pgxpool.Pool + + sync.RWMutex + sessionIDs map[*pgconn.PgConn]string // GUARDED_BY(RWMutex) +} + +// NewCanceler creates a Canceler with its own small connection pool derived +// from the given config. The config must come from pgxpool.ParseConfig. +func NewCanceler(ctx context.Context, config *pgxpool.Config) (*Canceler, error) { + config = config.Copy() + config.MinConns = 1 + config.MaxConns = defaultCancelPoolMaxConns + + p, err := pgxpool.NewWithConfig(ctx, config) + if err != nil { + return nil, fmt.Errorf("unable to create cancel pool: %w", err) + } + return &Canceler{pool: p, sessionIDs: make(map[*pgconn.PgConn]string)}, nil +} + +// RegisterSession captures the CockroachDB session ID of a connection so that +// its queries can later be canceled in-band. Intended to be called from +// pgxpool's AfterConnect hook (see InstallOn). +func (c *Canceler) RegisterSession(ctx context.Context, conn *pgx.Conn) error { + var sessionID string + if err := conn.QueryRow(ctx, "SHOW session_id").Scan(&sessionID); err != nil { + return fmt.Errorf("unable to read session id: %w", err) + } + if !validSessionID.MatchString(sessionID) { + return fmt.Errorf("unexpected session id format: %q", sessionID) + } + c.setSessionID(conn.PgConn(), sessionID) + return nil +} + +func (c *Canceler) setSessionID(pgConn *pgconn.PgConn, sessionID string) { + c.Lock() + defer c.Unlock() + c.sessionIDs[pgConn] = sessionID +} + +// UnregisterSession stops tracking a connection. Intended to be called from +// pgxpool's BeforeClose hook (see InstallOn). +func (c *Canceler) UnregisterSession(pgConn *pgconn.PgConn) { + c.Lock() + defer c.Unlock() + delete(c.sessionIDs, pgConn) +} + +// SessionID returns the tracked CockroachDB session ID for a connection. +func (c *Canceler) SessionID(pgConn *pgconn.PgConn) (string, bool) { + c.RLock() + defer c.RUnlock() + id, ok := c.sessionIDs[pgConn] + return id, ok +} + +// CancelSessionQueries synchronously cancels all queries currently running on +// the session of the given connection. A nil return guarantees the CANCEL +// statement completed on the server: any query that was active on the session +// has had its context canceled (it will observe SQLSTATE 57014), and nothing +// remains in flight that could affect a future query on the session. +func (c *Canceler) CancelSessionQueries(pgConn *pgconn.PgConn) error { + sessionID, ok := c.SessionID(pgConn) + if !ok { + return errors.New("no session registered for connection") + } + + ctx, cancel := context.WithTimeout(context.Background(), defaultCancelTimeout) + defer cancel() + + // The session ID is hex, validated at registration, and inlined rather + // than bound as a placeholder for compatibility with the bracketed SHOW + // subquery syntax. + stmt := fmt.Sprintf( + `CANCEL QUERIES IF EXISTS (SELECT query_id FROM [SHOW CLUSTER STATEMENTS] WHERE session_id = '%s')`, + sessionID, + ) + + start := time.Now() + _, err := c.pool.Exec(ctx, stmt) + cancellationDuration.Observe(time.Since(start).Seconds()) + if err != nil { + cancellationCounter.WithLabelValues("error").Inc() + return fmt.Errorf("unable to cancel queries for session %s: %w", sessionID, err) + } + cancellationCounter.WithLabelValues("ok").Inc() + return nil +} + +// InstallOn wires the canceler into a pool config: connections report their +// session IDs on connect, are untracked on close, and respond to context +// cancellation by canceling their queries in-band instead of being destroyed. +// Must be called before the config is used to construct a pool. +func (c *Canceler) InstallOn(config *pgxpool.Config) { + afterConnect := config.AfterConnect + 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) + } + + beforeClose := config.BeforeClose + config.BeforeClose = func(conn *pgx.Conn) { + if beforeClose != nil { + beforeClose(conn) + } + c.UnregisterSession(conn.PgConn()) + } + + config.ConnConfig.BuildContextWatcherHandler = func(pgConn *pgconn.PgConn) ctxwatch.Handler { + return NewCancelQueryContextWatcherHandler(pgConn, c, defaultCancelDeadlineDelay) + } +} + +// Close closes the cancel pool. +func (c *Canceler) Close() { + c.pool.Close() +} diff --git a/internal/datastore/crdb/pool/canceler_test.go b/internal/datastore/crdb/pool/canceler_test.go new file mode 100644 index 0000000000..7bb5330e7e --- /dev/null +++ b/internal/datastore/crdb/pool/canceler_test.go @@ -0,0 +1,143 @@ +package pool + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/require" +) + +func TestCancelerSessionRegistry(t *testing.T) { + c := &Canceler{sessionIDs: make(map[*pgconn.PgConn]string)} + pgc := &pgconn.PgConn{} + + _, ok := c.SessionID(pgc) + require.False(t, ok) + + c.setSessionID(pgc, "deadbeef01234567deadbeef01234567") + id, ok := c.SessionID(pgc) + require.True(t, ok) + require.Equal(t, "deadbeef01234567deadbeef01234567", id) + + c.UnregisterSession(pgc) + _, ok = c.SessionID(pgc) + require.False(t, ok) +} + +func TestCancelSessionQueriesUnknownSession(t *testing.T) { + c := &Canceler{sessionIDs: make(map[*pgconn.PgConn]string)} + err := c.CancelSessionQueries(&pgconn.PgConn{}) + require.Error(t, err, "unknown sessions must error so the handler poisons the connection") +} + +func TestSessionIDValidation(t *testing.T) { + require.True(t, validSessionID.MatchString("17a4c5b2e9d8f0a117a4c5b2e9d8f0a1")) + require.False(t, validSessionID.MatchString("")) + require.False(t, validSessionID.MatchString("17a4'; CANCEL SESSIONS --")) +} + +// TestCancelerInstallOnSetsHandlers verifies that InstallOn wires all three +// hooks (AfterConnect, BeforeClose, BuildContextWatcherHandler) onto a config +// that started with none. +func TestCancelerInstallOnSetsHandlers(t *testing.T) { + c := &Canceler{sessionIDs: make(map[*pgconn.PgConn]string)} + + config, err := pgxpool.ParseConfig("postgres://localhost/test") + require.NoError(t, err) + require.Nil(t, config.AfterConnect) + require.Nil(t, config.BeforeClose) + + c.InstallOn(config) + + require.NotNil(t, config.AfterConnect) + require.NotNil(t, config.BeforeClose) + require.NotNil(t, config.ConnConfig.BuildContextWatcherHandler) + + // Invoke BuildContextWatcherHandler to cover its body. + handler := config.ConnConfig.BuildContextWatcherHandler(&pgconn.PgConn{}) + require.NotNil(t, handler) +} + +// TestCancelerInstallOnComposesExistingAfterConnect verifies that an existing +// AfterConnect hook is called before RegisterSession, and that an error from +// that hook is propagated without calling RegisterSession. +func TestCancelerInstallOnComposesExistingAfterConnect(t *testing.T) { + c := &Canceler{sessionIDs: make(map[*pgconn.PgConn]string)} + + config, err := pgxpool.ParseConfig("postgres://localhost/test") + require.NoError(t, err) + + var called bool + config.AfterConnect = func(_ context.Context, _ *pgx.Conn) error { + called = true + return errors.New("existing hook error") + } + + c.InstallOn(config) + + // The existing hook returns an error; RegisterSession must not be reached. + err = config.AfterConnect(t.Context(), nil) + require.True(t, called, "existing AfterConnect must be invoked") + require.EqualError(t, err, "existing hook error") +} + +// TestCancelerInstallOnComposesExistingBeforeClose verifies that an existing +// BeforeClose hook is called as part of the composed hook. +func TestCancelerInstallOnComposesExistingBeforeClose(t *testing.T) { + c := &Canceler{sessionIDs: make(map[*pgconn.PgConn]string)} + + config, err := pgxpool.ParseConfig("postgres://localhost/test") + require.NoError(t, err) + + var called bool + config.BeforeClose = func(_ *pgx.Conn) { called = true } + + c.InstallOn(config) + + // &pgx.Conn{} has a nil pgConn field; PgConn() returns nil, and + // UnregisterSession(nil) is a safe no-op (delete from map with nil key). + config.BeforeClose(&pgx.Conn{}) + require.True(t, called, "existing BeforeClose must be invoked") +} + +// TestNewCancelerAndClose verifies that NewCanceler returns a valid Canceler +// and that Close drains the pool without panicking. The cancelled context +// stops the pool's background connection goroutines immediately so the test +// does not leak goroutines or attempt real TCP connections. +func TestNewCancelerAndClose(t *testing.T) { + config, err := pgxpool.ParseConfig("postgres://db:password@127.0.0.1:1/nonexistent") + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + canceler, err := NewCanceler(ctx, config) + require.NoError(t, err) + require.NotNil(t, canceler) + canceler.Close() +} + +// TestCancelerCancelSessionQueriesErrorsOnClosedPool verifies that +// CancelSessionQueries returns a wrapped error (and records the "error" +// outcome counter) when the underlying pool can no longer issue queries. +func TestCancelerCancelSessionQueriesErrorsOnClosedPool(t *testing.T) { + config, err := pgxpool.ParseConfig("postgres://db:password@127.0.0.1:1/nonexistent") + require.NoError(t, err) + config.MinConns = 0 + + p, err := pgxpool.NewWithConfig(t.Context(), config) + require.NoError(t, err) + p.Close() + + pgc := &pgconn.PgConn{} + c := &Canceler{pool: p, sessionIDs: make(map[*pgconn.PgConn]string)} + c.setSessionID(pgc, "deadbeef01234567deadbeef01234567") + + err = c.CancelSessionQueries(pgc) + require.Error(t, err) + require.Contains(t, err.Error(), "unable to cancel queries") +} diff --git a/internal/datastore/crdb/pool/pgx.go b/internal/datastore/crdb/pool/pgx.go index 15fed1d1c9..2377892ca6 100644 --- a/internal/datastore/crdb/pool/pgx.go +++ b/internal/datastore/crdb/pool/pgx.go @@ -33,11 +33,14 @@ import ( // OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION // WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -// beginFuncExec is copied directly from pgx +// beginFuncExec is adapted from pgx // src: https://github.com/jackc/pgx/blob/f59e8bf5551f403e6b7ec0912097bce85ea21351/tx.go#L410-L425 +// It diverges from upstream by severing the context for rollbacks: the +// surrounding context is often already canceled when rollback runs, and the +// rollback must still reach the server to keep the connection reusable. 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)) if rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) { err = rollbackErr } @@ -45,7 +48,7 @@ func beginFuncExec(ctx context.Context, tx pgx.Tx, fn func(pgx.Tx) error) (err e fErr := fn(tx) if fErr != nil { - _ = tx.Rollback(ctx) // ignore rollback error as there is already an error to return + _ = tx.Rollback(context.WithoutCancel(ctx)) // ignore rollback error as there is already an error to return return fErr } diff --git a/internal/datastore/crdb/pool/pgx_test.go b/internal/datastore/crdb/pool/pgx_test.go new file mode 100644 index 0000000000..cb6f6e4d28 --- /dev/null +++ b/internal/datastore/crdb/pool/pgx_test.go @@ -0,0 +1,43 @@ +package pool + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" +) + +// recordingTx records the context passed to Rollback. All other pgx.Tx +// methods panic via the embedded nil interface; beginFuncExec only calls +// Rollback and Commit. +type recordingTx struct { + pgx.Tx + rollbackCtxs []context.Context +} + +func (r *recordingTx) Rollback(ctx context.Context) error { + r.rollbackCtxs = append(r.rollbackCtxs, ctx) + return nil +} + +func (r *recordingTx) Commit(_ context.Context) error { + return nil +} + +func TestBeginFuncExecRollsBackWithSeveredContext(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + tx := &recordingTx{} + err := beginFuncExec(ctx, tx, func(pgx.Tx) error { + return errors.New("txn failed") + }) + require.Error(t, err) + + require.NotEmpty(t, tx.rollbackCtxs) + for _, rctx := range tx.rollbackCtxs { + require.NoError(t, rctx.Err(), "rollback must run with a severed context so it reaches the server") + } +} diff --git a/internal/datastore/crdb/pool/pool.go b/internal/datastore/crdb/pool/pool.go index 766de31e9a..2ec7feddb5 100644 --- a/internal/datastore/crdb/pool/pool.go +++ b/internal/datastore/crdb/pool/pool.go @@ -34,8 +34,18 @@ var resetHistogram = prometheus.NewHistogram(prometheus.HistogramOpts{ Buckets: []float64{0, 1, 2, 5, 10, 20, 50}, }) +// UnexpectedCancellationErrors counts queries that observed SQLSTATE 57014 +// (query canceled) although their own context was never canceled. With +// in-band query cancellation enabled this must remain zero: a nonzero value +// means a cancellation landed on a query it was not intended for. +var UnexpectedCancellationErrors = prometheus.NewCounter(prometheus.CounterOpts{ + Name: "crdb_unexpected_cancellation_errors_total", + Help: "Number of queries that received SQLSTATE 57014 (query canceled) although their own context was not canceled.", +}) + func init() { prometheus.MustRegister(resetHistogram) + prometheus.MustRegister(UnexpectedCancellationErrors) } type ctxDisableRetries struct{} @@ -452,6 +462,13 @@ func wrapRetryableError(ctx context.Context, err error) error { return err } + if sqlErrorCode(err) == CrdbQueryCanceledErrCode { + // The context checks above did not trigger, so this query was + // canceled without its own context being canceled. + UnexpectedCancellationErrors.Inc() + log.Ctx(ctx).Warn().Err(err).Msg("query canceled error received without a canceled context") + } + if IsResettableError(ctx, err) { return &ResettableError{Err: err} } diff --git a/internal/datastore/crdb/pool/sqlerrors.go b/internal/datastore/crdb/pool/sqlerrors.go index ca66c0e426..fecbe2d781 100644 --- a/internal/datastore/crdb/pool/sqlerrors.go +++ b/internal/datastore/crdb/pool/sqlerrors.go @@ -14,6 +14,9 @@ import ( const ( // https://www.cockroachlabs.com/docs/stable/common-errors.html#restart-transaction CrdbRetryErrCode = "40001" + // CrdbQueryCanceledErrCode is the SQLSTATE returned when a query is + // canceled, e.g. via CANCEL QUERIES or the pgwire cancel protocol. + CrdbQueryCanceledErrCode = "57014" // https://www.cockroachlabs.com/docs/stable/common-errors.html#result-is-ambiguous CrdbAmbiguousErrorCode = "40003" // https://www.cockroachlabs.com/docs/stable/node-shutdown.html#connection-retry-loop diff --git a/internal/datastore/crdb/pool/sqlerrors_test.go b/internal/datastore/crdb/pool/sqlerrors_test.go index 1a384933fb..e404af0237 100644 --- a/internal/datastore/crdb/pool/sqlerrors_test.go +++ b/internal/datastore/crdb/pool/sqlerrors_test.go @@ -1,11 +1,14 @@ package pool import ( + "context" "fmt" "reflect" "testing" "github.com/jackc/pgx/v5/pgconn" + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -447,3 +450,19 @@ func Test_sqlErrorCode(t *testing.T) { }) } } + +func TestUnexpectedCancellationTripwire(t *testing.T) { + before := testutil.ToFloat64(UnexpectedCancellationErrors) + + // A 57014 with a live context is a cancellation that hit the wrong query. + err := wrapRetryableError(t.Context(), &pgconn.PgError{Code: CrdbQueryCanceledErrCode}) + require.Error(t, err) + require.InEpsilon(t, before+1, testutil.ToFloat64(UnexpectedCancellationErrors), 1e-9) + + // A 57014 with a canceled context is expected and must not trip the wire. + ctx, cancel := context.WithCancel(t.Context()) + cancel() + err = wrapRetryableError(ctx, &pgconn.PgError{Code: CrdbQueryCanceledErrCode}) + require.ErrorIs(t, err, context.Canceled) + require.InEpsilon(t, before+1, testutil.ToFloat64(UnexpectedCancellationErrors), 1e-9) +} diff --git a/pkg/cmd/datastore/datastore.go b/pkg/cmd/datastore/datastore.go index 18b223c1c0..38e37c20ba 100644 --- a/pkg/cmd/datastore/datastore.go +++ b/pkg/cmd/datastore/datastore.go @@ -142,13 +142,14 @@ type Config struct { RequestHedgingQuantile float64 `debugmap:"visible"` // CRDB - FollowerReadDelay time.Duration `debugmap:"visible" default:"4800ms"` - MaxRetries int `debugmap:"visible" default:"10"` - OverlapKey string `debugmap:"visible" default:"key"` - OverlapStrategy string `debugmap:"visible" default:"static"` - EnableConnectionBalancing bool `debugmap:"visible" default:"true"` - ConnectRate time.Duration `debugmap:"visible" default:"100ms"` - WriteAcquisitionTimeout time.Duration `debugmap:"visible" default:"30ms"` + FollowerReadDelay time.Duration `debugmap:"visible" default:"4800ms"` + MaxRetries int `debugmap:"visible" default:"10"` + OverlapKey string `debugmap:"visible" default:"key"` + OverlapStrategy string `debugmap:"visible" default:"static"` + EnableConnectionBalancing bool `debugmap:"visible" default:"true"` + ConnectRate time.Duration `debugmap:"visible" default:"100ms"` + WriteAcquisitionTimeout time.Duration `debugmap:"visible" default:"30ms"` + ExperimentalCRDBQueryCancellation bool `debugmap:"visible"` // Postgres GCInterval time.Duration `debugmap:"visible" default:"3m"` @@ -365,6 +366,7 @@ func RegisterDatastoreFlagsWithPrefix(flagSet *pflag.FlagSet, prefix string, opt flagSet.BoolVar(&opts.DisableWatchSupport, flagName("datastore-disable-watch-support"), false, "disable watch support (only enable if you absolutely do not need watch)") flagSet.BoolVar(&opts.IncludeQueryParametersInTraces, flagName("datastore-include-query-parameters-in-traces"), false, "include query parameters in traces (Postgres and CockroachDB drivers only)") flagSet.DurationVar(&opts.WriteAcquisitionTimeout, flagName("write-conn-acquisition-timeout"), defaults.WriteAcquisitionTimeout, "amount of time that the server will wait for a connection to the datastore to become available when performing a write operation before throwing a ResourceExhausted error. 0 means wait indefinitely. (CockroachDB driver only)") + flagSet.BoolVar(&opts.ExperimentalCRDBQueryCancellation, flagName("datastore-experimental-crdb-query-cancellation"), defaults.ExperimentalCRDBQueryCancellation, "enables experimental in-band query cancellation: canceled requests cancel their in-flight queries via CockroachDB's CANCEL QUERIES statement instead of destroying the connection (CockroachDB driver only)") flagSet.BoolVar(&opts.RelationshipIntegrityEnabled, flagName("datastore-relationship-integrity-enabled"), false, "enables relationship integrity checks. (CockroachDB driver only)") flagSet.StringVar(&opts.RelationshipIntegrityCurrentKey.KeyID, flagName("datastore-relationship-integrity-current-key-id"), "", "current key id for relationship integrity checks") @@ -399,51 +401,52 @@ func RegisterDatastoreFlagsWithPrefix(flagSet *pflag.FlagSet, prefix string, opt func DefaultDatastoreConfig() *Config { return &Config{ - Engine: MemoryEngine, - GCWindow: 24 * time.Hour, - LegacyFuzzing: -1, - RevisionQuantization: 5 * time.Second, - MaxRevisionStalenessPercent: .1, // 10% - ReadConnPool: *DefaultReadConnPool(), - WriteConnPool: *DefaultWriteConnPool(), - ReadReplicaConnPool: *DefaultReadConnPool(), - OldReadReplicaConnPool: *DefaultReadConnPool(), - ReadReplicaURIs: []string{}, - ReadOnly: false, - MaxRetries: 10, - OverlapKey: "key", - OverlapStrategy: "static", - ConnectRate: 100 * time.Millisecond, - EnableConnectionBalancing: true, - GCInterval: 3 * time.Minute, - GCMaxOperationTime: 1 * time.Minute, - WatchBufferLength: 1024, - WatchChangeBufferMaximumSize: "15%", - WatchBufferWriteTimeout: 1 * time.Second, - WatchConnectTimeout: 1 * time.Second, - DisableWatchSupport: false, - EnableDatastoreMetrics: true, - DisableStats: false, - BootstrapFiles: []string{}, - BootstrapTimeout: 10 * time.Second, - BootstrapOverwrite: false, - SpannerCredentialsFile: "", - SpannerEmulatorHost: "", - TablePrefix: "", - MigrationPhase: "", - FollowerReadDelay: DefaultFollowerReadDelay, - SpannerMinSessions: 100, - SpannerMaxSessions: 400, - FilterMaximumIDCount: 100, - SpannerDatastoreMetricsOption: spanner.DatastoreMetricsOptionOpenTelemetry, - RelationshipIntegrityEnabled: false, - RelationshipIntegrityCurrentKey: RelIntegrityKey{}, - RelationshipIntegrityExpiredKeys: []string{}, - AllowedMigrations: []string{}, - ExperimentalColumnOptimization: true, - IncludeQueryParametersInTraces: false, - WriteAcquisitionTimeout: 30 * time.Millisecond, - CaveatTypeSet: caveattypes.Default.TypeSet, + Engine: MemoryEngine, + GCWindow: 24 * time.Hour, + LegacyFuzzing: -1, + RevisionQuantization: 5 * time.Second, + MaxRevisionStalenessPercent: .1, // 10% + ReadConnPool: *DefaultReadConnPool(), + WriteConnPool: *DefaultWriteConnPool(), + ReadReplicaConnPool: *DefaultReadConnPool(), + OldReadReplicaConnPool: *DefaultReadConnPool(), + ReadReplicaURIs: []string{}, + ReadOnly: false, + MaxRetries: 10, + OverlapKey: "key", + OverlapStrategy: "static", + ConnectRate: 100 * time.Millisecond, + EnableConnectionBalancing: true, + GCInterval: 3 * time.Minute, + GCMaxOperationTime: 1 * time.Minute, + WatchBufferLength: 1024, + WatchChangeBufferMaximumSize: "15%", + WatchBufferWriteTimeout: 1 * time.Second, + WatchConnectTimeout: 1 * time.Second, + DisableWatchSupport: false, + EnableDatastoreMetrics: true, + DisableStats: false, + BootstrapFiles: []string{}, + BootstrapTimeout: 10 * time.Second, + BootstrapOverwrite: false, + SpannerCredentialsFile: "", + SpannerEmulatorHost: "", + TablePrefix: "", + MigrationPhase: "", + FollowerReadDelay: DefaultFollowerReadDelay, + SpannerMinSessions: 100, + SpannerMaxSessions: 400, + FilterMaximumIDCount: 100, + SpannerDatastoreMetricsOption: spanner.DatastoreMetricsOptionOpenTelemetry, + RelationshipIntegrityEnabled: false, + RelationshipIntegrityCurrentKey: RelIntegrityKey{}, + RelationshipIntegrityExpiredKeys: []string{}, + AllowedMigrations: []string{}, + ExperimentalColumnOptimization: true, + IncludeQueryParametersInTraces: false, + WriteAcquisitionTimeout: 30 * time.Millisecond, + ExperimentalCRDBQueryCancellation: false, + CaveatTypeSet: caveattypes.Default.TypeSet, } } @@ -620,6 +623,7 @@ func newCRDBDatastore(ctx context.Context, opts Config) (datastore.Datastore, er crdb.ReadConnMaxLifetimeJitter(opts.ReadConnPool.MaxLifetimeJitter), crdb.ReadConnHealthCheckInterval(opts.ReadConnPool.HealthCheckInterval), crdb.WithAcquireTimeout(opts.WriteAcquisitionTimeout), + crdb.WithQueryCancellation(opts.ExperimentalCRDBQueryCancellation), crdb.WriteConnsMaxOpen(opts.WriteConnPool.MaxOpenConns), crdb.WriteConnsMinOpen(opts.WriteConnPool.MinOpenConns), crdb.WriteConnMaxIdleTime(opts.WriteConnPool.MaxIdleTime), diff --git a/pkg/cmd/datastore/zz_generated.options.go b/pkg/cmd/datastore/zz_generated.options.go index 9d814ebf72..1904f7b8f4 100644 --- a/pkg/cmd/datastore/zz_generated.options.go +++ b/pkg/cmd/datastore/zz_generated.options.go @@ -66,6 +66,7 @@ func (c *Config) ToOption() ConfigOption { to.EnableConnectionBalancing = c.EnableConnectionBalancing to.ConnectRate = c.ConnectRate to.WriteAcquisitionTimeout = c.WriteAcquisitionTimeout + to.ExperimentalCRDBQueryCancellation = c.ExperimentalCRDBQueryCancellation to.GCInterval = c.GCInterval to.GCMaxOperationTime = c.GCMaxOperationTime to.RelaxedIsolationLevel = c.RelaxedIsolationLevel @@ -231,6 +232,7 @@ func (c *Config) DebugMap() map[string]any { } else { debugMap["WriteAcquisitionTimeout"] = c.WriteAcquisitionTimeout } + debugMap["ExperimentalCRDBQueryCancellation"] = c.ExperimentalCRDBQueryCancellation if dm, ok := any(&c.GCInterval).(interface { DebugMap() map[string]any }); ok { @@ -613,6 +615,13 @@ func WithWriteAcquisitionTimeout(writeAcquisitionTimeout time.Duration) ConfigOp } } +// WithExperimentalCRDBQueryCancellation returns an option that can set ExperimentalCRDBQueryCancellation on a Config +func WithExperimentalCRDBQueryCancellation(experimentalCRDBQueryCancellation bool) ConfigOption { + return func(c *Config) { + c.ExperimentalCRDBQueryCancellation = experimentalCRDBQueryCancellation + } +} + // WithGCInterval returns an option that can set GCInterval on a Config func WithGCInterval(gCInterval time.Duration) ConfigOption { return func(c *Config) {