From 1af38ab3d3e0e7058f273bc21350304a1c9160ff Mon Sep 17 00:00:00 2001 From: Tanner Stirrat Date: Fri, 10 Jul 2026 08:16:42 -0600 Subject: [PATCH] chore: add retries to read path in postgres datastore impl --- internal/datastore/postgres/postgres.go | 26 ++++++++++++++ internal/datastore/postgres/revisions.go | 36 ++++++++++++++----- internal/datastore/postgres/revisions_test.go | 35 ++++++++++++++++++ 3 files changed, 89 insertions(+), 8 deletions(-) diff --git a/internal/datastore/postgres/postgres.go b/internal/datastore/postgres/postgres.go index 8ac1fbf88e..60f2385db6 100644 --- a/internal/datastore/postgres/postgres.go +++ b/internal/datastore/postgres/postgres.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "math/rand/v2" + "net" "os" "strconv" "sync/atomic" @@ -674,6 +675,31 @@ func errorRetryable(err error) bool { return false } +// readErrorRetryable returns whether an error from an idempotent, read-only query may +// be retried. Unlike errorRetryable, it also covers connection-level failures such as +// i/o timeouts: those surface when a pooled connection is torn down underneath the +// query, and pgx discards the connection afterward, so a retry acquires a fresh one. +// It never retries once the calling context itself is done. +func readErrorRetryable(ctx context.Context, err error) bool { + if ctx.Err() != nil { + return false + } + + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + + if pgconn.SafeToRetry(err) { + return true + } + + if netErr, ok := errors.AsType[net.Error](err); ok && netErr.Timeout() { + return true + } + + return common.IsResettableError(err) +} + func (pgd *pgDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { pgDriver, err := migrations.NewAlembicPostgresDriver(ctx, pgd.dburl, pgd.credentialsProvider, pgd.includeQueryParametersInTraces) if err != nil { diff --git a/internal/datastore/postgres/revisions.go b/internal/datastore/postgres/revisions.go index b800346a4a..7d29de6b8e 100644 --- a/internal/datastore/postgres/revisions.go +++ b/internal/datastore/postgres/revisions.go @@ -13,7 +13,6 @@ import ( "github.com/jackc/pgx/v5" dscommon "github.com/authzed/spicedb/internal/datastore/common" - "github.com/authzed/spicedb/internal/datastore/postgres/common" "github.com/authzed/spicedb/internal/datastore/postgres/schema" "github.com/authzed/spicedb/pkg/datastore" implv1 "github.com/authzed/spicedb/pkg/proto/impl/v1" @@ -128,13 +127,31 @@ const ( queryLatestXID = `SELECT max(xid)::text::integer FROM relation_tuple_transaction;` ) +// maxTransientReadRetries bounds the retries performed by queryRowWithRetries. +const maxTransientReadRetries = 3 + +// queryRowWithRetries runs a single-row query on the read pool, retrying a bounded +// number of times when the failure is a transient connection error (e.g. a pooled +// connection torn down underneath the query). Callers must only pass idempotent, +// read-only queries. +func (pgd *pgDatastore) queryRowWithRetries(ctx context.Context, sql string, scanFn func(row pgx.Row) error) error { + for retries := uint8(0); ; retries++ { + err := scanFn(pgd.readPool.QueryRow(ctx, sql)) + if err == nil || retries >= maxTransientReadRetries || !readErrorRetryable(ctx, err) { + return err + } + dscommon.SleepOnErr(ctx, err, retries) + } +} + func (pgd *pgDatastore) optimizedRevisionFunc(ctx context.Context) (datastore.Revision, time.Duration, string, error) { var revision xid8 var snapshot pgSnapshot var validForNanos time.Duration var schemaHash []byte - if err := pgd.readPool.QueryRow(ctx, pgd.optimizedRevisionQuery). - Scan(&revision, &snapshot, &validForNanos, &schemaHash); err != nil { + if err := pgd.queryRowWithRetries(ctx, pgd.optimizedRevisionQuery, func(row pgx.Row) error { + return row.Scan(&revision, &snapshot, &validForNanos, &schemaHash) + }); err != nil { return datastore.NoRevision, 0, "", fmt.Errorf(errRevision, err) } @@ -147,7 +164,7 @@ func (pgd *pgDatastore) HeadRevision(ctx context.Context) (datastore.RevisionWit ctx, span := tracer.Start(ctx, "HeadRevision") defer span.End() - result, schemaHash, err := pgd.getHeadRevisionWithHash(ctx, pgd.readPool) + result, schemaHash, err := pgd.getHeadRevisionWithHash(ctx) if err != nil { return datastore.RevisionWithSchemaHash{}, err } @@ -158,10 +175,12 @@ func (pgd *pgDatastore) HeadRevision(ctx context.Context) (datastore.RevisionWit return datastore.RevisionWithSchemaHash{Revision: *result, SchemaHash: string(schemaHash)}, nil } -func (pgd *pgDatastore) getHeadRevisionWithHash(ctx context.Context, querier common.Querier) (*postgresRevision, []byte, error) { +func (pgd *pgDatastore) getHeadRevisionWithHash(ctx context.Context) (*postgresRevision, []byte, error) { var snapshot pgSnapshot var schemaHash []byte - if err := querier.QueryRow(ctx, queryCurrentSnapshotWithHash).Scan(&snapshot, &schemaHash); err != nil { + if err := pgd.queryRowWithRetries(ctx, queryCurrentSnapshotWithHash, func(row pgx.Row) error { + return row.Scan(&snapshot, &schemaHash) + }); err != nil { if errors.Is(err, pgx.ErrNoRows) { return nil, nil, nil } @@ -180,8 +199,9 @@ func (pgd *pgDatastore) CheckRevision(ctx context.Context, revisionRaw datastore var minXid xid8 var minSnapshot, currentSnapshot pgSnapshot - if err := pgd.readPool.QueryRow(ctx, pgd.validTransactionQuery). - Scan(&minXid, &minSnapshot, ¤tSnapshot); err != nil { + if err := pgd.queryRowWithRetries(ctx, pgd.validTransactionQuery, func(row pgx.Row) error { + return row.Scan(&minXid, &minSnapshot, ¤tSnapshot) + }); err != nil { return fmt.Errorf(errCheckRevision, err) } diff --git a/internal/datastore/postgres/revisions_test.go b/internal/datastore/postgres/revisions_test.go index 6047488331..0a738cf884 100644 --- a/internal/datastore/postgres/revisions_test.go +++ b/internal/datastore/postgres/revisions_test.go @@ -1,12 +1,17 @@ package postgres import ( + "context" + "errors" "fmt" + "net" + "os" "strconv" "testing" "time" "github.com/ccoveille/go-safecast/v2" + "github.com/jackc/pgx/v5" "github.com/stretchr/testify/require" ) @@ -217,3 +222,33 @@ func FuzzRevision(f *testing.F) { } }) } + +func TestReadErrorRetryable(t *testing.T) { + canceledCtx, cancel := context.WithCancel(t.Context()) + cancel() + + ioTimeout := &net.OpError{Op: "read", Net: "tcp", Err: os.ErrDeadlineExceeded} + + testCases := []struct { + name string + ctx context.Context + err error + retryable bool + }{ + {"io timeout", t.Context(), ioTimeout, true}, + {"wrapped io timeout", t.Context(), fmt.Errorf("unable to find revision: %w", ioTimeout), true}, + {"connection reset", t.Context(), errors.New("connection reset by peer"), true}, + {"context canceled", t.Context(), context.Canceled, false}, + {"context deadline exceeded", t.Context(), context.DeadlineExceeded, false}, + {"wrapped context deadline exceeded", t.Context(), fmt.Errorf("query failed: %w", context.DeadlineExceeded), false}, + {"caller context done", canceledCtx, ioTimeout, false}, + {"no rows", t.Context(), pgx.ErrNoRows, false}, + {"generic error", t.Context(), errors.New("syntax error"), false}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.retryable, readErrorRetryable(tc.ctx, tc.err)) + }) + } +}