From d4252b228366aced8d3d388cc0a054eed9795172 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:19:40 -0700 Subject: [PATCH 01/11] feat(router): add the Policy interface and round-robin selection Introduce internal/router, the seam where "which eligible replica serves this read" becomes pluggable. A Policy ranks the replicas the proxy has already deemed eligible (healthy and past the fence) and picks one; eligibility stays upstream so a Policy never needs a database and can be tested against a simulated backend. Choose/Release bracket a read so a load-aware policy can track in-flight work, and NeedsFingerprint lets a caller skip computing a fingerprint for policies that ignore the query shape. Round-robin is the first, simplest implementation: an atomic rotation over the current candidate set. --- internal/router/policy.go | 44 ++++++++++++++++++++ internal/router/roundrobin.go | 33 +++++++++++++++ internal/router/roundrobin_test.go | 65 ++++++++++++++++++++++++++++++ 3 files changed, 142 insertions(+) create mode 100644 internal/router/policy.go create mode 100644 internal/router/roundrobin.go create mode 100644 internal/router/roundrobin_test.go diff --git a/internal/router/policy.go b/internal/router/policy.go new file mode 100644 index 0000000..47fac6e --- /dev/null +++ b/internal/router/policy.go @@ -0,0 +1,44 @@ +// Package router chooses which eligible replica serves a read. Eligibility — +// whether a replica is healthy and fresh enough under the fencing rules — is +// decided upstream in the proxy; a Policy only ranks the replicas that already +// qualify and picks one, balancing load across them. +// +// Policies are self-contained and safe for concurrent use: a load-aware policy +// tracks its own in-flight counts and latency estimates across Choose/Release +// calls, so the package has no database dependency and every policy can be +// exercised against a simulated backend in a plain unit test. +package router + +import "time" + +// Candidate is a replica eligible to serve a read, carrying the live signals a +// Policy may weigh. In-flight load and per-fingerprint latency are not supplied +// here — a Policy that cares about them accumulates them itself as reads are +// chosen and released. +type Candidate struct { + // Addr is the replica's network address and the Policy's return value. + Addr string + // LagBytes and LagSeconds are the replica's replication lag as of the most + // recent health poll. + LagBytes int64 + LagSeconds float64 +} + +// Policy selects one candidate to serve a read. Implementations must be safe for +// concurrent use by many sessions. Every Choose must be paired with exactly one +// Release of the returned address so that load-aware policies keep an accurate +// in-flight count; stateless policies treat Release as a no-op. +type Policy interface { + // Choose picks a candidate for a read whose query fingerprint is fingerprint + // (empty when NeedsFingerprint reports false) and returns its address, + // marking one read in flight on that backend. candidates is never empty. + Choose(candidates []Candidate, fingerprint string) string + // Release records that a read chosen earlier for addr has completed after + // latency, clearing its in-flight mark and folding latency into any estimate. + Release(addr, fingerprint string, latency time.Duration) + // NeedsFingerprint reports whether the policy consults the query fingerprint, + // so a caller can skip the cost of computing one when it does not. + NeedsFingerprint() bool + // Name identifies the policy in logs and metrics. + Name() string +} diff --git a/internal/router/roundrobin.go b/internal/router/roundrobin.go new file mode 100644 index 0000000..ed0e620 --- /dev/null +++ b/internal/router/roundrobin.go @@ -0,0 +1,33 @@ +package router + +import ( + "sync/atomic" + "time" +) + +// RoundRobin hands out candidates in rotation, ignoring load and latency. It is +// the simplest fair policy: over a stable candidate set every replica receives +// an equal share of reads. Because the eligible set can change between calls +// (a replica may fall behind the fence or fail health checks), the rotation +// indexes into whatever candidates it is given rather than pinning positions. +type RoundRobin struct { + next atomic.Uint64 +} + +// NewRoundRobin returns a ready round-robin policy. +func NewRoundRobin() *RoundRobin { return &RoundRobin{} } + +// Choose returns the next candidate in rotation. +func (p *RoundRobin) Choose(candidates []Candidate, _ string) string { + i := p.next.Add(1) - 1 + return candidates[i%uint64(len(candidates))].Addr +} + +// Release is a no-op: round-robin keeps no per-backend state. +func (p *RoundRobin) Release(string, string, time.Duration) {} + +// NeedsFingerprint reports false; round-robin ignores the query shape. +func (p *RoundRobin) NeedsFingerprint() bool { return false } + +// Name identifies the policy. +func (p *RoundRobin) Name() string { return "round-robin" } diff --git a/internal/router/roundrobin_test.go b/internal/router/roundrobin_test.go new file mode 100644 index 0000000..82ef601 --- /dev/null +++ b/internal/router/roundrobin_test.go @@ -0,0 +1,65 @@ +package router + +import ( + "sync" + "testing" +) + +func addrs(names ...string) []Candidate { + c := make([]Candidate, len(names)) + for i, n := range names { + c[i] = Candidate{Addr: n} + } + return c +} + +func TestRoundRobin_RotatesEvenly(t *testing.T) { + p := NewRoundRobin() + cands := addrs("a", "b", "c") + counts := map[string]int{} + for i := 0; i < 30; i++ { + counts[p.Choose(cands, "")]++ + } + for _, a := range []string{"a", "b", "c"} { + if counts[a] != 10 { + t.Errorf("addr %q served %d reads, want 10 (counts=%v)", a, counts[a], counts) + } + } +} + +func TestRoundRobin_Order(t *testing.T) { + p := NewRoundRobin() + cands := addrs("a", "b") + want := []string{"a", "b", "a", "b"} + for i, w := range want { + if got := p.Choose(cands, ""); got != w { + t.Errorf("choice %d = %q, want %q", i, got, w) + } + } +} + +func TestRoundRobin_SingleCandidate(t *testing.T) { + p := NewRoundRobin() + cands := addrs("only") + for i := 0; i < 5; i++ { + if got := p.Choose(cands, ""); got != "only" { + t.Fatalf("choice %d = %q, want only", i, got) + } + } +} + +func TestRoundRobin_ConcurrentIsRace(t *testing.T) { + p := NewRoundRobin() + cands := addrs("a", "b", "c", "d") + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 1000; j++ { + p.Release(p.Choose(cands, ""), "", 0) + } + }() + } + wg.Wait() +} From 369eed721a14c465ca97824e40d0fa0d839a3c34 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:20:24 -0700 Subject: [PATCH 02/11] feat(router): add least-in-flight selection Round-robin is blind: it keeps feeding a replica that has fallen behind on a slow query as readily as an idle one. Least-in-flight instead sends each read to the candidate with the fewest reads outstanding, so a backend whose work is piling up naturally stops being chosen until it drains. The policy keeps a per-address in-flight count, incremented on Choose and decremented on the paired Release, and breaks ties toward the earlier candidate for deterministic selection over a stable candidate order. --- internal/router/leastinflight.go | 59 ++++++++++++++++++ internal/router/leastinflight_test.go | 86 +++++++++++++++++++++++++++ 2 files changed, 145 insertions(+) create mode 100644 internal/router/leastinflight.go create mode 100644 internal/router/leastinflight_test.go diff --git a/internal/router/leastinflight.go b/internal/router/leastinflight.go new file mode 100644 index 0000000..f7ee7ad --- /dev/null +++ b/internal/router/leastinflight.go @@ -0,0 +1,59 @@ +package router + +import ( + "sync" + "time" +) + +// LeastInFlight sends each read to the candidate currently serving the fewest +// outstanding reads, spreading load by observed occupancy rather than by blind +// rotation. It adapts to replicas that are slow or briefly overloaded — a +// backend whose reads are piling up stops being chosen until it drains — without +// needing to know anything about the query itself. +// +// In-flight counts are kept per address, incremented on Choose and decremented +// on the paired Release. An address absent from the map is serving zero reads. +type LeastInFlight struct { + mu sync.Mutex + inFlight map[string]int +} + +// NewLeastInFlight returns a ready least-in-flight policy. +func NewLeastInFlight() *LeastInFlight { + return &LeastInFlight{inFlight: make(map[string]int)} +} + +// Choose returns the candidate with the fewest reads in flight, breaking ties in +// favor of the earlier candidate so selection is deterministic for a stable +// candidate order. +func (p *LeastInFlight) Choose(candidates []Candidate, _ string) string { + p.mu.Lock() + defer p.mu.Unlock() + best := candidates[0].Addr + bestN := p.inFlight[best] + for _, c := range candidates[1:] { + if n := p.inFlight[c.Addr]; n < bestN { + best, bestN = c.Addr, n + } + } + p.inFlight[best]++ + return best +} + +// Release decrements the in-flight count for addr. Latency is ignored; +// least-in-flight weighs occupancy, not duration. +func (p *LeastInFlight) Release(addr, _ string, _ time.Duration) { + p.mu.Lock() + defer p.mu.Unlock() + if p.inFlight[addr] <= 1 { + delete(p.inFlight, addr) + return + } + p.inFlight[addr]-- +} + +// NeedsFingerprint reports false; least-in-flight ignores the query shape. +func (p *LeastInFlight) NeedsFingerprint() bool { return false } + +// Name identifies the policy. +func (p *LeastInFlight) Name() string { return "least-in-flight" } diff --git a/internal/router/leastinflight_test.go b/internal/router/leastinflight_test.go new file mode 100644 index 0000000..5c1f507 --- /dev/null +++ b/internal/router/leastinflight_test.go @@ -0,0 +1,86 @@ +package router + +import ( + "sync" + "testing" +) + +func TestLeastInFlight_PicksIdlest(t *testing.T) { + p := NewLeastInFlight() + cands := addrs("a", "b", "c") + + // Load up "a" with two outstanding reads. + if got := p.Choose(cands, ""); got != "a" { // tie -> first + t.Fatalf("first choice = %q, want a", got) + } + // "a" now has 1; next idlest is a tie between b and c -> b. + if got := p.Choose(cands, ""); got != "b" { + t.Fatalf("second choice = %q, want b", got) + } + // b has 1; c has 0 -> c. + if got := p.Choose(cands, ""); got != "c" { + t.Fatalf("third choice = %q, want c", got) + } + // all at 1 -> tie, first (a). + if got := p.Choose(cands, ""); got != "a" { + t.Fatalf("fourth choice = %q, want a", got) + } +} + +func TestLeastInFlight_ReleaseFreesCapacity(t *testing.T) { + p := NewLeastInFlight() + cands := addrs("a", "b") + // Drive a to 3 in flight, b to 0 by always releasing b. + p.Choose(cands, "") // a=1 + p.Choose(addrs("a"), "") // a=2 + p.Choose(addrs("a"), "") // a=3 + if got := p.Choose(cands, ""); got != "b" { + t.Fatalf("with a=3,b=0 chose %q, want b", got) + } + // Release a twice: a=1, b=1 -> tie -> a. + p.Release("a", "", 0) + p.Release("a", "", 0) + if got := p.Choose(cands, ""); got != "a" { + t.Fatalf("after releasing a twice chose %q, want a", got) + } +} + +func TestLeastInFlight_ReleaseBalances(t *testing.T) { + p := NewLeastInFlight() + cands := addrs("a", "b", "c") + // Every read is released immediately, so counts stay at 0 and selection + // falls back to round-robin-like fairness via the tie-break order. + counts := map[string]int{} + for i := 0; i < 30; i++ { + addr := p.Choose(cands, "") + counts[addr]++ + p.Release(addr, "", 0) + } + if counts["a"] != 30 { + t.Errorf("with immediate release, ties always pick first: counts=%v", counts) + } + // The map should be empty again once every read is released. + if len(p.inFlight) != 0 { + t.Errorf("in-flight map not drained: %v", p.inFlight) + } +} + +func TestLeastInFlight_ConcurrentBalance(t *testing.T) { + p := NewLeastInFlight() + cands := addrs("a", "b", "c", "d") + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 500; j++ { + addr := p.Choose(cands, "") + p.Release(addr, "", 0) + } + }() + } + wg.Wait() + if len(p.inFlight) != 0 { + t.Errorf("in-flight map not drained after balanced load: %v", p.inFlight) + } +} From 64cdee9b51d84ad7b10f4923a9c6b6d5c4f90826 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:21:02 -0700 Subject: [PATCH 03/11] feat(classify): expose pg_query fingerprints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The scored routing policy keys its per-query-shape latency estimates by fingerprint, so it needs a stable identity for "queries like this one". pg_query already computes exactly that — a parse-tree hash with literals normalized away — so SELECT ... WHERE id = 1 and id = 42 collapse to one shape while a different table or clause set does not. Expose it from classify, which already owns the pg_query dependency, rather than leaking the parser into the router package; the router stays pure Go and database-free. An unparseable query returns "", which the scored policy treats as an unknown shape. --- internal/classify/classify.go | 12 ++++++++ internal/classify/fingerprint_test.go | 43 +++++++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 internal/classify/fingerprint_test.go diff --git a/internal/classify/classify.go b/internal/classify/classify.go index 0e417fe..3968ff1 100644 --- a/internal/classify/classify.go +++ b/internal/classify/classify.go @@ -44,6 +44,18 @@ var volatileFuncs = map[string]struct{}{ "clock_timestamp": {}, "timeofday": {}, "pg_sleep": {}, } +// Fingerprint returns pg_query's stable fingerprint of a query — a hash of its +// parse tree with literals normalized away, so queries that differ only in their +// constants share a fingerprint. Load-aware routing uses it to key per-shape +// latency estimates. An unparseable query has no fingerprint and returns "". +func Fingerprint(sql string) string { + fp, err := pg.Fingerprint(sql) + if err != nil { + return "" + } + return fp +} + // Classify returns the routing class of a simple-query string, which may contain // several statements. It is a Write if any statement is a Write. func Classify(sql string) Class { diff --git a/internal/classify/fingerprint_test.go b/internal/classify/fingerprint_test.go new file mode 100644 index 0000000..98f75ad --- /dev/null +++ b/internal/classify/fingerprint_test.go @@ -0,0 +1,43 @@ +package classify_test + +import ( + "testing" + + "github.com/sachhg/pgpilot/internal/classify" +) + +func TestFingerprint_StableAcrossLiterals(t *testing.T) { + // Queries differing only in constants share a fingerprint. + a := classify.Fingerprint("SELECT * FROM users WHERE id = 1") + b := classify.Fingerprint("SELECT * FROM users WHERE id = 42") + if a == "" { + t.Fatal("fingerprint was empty for a valid query") + } + if a != b { + t.Errorf("fingerprints differ across literals: %q vs %q", a, b) + } +} + +func TestFingerprint_DistinguishesShapes(t *testing.T) { + a := classify.Fingerprint("SELECT * FROM users WHERE id = 1") + b := classify.Fingerprint("SELECT * FROM orders WHERE id = 1") + if a == b { + t.Errorf("different table shapes shared a fingerprint: %q", a) + } +} + +func TestFingerprint_UnparseableIsEmpty(t *testing.T) { + if fp := classify.Fingerprint("SELECT FROM WHERE ("); fp != "" { + t.Errorf("unparseable query returned fingerprint %q, want empty", fp) + } +} + +func TestFingerprint_Deterministic(t *testing.T) { + sql := "SELECT a, b FROM t JOIN u ON t.id = u.id WHERE t.x > 3 ORDER BY a" + first := classify.Fingerprint(sql) + for i := 0; i < 5; i++ { + if got := classify.Fingerprint(sql); got != first { + t.Fatalf("fingerprint not deterministic: %q vs %q", got, first) + } + } +} From 39085b5a5df07c84cc58708b47acc3d90a4e4b18 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:22:39 -0700 Subject: [PATCH 04/11] feat(router): add the latency-scored policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Least-in-flight treats every read as equally costly, so a replica running two expensive scans looks as free as one running two index lookups. The scored policy fixes that by estimating each candidate's completion time: (inFlight + 1) * ewmaLatency(addr, fingerprint) + lagPenalty * lagSeconds, and choosing the minimum. Keying the EWMA by pg_query fingerprint is the crux — it lets the policy steer an expensive query shape away from a replica already loaded with that shape while still sending cheap queries there, and it learns per-replica speed differences for free as slower backends accrue higher averages. Weights are tunable; the defaults suit comparable replicas. Lag is folded in so that in bounded/relaxed mode a trailing replica is deprioritized even when idle. --- internal/router/scored.go | 117 ++++++++++++++++++++++++++++++ internal/router/scored_test.go | 127 +++++++++++++++++++++++++++++++++ 2 files changed, 244 insertions(+) create mode 100644 internal/router/scored.go create mode 100644 internal/router/scored_test.go diff --git a/internal/router/scored.go b/internal/router/scored.go new file mode 100644 index 0000000..7d45845 --- /dev/null +++ b/internal/router/scored.go @@ -0,0 +1,117 @@ +package router + +import ( + "math" + "sync" + "time" +) + +// ScoredWeights tunes the Scored policy. The defaults suit a cluster of +// comparable replicas; a deployment with wildly heterogeneous hardware or a +// staleness budget may want to raise LagPenaltyMs. +type ScoredWeights struct { + // Alpha is the EWMA smoothing factor in (0, 1]: higher reacts faster to a + // replica's recent latency, lower is steadier. + Alpha float64 + // DefaultServiceMs is the assumed per-read service time, in milliseconds, for + // a (replica, query-shape) pair not yet observed. + DefaultServiceMs float64 + // LagPenaltyMs is the estimated cost, in milliseconds, added per second of a + // replica's replication lag, discouraging reads to replicas that trail. + LagPenaltyMs float64 +} + +// DefaultScoredWeights returns the weights NewScored uses. +func DefaultScoredWeights() ScoredWeights { + return ScoredWeights{Alpha: 0.2, DefaultServiceMs: 1.0, LagPenaltyMs: 50.0} +} + +// Scored ranks candidates by an estimate of how quickly each would finish a new +// read, then picks the lowest. The estimate multiplies the replica's current +// in-flight count (plus the incoming read) by an exponentially weighted moving +// average of how long this query shape has taken on that replica, and adds a +// penalty for replication lag: +// +// score = (inFlight + 1) * ewmaLatency(addr, fingerprint) + lagPenalty * lagSeconds +// +// Keying latency by fingerprint is what lets it steer *expensive* shapes away +// from busy replicas: a costly query on a replica already running two of them +// scores far worse than the same query on an idle one, even when a cheap query +// would happily use the busy replica. It learns per-replica speed differences +// for free — a slower replica accrues higher EWMAs and wins fewer reads. +type Scored struct { + w ScoredWeights + + mu sync.Mutex + inFlight map[string]int + latency map[fpKey]float64 // milliseconds, per (addr, fingerprint) +} + +type fpKey struct { + addr string + fingerprint string +} + +// NewScored returns a scored policy with the default weights. +func NewScored() *Scored { return NewScoredWith(DefaultScoredWeights()) } + +// NewScoredWith returns a scored policy with explicit weights. +func NewScoredWith(w ScoredWeights) *Scored { + return &Scored{ + w: w, + inFlight: make(map[string]int), + latency: make(map[fpKey]float64), + } +} + +// Choose returns the candidate with the lowest estimated completion time for a +// read of the given fingerprint, breaking ties toward the earlier candidate. +func (p *Scored) Choose(candidates []Candidate, fingerprint string) string { + p.mu.Lock() + defer p.mu.Unlock() + best := "" + bestScore := math.Inf(1) + for _, c := range candidates { + service := p.serviceLocked(c.Addr, fingerprint) + score := float64(p.inFlight[c.Addr]+1)*service + p.w.LagPenaltyMs*c.LagSeconds + if score < bestScore { + best, bestScore = c.Addr, score + } + } + p.inFlight[best]++ + return best +} + +// Release clears the in-flight mark for addr and folds the read's measured +// latency into the EWMA for its (addr, fingerprint) pair. +func (p *Scored) Release(addr, fingerprint string, latency time.Duration) { + ms := float64(latency) / float64(time.Millisecond) + p.mu.Lock() + defer p.mu.Unlock() + if p.inFlight[addr] <= 1 { + delete(p.inFlight, addr) + } else { + p.inFlight[addr]-- + } + key := fpKey{addr, fingerprint} + prev, ok := p.latency[key] + if !ok { + prev = p.w.DefaultServiceMs + } + p.latency[key] = p.w.Alpha*ms + (1-p.w.Alpha)*prev +} + +// serviceLocked returns the current latency estimate for a (addr, fingerprint) +// pair, falling back to the default before any observation. Caller holds p.mu. +func (p *Scored) serviceLocked(addr, fingerprint string) float64 { + if v, ok := p.latency[fpKey{addr, fingerprint}]; ok { + return v + } + return p.w.DefaultServiceMs +} + +// NeedsFingerprint reports true; the scored policy keys latency by query shape. +func (p *Scored) NeedsFingerprint() bool { return true } + +// Name identifies the policy. +func (p *Scored) Name() string { return "scored" } diff --git a/internal/router/scored_test.go b/internal/router/scored_test.go new file mode 100644 index 0000000..f99e401 --- /dev/null +++ b/internal/router/scored_test.go @@ -0,0 +1,127 @@ +package router + +import ( + "sync" + "testing" + "time" +) + +// lagCand builds a candidate with a given replication lag. +func lagCand(addr string, lagSeconds float64) Candidate { + return Candidate{Addr: addr, LagSeconds: lagSeconds} +} + +func TestScored_ColdStartPrefersIdle(t *testing.T) { + p := NewScored() + cands := addrs("a", "b") + // First read: tie at equal cold estimates -> first candidate. + if got := p.Choose(cands, "fp"); got != "a" { + t.Fatalf("cold first choice = %q, want a", got) + } + // "a" now has one in flight; "b" is idle -> "b" is cheaper. + if got := p.Choose(cands, "fp"); got != "b" { + t.Fatalf("second choice = %q, want b (a is loaded)", got) + } +} + +func TestScored_SteersExpensiveShapeToFasterReplica(t *testing.T) { + p := NewScored() + cands := addrs("fast", "slow") + + // Teach the policy that fingerprint "big" is cheap on "fast" and dear on + // "slow" by observing several completions of each. + for i := 0; i < 20; i++ { + p.Choose([]Candidate{{Addr: "fast"}}, "big") + p.Release("fast", "big", 5*time.Millisecond) + p.Choose([]Candidate{{Addr: "slow"}}, "big") + p.Release("slow", "big", 200*time.Millisecond) + } + + // With both idle, a "big" read should now go to the fast replica. + got := map[string]int{} + for i := 0; i < 10; i++ { + addr := p.Choose(cands, "big") + got[addr]++ + p.Release(addr, "big", 5*time.Millisecond) // keep them idle between picks + } + if got["fast"] == 0 || got["slow"] > got["fast"] { + t.Errorf("expensive shape not steered to fast replica: %v", got) + } +} + +func TestScored_ShapeAwareNotGlobal(t *testing.T) { + p := NewScored() + // "slow" is dear for shape "big" but cheap for shape "small". + for i := 0; i < 20; i++ { + p.Choose([]Candidate{{Addr: "slow"}}, "big") + p.Release("slow", "big", 200*time.Millisecond) + p.Choose([]Candidate{{Addr: "slow"}}, "small") + p.Release("slow", "small", 1*time.Millisecond) + p.Choose([]Candidate{{Addr: "fast"}}, "big") + p.Release("fast", "big", 5*time.Millisecond) + p.Choose([]Candidate{{Addr: "fast"}}, "small") + p.Release("fast", "small", 5*time.Millisecond) + } + cands := addrs("fast", "slow") + // A "small" read should favor "slow" (1ms) over "fast" (5ms), proving the + // estimate is per shape and not a single per-replica number. + got := map[string]int{} + for i := 0; i < 10; i++ { + addr := p.Choose(cands, "small") + got[addr]++ + p.Release(addr, "small", 1*time.Millisecond) + } + if got["slow"] <= got["fast"] { + t.Errorf("small shape should favor slow replica for that shape: %v", got) + } +} + +func TestScored_PenalizesLag(t *testing.T) { + // Equal load and no latency history: the laggier replica should lose. + p := NewScored() + cands := []Candidate{lagCand("behind", 2.0), lagCand("current", 0.0)} + if got := p.Choose(cands, "fp"); got != "current" { + t.Errorf("lag not penalized: chose %q, want current", got) + } +} + +func TestScored_ReleaseDrainsInFlight(t *testing.T) { + p := NewScored() + cands := addrs("a", "b") + for i := 0; i < 50; i++ { + addr := p.Choose(cands, "fp") + p.Release(addr, "fp", time.Millisecond) + } + p.mu.Lock() + defer p.mu.Unlock() + if len(p.inFlight) != 0 { + t.Errorf("in-flight not drained: %v", p.inFlight) + } +} + +func TestScored_Metadata(t *testing.T) { + p := NewScored() + if !p.NeedsFingerprint() { + t.Error("scored policy should need the fingerprint") + } + if p.Name() != "scored" { + t.Errorf("Name = %q, want scored", p.Name()) + } +} + +func TestScored_ConcurrentIsRace(t *testing.T) { + p := NewScored() + cands := addrs("a", "b", "c") + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func(seed int) { + defer wg.Done() + for j := 0; j < 500; j++ { + addr := p.Choose(cands, "fp") + p.Release(addr, "fp", time.Duration(seed+j)*time.Microsecond) + } + }(i) + } + wg.Wait() +} From ab06ae4a00f414ef296ad2f468fe6e9d31bbab42 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:23:03 -0700 Subject: [PATCH 05/11] feat(router): construct policies by name The config file selects a policy by string, so give the package one place that maps a name to an implementation. An unknown name returns an error rather than a default, so a typo in routing.policy fails at startup instead of silently routing with the wrong strategy. --- internal/router/router.go | 28 ++++++++++++++++++++++++++++ internal/router/router_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) create mode 100644 internal/router/router.go create mode 100644 internal/router/router_test.go diff --git a/internal/router/router.go b/internal/router/router.go new file mode 100644 index 0000000..803cb93 --- /dev/null +++ b/internal/router/router.go @@ -0,0 +1,28 @@ +package router + +import "fmt" + +// Policy-name values accepted by New and the config file. +const ( + // PolicyRoundRobin rotates reads evenly across eligible replicas. + PolicyRoundRobin = "round-robin" + // PolicyLeastInFlight sends each read to the least-loaded eligible replica. + PolicyLeastInFlight = "least-in-flight" + // PolicyScored ranks eligible replicas by estimated completion time. + PolicyScored = "scored" +) + +// New constructs the named policy. An unknown name is an error so a typo in the +// config fails loudly at startup rather than silently falling back. +func New(name string) (Policy, error) { + switch name { + case PolicyRoundRobin: + return NewRoundRobin(), nil + case PolicyLeastInFlight: + return NewLeastInFlight(), nil + case PolicyScored: + return NewScored(), nil + default: + return nil, fmt.Errorf("router: unknown policy %q", name) + } +} diff --git a/internal/router/router_test.go b/internal/router/router_test.go new file mode 100644 index 0000000..82f3cfb --- /dev/null +++ b/internal/router/router_test.go @@ -0,0 +1,29 @@ +package router + +import "testing" + +func TestNew(t *testing.T) { + cases := []struct { + name string + wantName string + }{ + {PolicyRoundRobin, "round-robin"}, + {PolicyLeastInFlight, "least-in-flight"}, + {PolicyScored, "scored"}, + } + for _, tc := range cases { + p, err := New(tc.name) + if err != nil { + t.Fatalf("New(%q): %v", tc.name, err) + } + if p.Name() != tc.wantName { + t.Errorf("New(%q).Name() = %q, want %q", tc.name, p.Name(), tc.wantName) + } + } +} + +func TestNew_Unknown(t *testing.T) { + if _, err := New("bogus"); err == nil { + t.Error("New(bogus) succeeded, want an error") + } +} From 65db6152c9f2efcb8c26b68521a6399842ad8bab Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:26:10 -0700 Subject: [PATCH 06/11] bench(router): compare policies on a synthetic workload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a deterministic discrete-event simulator — three replicas (one 3x slow), a heavy-tailed shape mix, seeded exponential arrivals near 75% utilization — and drive every policy through the identical workload with no database. Each replica is a single FIFO server, and Choose/Release fire in simulated-time order so a policy's in-flight view matches real queue depth. The comparison test asserts the load-aware policies beat blind round-robin on tail latency and makespan, and logs the table: round-robin forces a third of all reads onto the slow replica and its queue explodes (p95 ~16s), while scored — which alone learns that heavy shapes are especially costly there — sends it the fewest reads and wins the p95. Companion micro-benchmarks show per-decision cost stays sub-microsecond (round-robin ~7ns, scored ~190ns). --- internal/router/sim_test.go | 203 +++++++++++++++++++++++++++++++ internal/router/workload_test.go | 96 +++++++++++++++ 2 files changed, 299 insertions(+) create mode 100644 internal/router/sim_test.go create mode 100644 internal/router/workload_test.go diff --git a/internal/router/sim_test.go b/internal/router/sim_test.go new file mode 100644 index 0000000..6848a23 --- /dev/null +++ b/internal/router/sim_test.go @@ -0,0 +1,203 @@ +package router + +import ( + "container/heap" + "math" + "math/rand" + "sort" + "time" +) + +// This file holds a small, fully deterministic discrete-event simulator used to +// compare policies on a synthetic mixed workload without a database. Each +// replica is modeled as a single FIFO server whose service time for a read is +// the query shape's base cost times the replica's slowness; requests arrive on a +// seeded schedule with a seeded mix of shapes. A policy's in-flight view stays +// exact because Choose and Release fire in simulated-time order. + +// shape is a query shape: a fingerprint and its base service cost in ms on a +// unit-speed replica. +type shape struct { + fp string + cost float64 +} + +// simReplica is a backend in the simulation. +type simReplica struct { + addr string + slowness float64 // service-time multiplier + lag float64 // replication lag in seconds +} + +// request is one arrival: when it arrives and what shape it is. +type request struct { + arrival float64 + shape shape +} + +// workload is a deterministic stream of requests over a fixed replica set. +type workload struct { + replicas []simReplica + requests []request +} + +// stats summarizes a simulation run. Latencies are in milliseconds. +type stats struct { + makespan float64 + mean float64 + p50 float64 + p95 float64 + p99 float64 + perReplica map[string]int +} + +// genWorkload builds a deterministic workload: three replicas (one 3x slow), a +// heavy-tailed shape mix (mostly cheap, a few very expensive), and exponential +// inter-arrival times, all driven by a seeded PRNG. +func genWorkload(seed int64, n int) workload { + rng := rand.New(rand.NewSource(seed)) + replicas := []simReplica{ + {addr: "r1", slowness: 1.0, lag: 0}, + {addr: "r2", slowness: 1.0, lag: 0}, + {addr: "r3", slowness: 3.0, lag: 0}, // a slower box + } + shapes := []shape{ + {fp: "cheap", cost: 1.0}, + {fp: "mid", cost: 12.0}, + {fp: "heavy", cost: 90.0}, + } + // Cumulative weights: ~72% cheap, ~22% mid, ~6% heavy. + weights := []float64{0.72, 0.94, 1.0} + + // Arrival rate (requests per ms) tuned so the two fast replicas run near + // ~75% utilization: busy enough to queue and separate the policies, stable + // enough that the tail reflects placement quality rather than a runaway + // queue. Mean service is ~8.8ms, so two fast replicas clear ~0.23 req/ms. + const rate = 0.16 + reqs := make([]request, n) + t := 0.0 + for i := range reqs { + t += rng.ExpFloat64() / rate + u := rng.Float64() + s := shapes[len(shapes)-1] + for j, w := range weights { + if u <= w { + s = shapes[j] + break + } + } + reqs[i] = request{arrival: t, shape: s} + } + return workload{replicas: replicas, requests: reqs} +} + +// releaseEvent is a scheduled completion. +type releaseEvent struct { + at float64 + addr string + fp string + latency float64 // ms +} + +type releaseHeap []releaseEvent + +func (h releaseHeap) Len() int { return len(h) } +func (h releaseHeap) Less(i, j int) bool { return h[i].at < h[j].at } +func (h releaseHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *releaseHeap) Push(x interface{}) { *h = append(*h, x.(releaseEvent)) } +func (h *releaseHeap) Pop() interface{} { + old := *h + n := len(old) + ev := old[n-1] + *h = old[:n-1] + return ev +} + +// simulate runs the workload through policy and returns latency statistics. It +// processes arrivals and releases in simulated-time order (releases first on a +// tie) so the policy's in-flight counts mirror real queue depth. +func simulate(policy Policy, w workload) stats { + cands := make([]Candidate, len(w.replicas)) + slow := make(map[string]float64, len(w.replicas)) + for i, r := range w.replicas { + cands[i] = Candidate{Addr: r.addr, LagSeconds: r.lag} + slow[r.addr] = r.slowness + } + + freeAt := make(map[string]float64, len(w.replicas)) + perReplica := make(map[string]int, len(w.replicas)) + latencies := make([]float64, 0, len(w.requests)) + pending := &releaseHeap{} + heap.Init(pending) + + ai := 0 + makespan := 0.0 + for ai < len(w.requests) || pending.Len() > 0 { + nextArrival := math.Inf(1) + if ai < len(w.requests) { + nextArrival = w.requests[ai].arrival + } + nextRelease := math.Inf(1) + if pending.Len() > 0 { + nextRelease = (*pending)[0].at + } + + if nextRelease <= nextArrival { + ev := heap.Pop(pending).(releaseEvent) + policy.Release(ev.addr, ev.fp, time.Duration(ev.latency*float64(time.Millisecond))) + continue + } + + req := w.requests[ai] + ai++ + addr := policy.Choose(cands, req.shape.fp) + svc := req.shape.cost * slow[addr] + start := req.arrival + if freeAt[addr] > start { + start = freeAt[addr] + } + finish := start + svc + freeAt[addr] = finish + if finish > makespan { + makespan = finish + } + lat := finish - req.arrival + latencies = append(latencies, lat) + perReplica[addr]++ + heap.Push(pending, releaseEvent{at: finish, addr: addr, fp: req.shape.fp, latency: lat}) + } + + return summarize(latencies, makespan, perReplica) +} + +func summarize(latencies []float64, makespan float64, perReplica map[string]int) stats { + sort.Float64s(latencies) + sum := 0.0 + for _, l := range latencies { + sum += l + } + n := len(latencies) + return stats{ + makespan: makespan, + mean: sum / float64(n), + p50: percentile(latencies, 0.50), + p95: percentile(latencies, 0.95), + p99: percentile(latencies, 0.99), + perReplica: perReplica, + } +} + +// percentile returns the p-quantile (0..1) of a sorted slice via nearest-rank. +func percentile(sorted []float64, p float64) float64 { + if len(sorted) == 0 { + return 0 + } + i := int(math.Ceil(p*float64(len(sorted)))) - 1 + if i < 0 { + i = 0 + } + if i >= len(sorted) { + i = len(sorted) - 1 + } + return sorted[i] +} diff --git a/internal/router/workload_test.go b/internal/router/workload_test.go new file mode 100644 index 0000000..ba14bcf --- /dev/null +++ b/internal/router/workload_test.go @@ -0,0 +1,96 @@ +package router + +import ( + "testing" + "time" +) + +// TestPolicies_WorkloadComparison runs every policy through the same +// deterministic mixed workload and reports latency statistics. It asserts the +// load-aware policies beat blind round-robin on both tail latency and makespan, +// and that the scored policy — which alone knows a query shape's cost — is at +// least as good as least-in-flight. The table it logs is the phase's +// policy comparison; run `go test -run WorkloadComparison -v ./internal/router`. +func TestPolicies_WorkloadComparison(t *testing.T) { + const ( + seed = 20260723 + n = 6000 + ) + w := genWorkload(seed, n) + + rr := simulate(NewRoundRobin(), genWorkload(seed, n)) + lif := simulate(NewLeastInFlight(), genWorkload(seed, n)) + sc := simulate(NewScored(), w) + + t.Logf("%-16s %10s %10s %10s %10s %10s %s", + "policy", "makespan", "mean", "p50", "p95", "p99", "per-replica") + for _, row := range []struct { + name string + s stats + }{ + {"round-robin", rr}, + {"least-in-flight", lif}, + {"scored", sc}, + } { + t.Logf("%-16s %9.1fms %9.2fms %9.2fms %9.2fms %9.2fms %v", + row.name, row.s.makespan, row.s.mean, row.s.p50, row.s.p95, row.s.p99, row.s.perReplica) + } + + // Load-aware policies must beat blind rotation on tail latency and makespan. + if lif.p95 >= rr.p95 { + t.Errorf("least-in-flight p95 %.2f not better than round-robin %.2f", lif.p95, rr.p95) + } + if sc.p95 >= rr.p95 { + t.Errorf("scored p95 %.2f not better than round-robin %.2f", sc.p95, rr.p95) + } + if sc.makespan >= rr.makespan { + t.Errorf("scored makespan %.1f not better than round-robin %.1f", sc.makespan, rr.makespan) + } + // Shape awareness should not lose to shape-blind load balancing. + if sc.p95 > lif.p95*1.02 { + t.Errorf("scored p95 %.2f worse than least-in-flight %.2f", sc.p95, lif.p95) + } + // Round-robin, ignoring the 3x-slow replica, should overload it; the + // load-aware policies should send it strictly fewer reads. + if lif.perReplica["r3"] >= rr.perReplica["r3"] { + t.Errorf("least-in-flight sent r3 %d reads, round-robin %d; expected fewer", + lif.perReplica["r3"], rr.perReplica["r3"]) + } +} + +// benchWorkload is generated once and reused so the benchmarks measure the +// policy, not workload generation. +var benchWorkload = genWorkload(1, 4000) + +func benchmarkWorkload(b *testing.B, make func() Policy) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + simulate(make(), benchWorkload) + } +} + +func BenchmarkWorkload_RoundRobin(b *testing.B) { + benchmarkWorkload(b, func() Policy { return NewRoundRobin() }) +} +func BenchmarkWorkload_LeastInFlight(b *testing.B) { + benchmarkWorkload(b, func() Policy { return NewLeastInFlight() }) +} +func BenchmarkWorkload_Scored(b *testing.B) { + benchmarkWorkload(b, func() Policy { return NewScored() }) +} + +// benchmarkChoose measures the cost of a single Choose/Release pair over a fixed +// candidate set, isolating per-decision overhead from the simulator. +func benchmarkChoose(b *testing.B, p Policy) { + cands := addrs("a", "b", "c", "d") + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + addr := p.Choose(cands, "fp") + p.Release(addr, "fp", time.Millisecond) + } +} + +func BenchmarkChoose_RoundRobin(b *testing.B) { benchmarkChoose(b, NewRoundRobin()) } +func BenchmarkChoose_LeastInFlight(b *testing.B) { benchmarkChoose(b, NewLeastInFlight()) } +func BenchmarkChoose_Scored(b *testing.B) { benchmarkChoose(b, NewScored()) } From f26d21116f348a15955e5879d7a12dcdec2cd81b Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:27:12 -0700 Subject: [PATCH 07/11] feat(config): select the routing policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a routing.policy field so a deployment can pick how reads spread across eligible replicas without a rebuild. It validates against the router's own policy names — an unknown value fails at load, like pool.mode and fencing.mode — and defaults to least-in-flight: a strong, zero-config choice that costs no per-read fingerprint parse, leaving scored as an opt-in for clusters with heterogeneous replicas or query costs. --- internal/config/config.go | 22 ++++++++++++++++++++++ internal/config/config_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/internal/config/config.go b/internal/config/config.go index 6719976..f8da03d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -10,6 +10,8 @@ import ( "fmt" "os" "time" + + "github.com/sachhg/pgpilot/internal/router" ) // Config is the top-level pgpilot configuration. @@ -21,6 +23,17 @@ type Config struct { Pool Pool `json:"pool"` Health Health `json:"health"` Fencing Fencing `json:"fencing"` + Routing Routing `json:"routing"` +} + +// Routing configures how reads are balanced across the replicas eligible to +// serve them (those healthy and fresh enough under the fencing rules). +type Routing struct { + // Policy is "round-robin" (even rotation), "least-in-flight" (fewest + // outstanding reads wins), or "scored" (rank by estimated completion time, + // learning per-query-shape latency). "scored" costs one pg_query fingerprint + // per read, so "least-in-flight" is the default. + Policy string `json:"policy"` } // Fencing-mode values. @@ -157,6 +170,9 @@ func (c *Config) applyDefaults() { if c.Fencing.BoundedMs == 0 { c.Fencing.BoundedMs = 100 } + if c.Routing.Policy == "" { + c.Routing.Policy = router.PolicyLeastInFlight + } } func (c *Config) validate() error { @@ -191,6 +207,12 @@ func (c *Config) validate() error { return fmt.Errorf("config: fencing.mode must be %q, %q, or %q, got %q", FenceStrict, FenceBounded, FenceRelaxed, c.Fencing.Mode) } + switch c.Routing.Policy { + case router.PolicyRoundRobin, router.PolicyLeastInFlight, router.PolicyScored: + default: + return fmt.Errorf("config: routing.policy must be %q, %q, or %q, got %q", + router.PolicyRoundRobin, router.PolicyLeastInFlight, router.PolicyScored, c.Routing.Policy) + } return nil } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 66b6e75..c7aaf86 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -139,6 +139,35 @@ func TestLoad_Fencing(t *testing.T) { } } +func TestLoad_Routing(t *testing.T) { + path := writeConfig(t, `{ + "listen": "x", "primary": "y", + "users": [{"name": "u", "password": "p"}], + "routing": {"policy": "scored"} + }`) + c, err := config.Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if c.Routing.Policy != "scored" { + t.Errorf("routing policy = %q, want scored", c.Routing.Policy) + } + + def := writeConfig(t, `{"listen":"x","primary":"y","users":[{"name":"u","password":"p"}]}`) + dc, err := config.Load(def) + if err != nil { + t.Fatalf("Load default: %v", err) + } + if dc.Routing.Policy != "least-in-flight" { + t.Errorf("default routing policy = %q, want least-in-flight", dc.Routing.Policy) + } + + bad := writeConfig(t, `{"listen":"x","primary":"y","users":[{"name":"u","password":"p"}],"routing":{"policy":"nope"}}`) + if _, err := config.Load(bad); err == nil { + t.Error("expected an error for an invalid routing policy") + } +} + func TestLoad_HealthDefaults(t *testing.T) { path := writeConfig(t, `{"listen":"x","primary":"y","users":[{"name":"u","password":"p"}]}`) c, err := config.Load(path) From 66776a33eca04436d07019331df34e47356b6a36 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:30:23 -0700 Subject: [PATCH 08/11] feat(proxy): route reads through the policy engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace "first eligible replica" with a real choice. chooseReadTarget now gathers every replica eligible under the fence and hands them, with their lag, to the shared routing policy, which picks one. The policy instance is built once and shared across all sessions so a load-aware policy sees the whole proxy's in-flight work. Each replica-routed read is bracketed by a Choose and exactly one Release carrying its measured latency, across every exit path — acquire failure, write error, and normal completion at ReadyForQuery — so in-flight counts and the scored policy's per-shape EWMA stay accurate and never leak. Reads that fall back to the primary and in-transaction reads never touch the policy. The fingerprint is computed only when the chosen policy needs it. --- cmd/pgpilot/main.go | 7 +++++ internal/proxy/proxy.go | 31 +++++++++++++++++++--- internal/proxy/session.go | 55 ++++++++++++++++++++++++++++++++++----- 3 files changed, 84 insertions(+), 9 deletions(-) diff --git a/cmd/pgpilot/main.go b/cmd/pgpilot/main.go index f384adb..adb6a11 100644 --- a/cmd/pgpilot/main.go +++ b/cmd/pgpilot/main.go @@ -19,6 +19,7 @@ import ( "github.com/sachhg/pgpilot/internal/config" "github.com/sachhg/pgpilot/internal/proxy" "github.com/sachhg/pgpilot/internal/registry" + "github.com/sachhg/pgpilot/internal/router" ) // version is the build version, overridden via -ldflags "-X main.version=...". @@ -91,11 +92,17 @@ func run(args []string) error { reg.Start(ctx, backendsFrom(cfg)) go reloadOnHUP(ctx, logger, *configPath, reg) + policy, err := router.New(cfg.Routing.Policy) + if err != nil { + return err + } + srv := proxy.New(proxy.Config{ ListenAddr: cfg.Listen, Users: cfg, Manager: mgr, Registry: reg, + Policy: policy, Logger: logger, }) addr, err := srv.Listen() diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index f3e9611..c0f412e 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -18,6 +18,7 @@ import ( "github.com/sachhg/pgpilot/internal/config" "github.com/sachhg/pgpilot/internal/protocol" "github.com/sachhg/pgpilot/internal/registry" + "github.com/sachhg/pgpilot/internal/router" ) // Config configures a proxy Server. @@ -32,6 +33,9 @@ type Config struct { // Registry reports backend health and lag for read routing. Nil disables // routing (every session is served by the primary). Registry *registry.Registry + // Policy selects which eligible replica serves each read. Nil resolves the + // policy named in Users.Routing.Policy, defaulting to least-in-flight. + Policy router.Policy // Logger receives structured logs. Nil selects slog.Default. Logger *slog.Logger } @@ -39,8 +43,9 @@ type Config struct { // Server is pgpilot's client-facing proxy. Construct it with New, then call // Listen and Serve. type Server struct { - cfg Config - log *slog.Logger + cfg Config + log *slog.Logger + policy router.Policy sessions atomic.Uint64 wg sync.WaitGroup @@ -55,7 +60,26 @@ func New(cfg Config) *Server { if logger == nil { logger = slog.Default() } - return &Server{cfg: cfg, log: logger} + return &Server{cfg: cfg, log: logger, policy: resolvePolicy(cfg)} +} + +// resolvePolicy returns cfg.Policy, or constructs the policy named in the config +// (defaulting to least-in-flight) when the caller did not inject one. The single +// returned instance is shared across every session, since load-aware policies +// track in-flight work across the whole proxy. +func resolvePolicy(cfg Config) router.Policy { + if cfg.Policy != nil { + return cfg.Policy + } + name := router.PolicyLeastInFlight + if cfg.Users != nil && cfg.Users.Routing.Policy != "" { + name = cfg.Users.Routing.Policy + } + p, err := router.New(name) + if err != nil { + return router.NewLeastInFlight() + } + return p } // Listen binds the configured listen address and returns the resolved address. @@ -128,6 +152,7 @@ func (s *Server) handle(ctx context.Context, client net.Conn) { cfg: s.cfg.Users, manager: s.cfg.Manager, registry: s.cfg.Registry, + policy: s.policy, log: log, tracker: &protocol.TxTracker{}, } diff --git a/internal/proxy/session.go b/internal/proxy/session.go index 752b3c0..e2d1c45 100644 --- a/internal/proxy/session.go +++ b/internal/proxy/session.go @@ -21,6 +21,7 @@ import ( "github.com/sachhg/pgpilot/internal/detect" "github.com/sachhg/pgpilot/internal/protocol" "github.com/sachhg/pgpilot/internal/registry" + "github.com/sachhg/pgpilot/internal/router" "github.com/sachhg/pgpilot/internal/scram" ) @@ -40,6 +41,7 @@ type session struct { cfg *config.Config manager *backend.Manager registry *registry.Registry + policy router.Policy log *slog.Logger tracker *protocol.TxTracker } @@ -428,14 +430,34 @@ func (s *session) serveRouting(ctx context.Context, user, database string) error _, breaks = detect.BreaksTxPooling(sql) } + // A read that the policy places on a replica owes exactly one Release, + // paired with this Choose, so its in-flight count and latency estimate + // stay accurate. routedAddr is empty for writes, in-transaction reads, + // and reads that fall back to the primary (no Choose happened). + var ( + routedAddr string + routedFP string + readStart time.Time + ) + // Choose the target backend when starting a new transaction. if held == nil { target := primary if class == classify.Read { - target = s.chooseReadTarget(fence) + fp := "" + if s.policy.NeedsFingerprint() { + fp = classify.Fingerprint(sql) + } + var routed bool + if target, routed = s.chooseReadTarget(fence, fp); routed { + routedAddr, routedFP = target, fp + } } held, err = s.manager.AcquireAt(ctx, user, database, target) if err != nil { + if routedAddr != "" { + s.policy.Release(routedAddr, routedFP, 0) + } return s.reject("53300", "could not acquire a backend connection") } heldAddr = target @@ -444,12 +466,22 @@ func (s *session) serveRouting(ctx context.Context, user, database string) error wroteThisTx = true } + if routedAddr != "" { + readStart = time.Now() + } if _, err := held.NetConn().Write(full); err != nil { + if routedAddr != "" { + s.policy.Release(routedAddr, routedFP, time.Since(readStart)) + } s.manager.DiscardAt(user, database, heldAddr, held) held = nil return nil } status, err := relayResponse(s.client, held.NetConn()) + if routedAddr != "" { + // The read has run to ReadyForQuery; release it whatever the outcome. + s.policy.Release(routedAddr, routedFP, time.Since(readStart)) + } if err != nil { s.manager.DiscardAt(user, database, heldAddr, held) held = nil @@ -480,15 +512,26 @@ func (s *session) serveRouting(ctx context.Context, user, database string) error } } -// chooseReadTarget returns the address of a replica eligible to serve a read -// under the current fence and fencing mode, or the primary if none qualifies. -func (s *session) chooseReadTarget(fence uint64) string { +// chooseReadTarget picks a backend to serve a read under the current fence and +// fencing mode. It collects the replicas eligible to serve the read and lets the +// routing policy choose among them; routed reports whether the policy was +// consulted (true) or the read fell back to the primary because no replica +// qualified (false), so the caller knows whether it owes the policy a Release. +func (s *session) chooseReadTarget(fence uint64, fingerprint string) (addr string, routed bool) { + var candidates []router.Candidate for _, st := range s.registry.Snapshot() { if st.Role == registry.RoleReplica && st.Healthy && s.replicaEligible(st, fence) { - return st.Addr + candidates = append(candidates, router.Candidate{ + Addr: st.Addr, + LagBytes: st.LagBytes, + LagSeconds: st.LagSeconds, + }) } } - return s.cfg.Primary + if len(candidates) == 0 { + return s.cfg.Primary, false + } + return s.policy.Choose(candidates, fingerprint), true } func (s *session) replicaEligible(st registry.Status, fence uint64) bool { From a4c8723a0dd60ef3de8096bec74d090585f48ee9 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:33:40 -0700 Subject: [PATCH 09/11] test(integration): assert reads spread across replicas Prove the policy distributes reads end to end, not just in the simulator. Under round-robin (least-in-flight would keep sending every sequential, immediately released read to the same idle replica), six reads through the proxy must land on more than one backend; each is identified by inet_server_addr(), its distinct container IP. startRoutingProxy gains a policy parameter so a test can pin the policy under test. --- test/integration/fencing_integration_test.go | 13 +++--- test/integration/routing_integration_test.go | 42 ++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) create mode 100644 test/integration/routing_integration_test.go diff --git a/test/integration/fencing_integration_test.go b/test/integration/fencing_integration_test.go index ce0fd7b..40304b8 100644 --- a/test/integration/fencing_integration_test.go +++ b/test/integration/fencing_integration_test.go @@ -14,14 +14,15 @@ import ( "github.com/sachhg/pgpilot/internal/config" "github.com/sachhg/pgpilot/internal/proxy" "github.com/sachhg/pgpilot/internal/registry" + "github.com/sachhg/pgpilot/internal/router" ) var replicaHostPorts = []string{"127.0.0.1:55433", "127.0.0.1:55434"} // startRoutingProxy launches an in-process routing proxy (primary + 2 replicas) // in the given fencing mode and returns its port once the registry has polled -// every backend healthy. -func startRoutingProxy(t *testing.T, mode string) int { +// every backend healthy. A nil policy leaves the proxy to resolve its default. +func startRoutingProxy(t *testing.T, mode string, policy router.Policy) int { t.Helper() cfg := &config.Config{ Primary: primaryHostPort, @@ -52,7 +53,7 @@ func startRoutingProxy(t *testing.T, mode string) int { {Name: "replica2", Addr: replicaHostPorts[1]}, }) - srv := proxy.New(proxy.Config{ListenAddr: "0.0.0.0:0", Users: cfg, Manager: mgr, Registry: reg, Logger: log}) + srv := proxy.New(proxy.Config{ListenAddr: "0.0.0.0:0", Users: cfg, Manager: mgr, Registry: reg, Policy: policy, Logger: log}) addr, err := srv.Listen() if err != nil { t.Fatalf("proxy listen: %v", err) @@ -106,7 +107,7 @@ func resumeReplication(t *testing.T, compose string) { func TestProxy_RoutesReadsToReplicas(t *testing.T) { compose := requireCluster(t) resumeReplication(t, compose) - port := startRoutingProxy(t, config.FenceStrict) + port := startRoutingProxy(t, config.FenceStrict, nil) // A read with no write fence is served by a replica. if got, err := runPsql(compose, "host.docker.internal", port, "SELECT pg_is_in_recovery()"); err != nil { @@ -148,7 +149,7 @@ func TestProxy_LSNFencing_StrictNeverStale(t *testing.T) { // Strict: write 'new' (fence advances on the primary) then read — the frozen // replicas are behind the fence, so the read falls back to the primary. - strict := startRoutingProxy(t, config.FenceStrict) + strict := startRoutingProxy(t, config.FenceStrict, nil) got, err := runPsqlSession(compose, "host.docker.internal", strict, "UPDATE fence_test SET v = 'new' WHERE id = 1", "SELECT v FROM fence_test WHERE id = 1") @@ -160,7 +161,7 @@ func TestProxy_LSNFencing_StrictNeverStale(t *testing.T) { } // Relaxed: the same read is served by a frozen replica and is stale. - relaxed := startRoutingProxy(t, config.FenceRelaxed) + relaxed := startRoutingProxy(t, config.FenceRelaxed, nil) rgot, err := runPsql(compose, "host.docker.internal", relaxed, "SELECT v FROM fence_test WHERE id = 1") if err != nil { t.Fatalf("relaxed session: %v", err) diff --git a/test/integration/routing_integration_test.go b/test/integration/routing_integration_test.go new file mode 100644 index 0000000..06f1278 --- /dev/null +++ b/test/integration/routing_integration_test.go @@ -0,0 +1,42 @@ +//go:build integration + +package integration + +import ( + "testing" + + "github.com/sachhg/pgpilot/internal/config" + "github.com/sachhg/pgpilot/internal/router" +) + +// TestProxy_RoundRobinSpreadsReadsAcrossReplicas proves the routing policy +// actually distributes reads over the whole replica set, end to end. Under the +// round-robin policy (least-in-flight would send every sequential, immediately +// released read to the same idle replica), a handful of reads must land on more +// than one replica. Each backend is identified by inet_server_addr(), which +// returns its distinct container IP. +func TestProxy_RoundRobinSpreadsReadsAcrossReplicas(t *testing.T) { + compose := requireCluster(t) + resumeReplication(t, compose) + // Relaxed fencing keeps every healthy replica eligible, isolating the test to + // the policy's distribution rather than the fence. + port := startRoutingProxy(t, config.FenceRelaxed, router.NewRoundRobin()) + + seen := map[string]int{} + const reads = 6 + for i := 0; i < reads; i++ { + addr, err := runPsql(compose, "host.docker.internal", port, "SELECT inet_server_addr()") + if err != nil { + t.Fatalf("read %d: %v", i, err) + } + if addr == "" { + t.Fatalf("read %d returned an empty server address", i) + } + seen[addr]++ + } + if len(seen) < 2 { + t.Errorf("round-robin sent %d reads to a single backend %v; want them spread across both replicas", reads, seen) + } else { + t.Logf("round-robin spread %d reads across %d replicas: %v", reads, len(seen), seen) + } +} From c7f6f41471d96a99da3446a99663f402f6008c68 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:34:19 -0700 Subject: [PATCH 10/11] docs(adr): record the routing policy engine design Explain why replica selection sits behind a Policy interface separate from eligibility, what each of the three policies optimizes, the scored formula and why its EWMA is keyed by pg_query fingerprint, and why least-in-flight is the default while scored is opt-in. Records the alternatives rejected (random, lag-only, power-of-two-choices) and the known trade-offs: the scored policy's median-for-tail trade and its queue-inclusive latency estimate. --- docs/adr/0009-routing-policy-engine.md | 79 ++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 docs/adr/0009-routing-policy-engine.md diff --git a/docs/adr/0009-routing-policy-engine.md b/docs/adr/0009-routing-policy-engine.md new file mode 100644 index 0000000..db3dbf2 --- /dev/null +++ b/docs/adr/0009-routing-policy-engine.md @@ -0,0 +1,79 @@ +# 9. Routing policy engine + +Date: 2026-07-23 + +## Status + +Accepted. + +## Context + +Phase 7 made pgpilot route reads to replicas that have replayed past the +session's fence, but among the replicas that qualify it simply took the first +one. On a cluster with more than one eligible replica that concentrates every +read on the same backend, wasting the rest and ignoring that replicas differ in +load, speed, and how expensive a given query is on each. This phase decides +*which* eligible replica serves a read. + +Eligibility (healthy, and fresh enough under the fencing rules) already lives in +the proxy and is a correctness question. Selection among the eligible is a +performance question with real alternatives, so it belongs behind an interface. + +## Decision + +**A `Policy` interface, separate from eligibility.** `internal/router` defines +`Policy.Choose(candidates, fingerprint) addr` over the replicas the proxy has +already deemed eligible. Because eligibility is decided upstream, a policy never +needs a database or the registry: it is pure Go, exercised in unit tests and a +deterministic simulator with no cluster. Every `Choose` is paired with one +`Release(addr, fingerprint, latency)` so a load-aware policy tracks in-flight +work and observed latency; the single policy instance is shared across all +sessions so it sees the whole proxy's load. + +**Three implementations, increasing in sophistication.** + +- **round-robin** — even rotation. Fair over a stable candidate set, keeps no + state, ignores load and query cost. +- **least-in-flight** — send each read to the eligible replica with the fewest + outstanding reads. Adapts to a slow or briefly overloaded replica for free, + because its reads pile up and it stops being chosen until it drains. This is + the default: strong, zero-config, and it costs no per-read parse. +- **scored** — rank candidates by estimated completion time, + `(inFlight + 1) * ewmaLatency(addr, fingerprint) + lagPenalty * lagSeconds`, + and take the minimum. Keying the EWMA by pg_query fingerprint is the point: it + lets the policy steer an *expensive* query shape away from a replica already + busy with that shape while still sending cheap queries there, and it learns + per-replica speed differences as slower backends accrue higher averages. + +**least-in-flight is the default, scored is opt-in.** The scored policy needs a +pg_query fingerprint per read — a cgo parse — and only pays off when replicas or +query costs are heterogeneous. Least-in-flight already handles slow and +overloaded replicas without that cost, so it is the safe default; `routing.policy` +selects another, validated at load time against the router's own names. + +## Consequences + +- A deterministic discrete-event simulator (seeded arrivals, a heavy-tailed + shape mix, one 3x-slow replica, each replica a single FIFO server) compares + the policies with no database. It shows the failure it is meant to: blind + round-robin forces a third of all reads onto the slow replica and its queue + explodes (p95 into the seconds), while least-in-flight and scored back off it — + and scored, alone knowing heavy shapes are dear there, sends it the fewest + reads and wins the tail. Micro-benchmarks confirm per-decision cost stays + sub-microsecond. +- The scored policy trades slightly higher median latency for a better tail: it + concentrates load on the fast replicas to spare the slow one, so the fast + replicas queue a little more at p50 while p95/p99 improve. That is the right + trade for a router, but it is a trade. +- The EWMA folds in queueing time, not pure service time, because that is what + the proxy can measure — the latency from dispatch to ReadyForQuery. Under load + the estimate therefore runs a little hot, which reinforces the correct steering + but is not a clean per-shape cost model. A server-reported cost would be + cleaner and is future work. +- The scored weights are constants tuned for comparable replicas. A cluster with + a large staleness budget or very heterogeneous hardware may want them tunable + from config; they are exported so that remains an additive change. +- Rejected: random choice (no load awareness), lag-only weighting (blind to load + and query cost — the very thing that overloads a replica), and power-of-two- + choices (a good cheap approximation of least-in-flight, but with only a handful + of replicas exact least-in-flight is affordable and strictly better). From 90094b3e29a6e4e318c24d46135448fba24717b5 Mon Sep 17 00:00:00 2001 From: sachhg Date: Thu, 23 Jul 2026 21:35:39 -0700 Subject: [PATCH 11/11] docs: mark Phase 8 done and document routing policies Flip the roadmap to Phase 8 done / Phase 9 next, describe the three routing policies and the deterministic comparison in the README, add the CHANGELOG entry, and show routing.policy in the example config. --- CHANGELOG.md | 10 ++++++++++ README.md | 31 ++++++++++++++++++++++++++----- pgpilot.example.json | 3 +++ 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f27315b..023b58e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -52,6 +52,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 integration test proves that, with replication paused via `pg_wal_replay_pause()`, a write followed by a read never observes a stale value under strict mode. Rationale in `docs/adr/0008-lsn-fencing.md`. +- Routing policy engine (Phase 8): `internal/router` makes replica selection + pluggable behind a `Policy` interface, choosing among the replicas the fence + has already deemed eligible. Three policies ship — `round-robin`, + `least-in-flight` (the default), and `scored`, which ranks replicas by + estimated completion time using an EWMA of each query shape's latency keyed by + pg_query fingerprint, steering expensive shapes away from busy replicas. + `routing.policy` selects one. A deterministic discrete-event simulator + compares the policies on a synthetic mixed workload with no database, and an + integration test asserts reads spread across both replicas. Rationale in + `docs/adr/0009-routing-policy-engine.md`. ### Dependencies diff --git a/README.md b/README.md index 675127f..1b36455 100644 --- a/README.md +++ b/README.md @@ -42,8 +42,9 @@ The goal is to do one thing — correct, observable read/write routing — well. Early development, built in phases (see the roadmap). Not production-ready yet. pgpilot now **routes**: it authenticates each client with SCRAM-SHA-256, pools connections, classifies each query, and sends writes to the primary and reads to -a replica — enforcing read-your-writes with a per-session LSN fence. A routing -*policy* (scoring replicas by lag and load) and observability come next. +a replica — enforcing read-your-writes with a per-session LSN fence — and +balances reads across eligible replicas with a selectable routing policy +(round-robin, least-in-flight, or latency-scored). Observability comes next. ## Roadmap @@ -57,8 +58,8 @@ a replica — enforcing read-your-writes with a per-session LSN fence. A routing | 5 | Query classification (read vs. write via pg_query) | done | | 6 | Replica registry, health polling, circuit breakers | done | | 7 | LSN fencing | done | -| 8 | Routing policy engine | next | -| 9 | Observability (Prometheus, structured logs, pprof) | | +| 8 | Routing policy engine | done | +| 9 | Observability (Prometheus, structured logs, pprof) | next | | 10 | Fault-injection harness | | | 11 | Benchmarks vs. direct connection and pgbouncer | | | 12 | Docs and the v0.1.0 release | | @@ -106,7 +107,8 @@ pgpilot reads a JSON config file (see [`pgpilot.example.json`](pgpilot.example.j "users": [{"name": "pgpilot", "password": "pgpilot"}], "pool": {"mode": "session", "max_size": 10, "acquire_timeout": "5s", "idle_timeout": "5m"}, "health": {"interval": "1s", "failure_threshold": 3, "base_backoff": "1s", "max_backoff": "30s"}, - "fencing": {"mode": "strict", "bounded_ms": 100} + "fencing": {"mode": "strict", "bounded_ms": 100}, + "routing": {"policy": "least-in-flight"} } ``` @@ -142,6 +144,25 @@ primary). `fencing.mode` selects the trade-off: stale value under strict mode. Design in [`docs/adr/0008-lsn-fencing.md`](docs/adr/0008-lsn-fencing.md). +### Routing policies + +Fencing decides which replicas *may* serve a read; `routing.policy` decides which +one *does* when more than one qualifies: + +- **`round-robin`** — even rotation across eligible replicas. +- **`least-in-flight`** (default) — the eligible replica with the fewest reads + outstanding, so a slow or overloaded replica sheds load until it drains. +- **`scored`** — ranks replicas by estimated completion time, + `(inFlight + 1) * ewmaLatency(addr, shape) + lagPenalty * lag`, learning each + query shape's cost per replica (keyed by pg_query fingerprint) so it steers + expensive shapes away from busy replicas. It costs one fingerprint parse per + read, which is why it is opt-in. + +`go test -run WorkloadComparison -v ./internal/router` runs a deterministic +simulation comparing the policies on a synthetic mixed workload with no database. +Design in +[`docs/adr/0009-routing-policy-engine.md`](docs/adr/0009-routing-policy-engine.md). + ### Health and replication lag A background poller (`internal/registry`) tracks each backend's role and diff --git a/pgpilot.example.json b/pgpilot.example.json index 53b3bdc..3d51777 100644 --- a/pgpilot.example.json +++ b/pgpilot.example.json @@ -21,5 +21,8 @@ "fencing": { "mode": "strict", "bounded_ms": 100 + }, + "routing": { + "policy": "least-in-flight" } }