fix(db): back off serializable transaction retries

GitOrigin-RevId: 76591f66b30fb229ca506daacc33a820cadd82b1

Author: gm-e2b

packages/db/pkg/pool/session.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math/rand/v2"
 	"time"
 
 	"github.com/jackc/pgerrcode"
@@ -13,9 +14,11 @@ import (
 )
 
 const (
-	sessionReleaseTimeout = 5 * time.Second
-	serializableAttempts  = 10
-	serializableRetryStep = 5 * time.Millisecond
+	sessionReleaseTimeout    = 5 * time.Second
+	serializableAttempts     = 10
+	serializableRetryBase    = 10 * time.Millisecond
+	serializableRetryMax     = 500 * time.Millisecond
+	serializableRetryTimeout = 5 * time.Second
 )
 
 var (
@@ -92,6 +95,7 @@ func (l *AdvisoryLock) InSerializableTx(ctx context.Context, fn func(context.Con
 // InSerializableTxReturn1 runs fn in a SERIALIZABLE transaction on the session
 // that holds the lock. Serialization failures and deadlocks replay the whole
 // callback; only the value produced by the attempt that commits is returned.
+// Attempts and backoff share a five-second budget, bounded by the caller's deadline.
 func InSerializableTxReturn1[T any](
 	ctx context.Context,
 	lock *AdvisoryLock,
@@ -102,8 +106,15 @@ func InSerializableTxReturn1[T any](
 		return zero, errLockReleased
 	}
 
+	ctx, cancel := context.WithTimeout(ctx, serializableRetryTimeout)
+	defer cancel()
+
 	var conflict error
 	for attempt := 1; attempt <= serializableAttempts; attempt++ {
+		if err := ctx.Err(); err != nil {
+			return zero, errors.Join(err, conflict)
+		}
+
 		value, err := runInTxReturn1(ctx, lock.conn, fn)
 		switch {
 		case err == nil:
@@ -165,7 +176,7 @@ func isSerializationConflict(err error) bool {
 }
 
 func waitToReplay(ctx context.Context, attempt int) error {
-	timer := time.NewTimer(serializableRetryStep * time.Duration(attempt))
+	timer := time.NewTimer(serializableRetryDelay(attempt))
 	defer timer.Stop()
 
 	select {
@@ -176,6 +187,14 @@ func waitToReplay(ctx context.Context, attempt int) error {
 	}
 }
 
+func serializableRetryDelay(attempt int) time.Duration {
+	// Equal jitter keeps retries apart without allowing an immediate replay.
+	// attempt is bounded by serializableAttempts, so the shift cannot overflow.
+	window := min(2*serializableRetryBase<<(attempt-1), serializableRetryMax)
+
+	return window/2 + rand.N(window/2)
+}
+
 // Release unlocks the session and returns its connection to the pool. It is
 // idempotent. If the unlock result is unknown, Release destroys the session so
 // a connection carrying an unknown lock cannot re-enter the pool.

packages/db/pkg/pool/session_integration_test.go

@@ -18,6 +18,36 @@ import (
 
 const testPostgresImage = "postgres:18-alpine"
 
+func TestSerializableRetryDelay(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		attempt int
+		minimum time.Duration
+		maximum time.Duration
+	}{
+		{1, 10 * time.Millisecond, 20 * time.Millisecond},
+		{2, 20 * time.Millisecond, 40 * time.Millisecond},
+		{3, 40 * time.Millisecond, 80 * time.Millisecond},
+		{5, 160 * time.Millisecond, 320 * time.Millisecond},
+		{6, 250 * time.Millisecond, 500 * time.Millisecond},
+		{9, 250 * time.Millisecond, 500 * time.Millisecond},
+	} {
+		t.Run(fmt.Sprintf("attempt-%d", tc.attempt), func(t *testing.T) {
+			t.Parallel()
+
+			delays := make(map[time.Duration]struct{})
+			for range 32 {
+				delay := serializableRetryDelay(tc.attempt)
+				require.GreaterOrEqual(t, delay, tc.minimum)
+				require.Less(t, delay, tc.maximum)
+				delays[delay] = struct{}{}
+			}
+			assert.Greater(t, len(delays), 1, "retries must not use a deterministic delay")
+		})
+	}
+}
+
 func TestConnectChecksAndConfiguresPool(t *testing.T) {
 	t.Parallel()
 
@@ -96,6 +126,9 @@ func TestAdvisoryLockRetriesSerializableTransactionOnItsSession(t *testing.T) {
 	attempts := 0
 	result, err := InSerializableTxReturn1(t.Context(), lock, func(ctx context.Context, tx pgx.Tx) (transactionResult, error) {
 		attempts++
+		deadline, ok := ctx.Deadline()
+		require.True(t, ok, "transaction attempts must have a deadline")
+		assert.LessOrEqual(t, time.Until(deadline), 5*time.Second)
 
 		var value int
 		var backend int
@@ -194,6 +227,17 @@ func TestSerializableTransactionCleansUpAndBoundsReplays(t *testing.T) {
 	).Scan(&value))
 	assert.Equal(t, 3, value)
 
+	deadline := time.Now().Add(time.Second)
+	shorter, cancelShorter := context.WithDeadline(t.Context(), deadline)
+	defer cancelShorter()
+	require.NoError(t, lock.InSerializableTx(shorter, func(ctx context.Context, _ pgx.Tx) error {
+		actual, ok := ctx.Deadline()
+		require.True(t, ok)
+		assert.Equal(t, deadline, actual, "transaction budget must not extend the caller's deadline")
+
+		return nil
+	}))
+
 	attempts := 0
 	err = lock.InSerializableTx(t.Context(), func(context.Context, pgx.Tx) error {
 		attempts++