chore: add retries to read path in postgres datastore impl

internal/datastore/postgres/postgres.go

@@ -6,6 +6,7 @@
 	"errors"
 	"fmt"
 	"math/rand/v2"
+	"net"
 	"os"
 	"strconv"
 	"sync/atomic"
@@ -674,6 +675,31 @@
 	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 {

internal/datastore/postgres/revisions.go

@@ -13,7 +13,6 @@
 	"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 @@
 	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 @@
 	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 @@
 	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 @@
 
 	var minXid xid8
 	var minSnapshot, currentSnapshot pgSnapshot
-	if err := pgd.readPool.QueryRow(ctx, pgd.validTransactionQuery).
-		Scan(&minXid, &minSnapshot, &currentSnapshot); err != nil {
+	if err := pgd.queryRowWithRetries(ctx, pgd.validTransactionQuery, func(row pgx.Row) error {
+		return row.Scan(&minXid, &minSnapshot, &currentSnapshot)
+	}); err != nil {
 		return fmt.Errorf(errCheckRevision, err)
 	}
 

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))
+		})
+	}
+}