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/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/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). 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) + } + } +} 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) 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 { 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) + } +} 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() +} 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") + } +} 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() +} 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()) } 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" } } 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) + } +}