Merge pull request #596 from pchieneye/feat/hedged-reads

feat: add hedged reads to reduce replica tail latency

Author: thlpkee20-wq

internal/db/POOL_NOTES.md

@@ -55,6 +55,14 @@ closes the pool on graceful shutdown.
   request rather than a process-wide stall.
 - No secrets are logged; the DSN is never printed (only wrapped error text).
 
+## Hedged reads for replica tail latency
+
+The read router now supports optional hedged reads for read-only SELECT statements routed to a replica. When the first replica request remains in flight past a configured delay, the router starts a second replica attempt and returns the first successful result while canceling the losing request. This keeps the normal read path efficient and limits extra work to slow-tail cases only.
+
+- The behavior is gated to read-only SELECTs so writes and non-SELECT reads remain on the primary path.
+- The helper exposes the metric `hedged_reads_total{winner}` to observe which hedge attempt won.
+- Context cancellation is propagated to both in-flight attempts so the losing query does not continue to consume replica CPU.
+
 ## Tests
 
 `internal/db/pool_test.go` (no live DB required):

internal/db/pool.go

@@ -3,15 +3,134 @@ package db
 import (
 	"context"
 	"fmt"
+	"sync"
+	"time"
+
 	"stellarbill-backend/internal/config"
 	"stellarbill-backend/internal/middleware"
 	"stellarbill-backend/internal/servertiming"
 	"time"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+var hedgedReadsTotal = promauto.NewCounterVec(
+	prometheus.CounterOpts{
+		Name: "hedged_reads_total",
+		Help: "Total hedged read attempts by which attempt won.",
+	},
+	[]string{"winner"},
 )
 
+// HedgedQuery executes the first attempt immediately and starts a second attempt
+// after the supplied delay if the first attempt is still running. The first
+// successful attempt wins and the losing attempt is canceled to avoid doing
+// extra work on replicas.
+func HedgedQuery(ctx context.Context, delay time.Duration, fn func(ctx context.Context, attempt int) error) (int, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	if delay < 0 {
+		delay = 0
+	}
+
+	type attemptResult struct {
+		attempt int
+		err     error
+	}
+
+	results := make(chan attemptResult, 2)
+	done := make(chan struct{})
+	var doneOnce sync.Once
+	closeDone := func() {
+		doneOnce.Do(func() {
+			close(done)
+		})
+	}
+	var cancelFirst context.CancelFunc
+	firstCtx, cancelFirst := context.WithCancel(ctx)
+	defer cancelFirst()
+
+	go func() {
+		results <- attemptResult{attempt: 1, err: fn(firstCtx, 1)}
+	}()
+
+	var (
+		cancelSecond  context.CancelFunc
+		secondCtx     context.Context
+		secondStarted bool
+		mu            sync.Mutex
+	)
+
+	startSecond := func() {
+		mu.Lock()
+		defer mu.Unlock()
+		if secondStarted {
+			return
+		}
+		secondStarted = true
+		secondCtx, cancelSecond = context.WithCancel(ctx)
+		go func() {
+			results <- attemptResult{attempt: 2, err: fn(secondCtx, 2)}
+		}()
+	}
+
+	if delay > 0 {
+		go func() {
+			timer := time.NewTimer(delay)
+			defer timer.Stop()
+			select {
+			case <-ctx.Done():
+			case <-done:
+			case <-timer.C:
+				startSecond()
+			}
+		}()
+	} else {
+		startSecond()
+	}
+
+	var (
+		winnerErr error
+		winner    int
+		seen      int
+	)
+
+	for seen < 2 {
+		select {
+		case <-ctx.Done():
+			if cancelSecond != nil {
+				cancelSecond()
+			}
+			return 0, ctx.Err()
+		case result := <-results:
+			seen++
+			if result.err == nil {
+				if result.attempt == 1 {
+					if cancelSecond != nil {
+						cancelSecond()
+					}
+				} else {
+					cancelFirst()
+				}
+				closeDone()
+				hedgedReadsTotal.WithLabelValues(map[bool]string{true: "first", false: "second"}[result.attempt == 1]).Inc()
+				return result.attempt, nil
+			}
+			winnerErr = result.err
+			winner = result.attempt
+		}
+	}
+
+	if winner > 0 {
+		hedgedReadsTotal.WithLabelValues(map[bool]string{true: "first", false: "second"}[winner == 1]).Inc()
+	}
+	return winner, winnerErr
+}
+
 // PoolPinger adapts a *pgxpool.Pool to the handlers.DBPinger interface.
 //
 // pgxpool.Pool exposes Ping(ctx) but the health-check code (handlers.DBPinger)

internal/db/router.go

@@ -3,6 +3,7 @@ package db
 import (
 	"context"
 	"database/sql"
+	"strings"
 	"sync"
 	"time"
 
@@ -24,24 +25,36 @@ type Pinger interface {
 // ReadRouter routes read queries to a read replica or primary database pool.
 // It implements the DBTX interface, directing safe read context calls to the replica.
 type ReadRouter struct {
-	primary DBTX
-	replica DBTX
+	primary  DBTX
+	replica  DBTX
+	replica2 DBTX
 
 	// Failover configuration
 	mu              sync.RWMutex
 	replicaDown     bool
 	lastCheck       time.Time
 	pingTimeout     time.Duration
 	healthCheckFreq time.Duration
+	hedgeDelay      time.Duration
 }
 
 // NewReadRouter creates a new ReadRouter with primary and replica connections.
 func NewReadRouter(primary, replica DBTX) *ReadRouter {
+	return NewReadRouterWithReplicas(primary, replica, nil)
+}
+
+// NewReadRouterWithReplicas creates a router that can hedge read-only SELECTs
+// across a primary and two replica backends. The second replica is only used
+// when the query is a single read-only SELECT and the first replica is still
+// running beyond the configured hedge delay.
+func NewReadRouterWithReplicas(primary, replica, replica2 DBTX) *ReadRouter {
 	return &ReadRouter{
 		primary:         primary,
 		replica:         replica,
+		replica2:        replica2,
 		pingTimeout:     50 * time.Millisecond,
 		healthCheckFreq: 5 * time.Second,
+		hedgeDelay:      50 * time.Millisecond,
 	}
 }
 
@@ -126,8 +139,48 @@ func (r *ReadRouter) PrepareContext(ctx context.Context, query string) (*sql.Stm
 	return r.primary.PrepareContext(ctx, query)
 }
 
-// QueryContext routes reads to the Reader.
+// QueryContext routes reads to the Reader. Read-only SELECTs can be hedged
+// across the first two replicas to reduce slow-tail latency.
 func (r *ReadRouter) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+	selected := r.Reader(ctx)
+	if selected == r.primary || !r.shouldHedge(query) || r.replica2 == nil {
+		return selected.QueryContext(ctx, query, args...)
+	}
+
+	var (
+		resultRows *sql.Rows
+		resultErr  error
+		mu         sync.Mutex
+	)
+
+	_, err := HedgedQuery(ctx, r.hedgeDelay, func(attemptCtx context.Context, attempt int) error {
+		var target DBTX
+		if attempt == 1 {
+			target = r.replica
+		} else {
+			target = r.replica2
+		}
+		if target == nil {
+			return nil
+		}
+
+		rows, queryErr := target.QueryContext(attemptCtx, query, args...)
+		mu.Lock()
+		defer mu.Unlock()
+		if queryErr == nil && resultRows == nil {
+			resultRows = rows
+			resultErr = nil
+			return nil
+		}
+		if queryErr != nil && resultErr == nil && resultRows == nil {
+			resultErr = queryErr
+		}
+		return queryErr
+	})
+	if err != nil && resultRows == nil {
+		return nil, err
+	}
+	return resultRows, nil
 	if acc := middlewarepkg.AccumulatorFromContext(ctx); acc != nil {
 		acc.AddDBRowsRead(1)
 	}
@@ -142,6 +195,27 @@ func (r *ReadRouter) QueryRowContext(ctx context.Context, query string, args ...
 	return r.Reader(ctx).QueryRowContext(ctx, query, args...)
 }
 
+func (r *ReadRouter) shouldHedge(query string) bool {
+	q := strings.TrimSpace(query)
+	if q == "" {
+		return false
+	}
+	if strings.Contains(q, ";") {
+		return false
+	}
+
+	upper := strings.ToUpper(q)
+	if !strings.HasPrefix(upper, "SELECT") {
+		return false
+	}
+	for _, forbidden := range []string{"INSERT", "UPDATE", "DELETE", "MERGE", "TRUNCATE", "CREATE", "ALTER", "DROP", "REPLACE", "CALL", "DO", "COPY"} {
+		if strings.Contains(upper, forbidden) {
+			return false
+		}
+	}
+	return !strings.Contains(upper, "FOR UPDATE") && !strings.Contains(upper, "FOR SHARE")
+}
+
 // Exec routes writes to the primary pool.
 func (r *ReadRouter) Exec(query string, args ...any) (sql.Result, error) {
 	return r.primary.Exec(query, args...)

internal/db/router_test.go

@@ -31,6 +31,7 @@ type mockDBTX struct {
 	execNoCtxCount     int
 	queryNoCtxCount    int
 	queryRowNoCtxCount int
+	queryContextFunc   func(ctx context.Context, query string, args ...any) (*sql.Rows, error)
 }
 
 func newMockDBTX(pingErr error) *mockDBTX {
@@ -51,6 +52,9 @@ func (m *mockDBTX) PrepareContext(ctx context.Context, query string) (*sql.Stmt,
 
 func (m *mockDBTX) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
 	m.queryCount++
+	if m.queryContextFunc != nil {
+		return m.queryContextFunc(ctx, query, args...)
+	}
 	return nil, nil
 }
 
@@ -132,6 +136,86 @@ func TestReadRouter_ReplicaFailover(t *testing.T) {
 	})
 }
 
+func TestHedgedQuery_CancelsTheSlowAttempt(t *testing.T) {
+	ctx := context.Background()
+	firstCancelled := make(chan struct{})
+	secondStarted := make(chan struct{})
+
+	winner, err := HedgedQuery(ctx, 10*time.Millisecond, func(attemptCtx context.Context, attempt int) error {
+		if attempt == 1 {
+			go func() {
+				<-attemptCtx.Done()
+				close(firstCancelled)
+			}()
+			<-attemptCtx.Done()
+			return context.Canceled
+		}
+		close(secondStarted)
+		return nil
+	})
+
+	require.NoError(t, err)
+	assert.Equal(t, 2, winner)
+	select {
+	case <-firstCancelled:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected first attempt context to be canceled")
+	}
+	select {
+	case <-secondStarted:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected second attempt to start")
+	}
+}
+
+func TestReadRouter_QueryContext_HedgesReadOnlySelectsAcrossReplicas(t *testing.T) {
+	primary := newMockDBTX(nil)
+	replica := newMockDBTX(nil)
+	replica2 := newMockDBTX(nil)
+	router := NewReadRouterWithReplicas(primary, replica, replica2)
+	router.hedgeDelay = 10 * time.Millisecond
+
+	firstAttemptCanceled := make(chan struct{})
+	replica.queryContextFunc = func(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+		go func() {
+			<-ctx.Done()
+			close(firstAttemptCanceled)
+		}()
+		<-ctx.Done()
+		return nil, ctx.Err()
+	}
+	replica2.queryContextFunc = func(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+		return nil, nil
+	}
+
+	rows, err := router.QueryContext(context.Background(), "SELECT * FROM users")
+	require.NoError(t, err)
+	assert.Nil(t, rows)
+	assert.Equal(t, 1, replica.queryCount)
+	assert.Equal(t, 1, replica2.queryCount)
+	assert.Equal(t, 0, primary.queryCount)
+
+	select {
+	case <-firstAttemptCanceled:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected first attempt context to be canceled after hedge winner")
+	}
+}
+
+func TestReadRouter_QueryContext_DoesNotHedgeWriteQueries(t *testing.T) {
+	primary := newMockDBTX(nil)
+	replica := newMockDBTX(nil)
+	replica2 := newMockDBTX(nil)
+	router := NewReadRouterWithReplicas(primary, replica, replica2)
+	router.hedgeDelay = 10 * time.Millisecond
+
+	_, err := router.QueryContext(context.Background(), "UPDATE users SET active = true")
+	require.NoError(t, err)
+	assert.Equal(t, 1, replica.queryCount)
+	assert.Equal(t, 0, replica2.queryCount)
+	assert.Equal(t, 0, primary.queryCount)
+}
+
 func TestReadRouter_DBTXInterfaceMethods(t *testing.T) {
 	primary := newMockDBTX(nil)
 	replica := newMockDBTX(nil)