From f30b56891297030bfb082abd1eb1cbb149774868 Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 14:25:36 -0700 Subject: [PATCH 1/2] feat(pool): add a bounded, health-checked connection pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit internal/pool caps the number of open backend connections, hands idle ones back out after an optional health check, reaps connections idle past a timeout, and applies backpressure — a bounded wait with AcquireTimeout, or an immediate ErrExhausted once MaxWaiters is reached — instead of queueing acquirers without bound. Acquire never opens more than MaxSize connections even under concurrency, and Close drains idle connections and wakes waiters. This is the primitive the session- and transaction-pooling layers will build on. Covered by race-enabled unit tests, including a concurrency test asserting the cap is never exceeded. --- internal/pool/pool.go | 343 +++++++++++++++++++++++++++++++++++++ internal/pool/pool_test.go | 337 ++++++++++++++++++++++++++++++++++++ 2 files changed, 680 insertions(+) create mode 100644 internal/pool/pool.go create mode 100644 internal/pool/pool_test.go diff --git a/internal/pool/pool.go b/internal/pool/pool.go new file mode 100644 index 0000000..3ce5f63 --- /dev/null +++ b/internal/pool/pool.go @@ -0,0 +1,343 @@ +// Package pool provides a bounded, health-checked pool of backend connections. +// +// A Pool caps the number of open connections, hands idle connections back out +// after an optional health check, reaps connections that sit idle too long, and +// applies backpressure — a bounded wait with a timeout, or an immediate refusal +// once a waiter limit is reached — instead of queueing acquirers without limit. +// +// Each connection returned by Acquire must be returned exactly once with either +// Release (reuse it) or Discard (throw it away, e.g. after a protocol error). +package pool + +import ( + "context" + "errors" + "fmt" + "sync" + "time" +) + +// Conn is a backend connection managed by a Pool. Close must be safe to call +// from the pool's own goroutines. +type Conn interface { + Close() error +} + +// Errors returned by Acquire. +var ( + // ErrClosed is returned by Acquire after the pool has been closed. + ErrClosed = errors.New("pool: closed") + // ErrTimeout is returned when Acquire waits longer than AcquireTimeout. + ErrTimeout = errors.New("pool: acquire timed out") + // ErrExhausted is returned when the pool is at capacity and the waiter + // limit (MaxWaiters) has been reached. + ErrExhausted = errors.New("pool: exhausted") +) + +// Config configures a Pool. +type Config struct { + // MaxSize is the maximum number of open connections. Must be > 0. + MaxSize int + // MaxWaiters caps the goroutines that may wait for a connection once MaxSize + // is reached; further Acquire calls fail immediately with ErrExhausted. Zero + // means unlimited waiters (each still bounded by AcquireTimeout). + MaxWaiters int + // AcquireTimeout bounds how long Acquire waits for a connection. Zero waits + // until the caller's context is done. + AcquireTimeout time.Duration + // IdleTimeout closes connections that sit idle in the pool longer than this. + // Zero disables idle reaping. + IdleTimeout time.Duration + // New opens a fresh backend connection. Required. + New func(ctx context.Context) (Conn, error) + // Ping checks that a reused connection is still usable before Acquire hands + // it out; a non-nil error discards the connection and a replacement is used. + // Nil skips the check. + Ping func(ctx context.Context, c Conn) error +} + +// Pool is a bounded pool of backend connections. The zero value is not usable; +// construct one with New. +type Pool struct { + cfg Config + + mu sync.Mutex + idle []idleConn + open int // total open: idle + handed out + waiters []chan Conn // FIFO queue of goroutines waiting for a connection + closed bool + + stopReaper chan struct{} + reaperDone chan struct{} +} + +type idleConn struct { + conn Conn + idleSince time.Time +} + +// New constructs a Pool from cfg. +func New(cfg Config) (*Pool, error) { + if cfg.MaxSize <= 0 { + return nil, fmt.Errorf("pool: MaxSize must be > 0, got %d", cfg.MaxSize) + } + if cfg.New == nil { + return nil, errors.New("pool: New is required") + } + p := &Pool{cfg: cfg} + if cfg.IdleTimeout > 0 { + p.stopReaper = make(chan struct{}) + p.reaperDone = make(chan struct{}) + go p.reap() + } + return p, nil +} + +// Acquire returns a healthy connection, creating one if the pool is below +// MaxSize, or waiting for one to be released otherwise. It respects +// AcquireTimeout and ctx cancellation, and never opens more than MaxSize +// connections. +func (p *Pool) Acquire(ctx context.Context) (Conn, error) { + var deadline time.Time + if p.cfg.AcquireTimeout > 0 { + deadline = time.Now().Add(p.cfg.AcquireTimeout) + } + for { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return nil, ErrClosed + } + + // 1. Reuse an idle connection if one is available. + if n := len(p.idle); n > 0 { + ic := p.idle[n-1] + p.idle = p.idle[:n-1] + if p.cfg.IdleTimeout > 0 && time.Since(ic.idleSince) >= p.cfg.IdleTimeout { + p.open-- + p.mu.Unlock() + _ = ic.conn.Close() + continue + } + p.mu.Unlock() + if c, ok := p.checkReused(ctx, ic.conn); ok { + return c, nil + } + continue + } + + // 2. Open a new connection if under capacity. + if p.open < p.cfg.MaxSize { + p.open++ + p.mu.Unlock() + c, err := p.cfg.New(ctx) + if err != nil { + p.mu.Lock() + p.open-- + p.mu.Unlock() + return nil, err + } + return c, nil + } + + // 3. At capacity: wait for a released connection (backpressure). + if p.cfg.MaxWaiters > 0 && len(p.waiters) >= p.cfg.MaxWaiters { + p.mu.Unlock() + return nil, ErrExhausted + } + w := make(chan Conn, 1) + p.waiters = append(p.waiters, w) + p.mu.Unlock() + + c, err := p.waitFor(ctx, w, deadline) + if err != nil { + return nil, err + } + if hc, ok := p.checkReused(ctx, c); ok { + return hc, nil + } + // The released connection failed its health check and was discarded; + // loop to open a fresh one in its place. + } +} + +// checkReused health-checks a reused connection. On failure it discards the +// connection (closes it and frees its slot) and returns ok=false so the caller +// tries again. +func (p *Pool) checkReused(ctx context.Context, c Conn) (Conn, bool) { + if p.cfg.Ping == nil { + return c, true + } + if err := p.cfg.Ping(ctx, c); err != nil { + p.mu.Lock() + p.open-- + p.mu.Unlock() + _ = c.Close() + return nil, false + } + return c, true +} + +func (p *Pool) waitFor(ctx context.Context, w chan Conn, deadline time.Time) (Conn, error) { + var timeout <-chan time.Time + if !deadline.IsZero() { + t := time.NewTimer(time.Until(deadline)) + defer t.Stop() + timeout = t.C + } + select { + case c := <-w: + if c == nil { + return nil, ErrClosed + } + return c, nil + case <-timeout: + return nil, p.abandonWaiter(w, ErrTimeout) + case <-ctx.Done(): + return nil, p.abandonWaiter(w, ctx.Err()) + } +} + +// abandonWaiter removes w from the waiter queue and returns err. If a releaser +// already delivered a connection to w, that connection is returned to the pool +// instead of being lost. +func (p *Pool) abandonWaiter(w chan Conn, err error) error { + p.mu.Lock() + for i, x := range p.waiters { + if x == w { + p.waiters = append(p.waiters[:i], p.waiters[i+1:]...) + p.mu.Unlock() + return err + } + } + p.mu.Unlock() + if c := <-w; c != nil { + p.Release(c) + } + return err +} + +// Release returns a connection to the pool for reuse. +func (p *Pool) Release(c Conn) { + if c == nil { + return + } + p.mu.Lock() + if p.closed { + p.open-- + p.mu.Unlock() + _ = c.Close() + return + } + if len(p.waiters) > 0 { + w := p.waiters[0] + p.waiters = p.waiters[1:] + w <- c + p.mu.Unlock() + return + } + p.idle = append(p.idle, idleConn{conn: c, idleSince: time.Now()}) + p.mu.Unlock() +} + +// Discard closes a connection and frees its slot instead of returning it to the +// pool. Use it for a connection left in an unknown state, e.g. after an error. +func (p *Pool) Discard(c Conn) { + if c == nil { + return + } + p.mu.Lock() + p.open-- + p.mu.Unlock() + _ = c.Close() +} + +// Close drains the pool: it closes every idle connection, wakes waiters with +// ErrClosed, and rejects further Acquire calls. Connections still handed out are +// closed when they are Released or Discarded. +func (p *Pool) Close() error { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return nil + } + p.closed = true + idle := p.idle + p.idle = nil + waiters := p.waiters + p.waiters = nil + p.mu.Unlock() + + if p.stopReaper != nil { + close(p.stopReaper) + <-p.reaperDone + } + for _, w := range waiters { + w <- nil // wakes the waiter, which returns ErrClosed + } + for _, ic := range idle { + _ = ic.conn.Close() + } + return nil +} + +// Stats is a snapshot of a pool's connection counts. +type Stats struct { + Open int // total open connections (idle + in use) + Idle int // connections sitting idle in the pool + InUse int // connections currently handed out + Waiters int // goroutines waiting for a connection +} + +// Stats returns a snapshot of the pool's current counts. +func (p *Pool) Stats() Stats { + p.mu.Lock() + defer p.mu.Unlock() + return Stats{ + Open: p.open, + Idle: len(p.idle), + InUse: p.open - len(p.idle), + Waiters: len(p.waiters), + } +} + +func (p *Pool) reap() { + defer close(p.reaperDone) + interval := p.cfg.IdleTimeout + if interval > 30*time.Second { + interval = 30 * time.Second + } + if interval < 10*time.Millisecond { + interval = 10 * time.Millisecond + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-p.stopReaper: + return + case <-ticker.C: + p.reapOnce() + } + } +} + +func (p *Pool) reapOnce() { + now := time.Now() + var toClose []Conn + p.mu.Lock() + kept := make([]idleConn, 0, len(p.idle)) + for _, ic := range p.idle { + if now.Sub(ic.idleSince) >= p.cfg.IdleTimeout { + p.open-- + toClose = append(toClose, ic.conn) + } else { + kept = append(kept, ic) + } + } + p.idle = kept + p.mu.Unlock() + for _, c := range toClose { + _ = c.Close() + } +} diff --git a/internal/pool/pool_test.go b/internal/pool/pool_test.go new file mode 100644 index 0000000..5222232 --- /dev/null +++ b/internal/pool/pool_test.go @@ -0,0 +1,337 @@ +package pool_test + +import ( + "context" + "errors" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/sachhg/pgpilot/internal/pool" +) + +type fakeConn struct { + id int + closed atomic.Bool + pingErr error +} + +func (c *fakeConn) Close() error { + c.closed.Store(true) + return nil +} + +type factory struct { + mu sync.Mutex + created int + newErr error +} + +func (f *factory) New(_ context.Context) (pool.Conn, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.newErr != nil { + return nil, f.newErr + } + f.created++ + return &fakeConn{id: f.created}, nil +} + +func (f *factory) count() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.created +} + +func pingByConn(_ context.Context, c pool.Conn) error { + return c.(*fakeConn).pingErr +} + +func waitFor(t *testing.T, cond func() bool) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(2 * time.Millisecond) + } + t.Fatal("condition not met within deadline") +} + +func TestNew_InvalidConfig(t *testing.T) { + if _, err := pool.New(pool.Config{MaxSize: 0, New: (&factory{}).New}); err == nil { + t.Error("expected an error for MaxSize <= 0") + } + if _, err := pool.New(pool.Config{MaxSize: 1}); err == nil { + t.Error("expected an error for a nil New") + } +} + +func TestAcquireRelease_ReusesConnection(t *testing.T) { + f := &factory{} + p, err := pool.New(pool.Config{MaxSize: 2, New: f.New}) + if err != nil { + t.Fatal(err) + } + defer p.Close() + + c1, err := p.Acquire(context.Background()) + if err != nil { + t.Fatalf("acquire: %v", err) + } + p.Release(c1) + c2, err := p.Acquire(context.Background()) + if err != nil { + t.Fatalf("acquire again: %v", err) + } + if c1 != c2 { + t.Error("expected the released connection to be reused") + } + if f.count() != 1 { + t.Errorf("created %d connections, want 1", f.count()) + } +} + +func TestMaxSize_BackpressureTimeout(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 2, AcquireTimeout: 50 * time.Millisecond, New: f.New}) + defer p.Close() + + a, _ := p.Acquire(context.Background()) + _, _ = p.Acquire(context.Background()) + + if _, err := p.Acquire(context.Background()); !errors.Is(err, pool.ErrTimeout) { + t.Fatalf("third acquire err = %v, want ErrTimeout", err) + } + if s := p.Stats(); s.Open != 2 || s.InUse != 2 { + t.Fatalf("stats = %+v, want Open=2 InUse=2", s) + } + + p.Release(a) + if _, err := p.Acquire(context.Background()); err != nil { + t.Fatalf("acquire after release: %v", err) + } + if f.count() != 2 { + t.Errorf("created %d, want 2 (never exceed MaxSize)", f.count()) + } +} + +func TestWaiter_GetsReleasedConnection(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 1, AcquireTimeout: 2 * time.Second, New: f.New}) + defer p.Close() + + c1, _ := p.Acquire(context.Background()) + got := make(chan pool.Conn, 1) + go func() { + c, err := p.Acquire(context.Background()) + if err != nil { + got <- nil + return + } + got <- c + }() + waitFor(t, func() bool { return p.Stats().Waiters == 1 }) + + p.Release(c1) + select { + case c := <-got: + if c != c1 { + t.Error("waiter should receive the released connection") + } + case <-time.After(2 * time.Second): + t.Fatal("waiter was not served") + } +} + +func TestMaxWaiters_RefusesImmediately(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 1, MaxWaiters: 1, AcquireTimeout: time.Second, New: f.New}) + defer p.Close() + + _, _ = p.Acquire(context.Background()) + go func() { _, _ = p.Acquire(context.Background()) }() // fills the single waiter slot + waitFor(t, func() bool { return p.Stats().Waiters == 1 }) + + if _, err := p.Acquire(context.Background()); !errors.Is(err, pool.ErrExhausted) { + t.Fatalf("err = %v, want ErrExhausted", err) + } +} + +func TestPing_DiscardsUnhealthyConnection(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 1, New: f.New, Ping: pingByConn}) + defer p.Close() + + c1, _ := p.Acquire(context.Background()) + c1.(*fakeConn).pingErr = errors.New("dead") + p.Release(c1) + + c2, err := p.Acquire(context.Background()) + if err != nil { + t.Fatalf("acquire: %v", err) + } + if c2 == c1 { + t.Error("unhealthy connection should not be reused") + } + if !c1.(*fakeConn).closed.Load() { + t.Error("unhealthy connection should be closed") + } + if f.count() != 2 { + t.Errorf("created %d, want 2 (replacement opened)", f.count()) + } +} + +func TestIdleTimeout_ReapsIdleConnections(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 2, IdleTimeout: 40 * time.Millisecond, New: f.New}) + defer p.Close() + + c1, _ := p.Acquire(context.Background()) + c2, _ := p.Acquire(context.Background()) + p.Release(c1) + p.Release(c2) + + waitFor(t, func() bool { return p.Stats().Open == 0 }) + if !c1.(*fakeConn).closed.Load() || !c2.(*fakeConn).closed.Load() { + t.Error("idle connections should be closed by the reaper") + } +} + +func TestClose_ClosesIdleAndRejects(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 2, New: f.New}) + + c1, _ := p.Acquire(context.Background()) + p.Release(c1) + _ = p.Close() + + if !c1.(*fakeConn).closed.Load() { + t.Error("idle connection should be closed on Close") + } + if _, err := p.Acquire(context.Background()); !errors.Is(err, pool.ErrClosed) { + t.Fatalf("acquire after close = %v, want ErrClosed", err) + } +} + +func TestClose_WakesWaiters(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 1, AcquireTimeout: 5 * time.Second, New: f.New}) + + _, _ = p.Acquire(context.Background()) + errc := make(chan error, 1) + go func() { + _, err := p.Acquire(context.Background()) + errc <- err + }() + waitFor(t, func() bool { return p.Stats().Waiters == 1 }) + + _ = p.Close() + select { + case err := <-errc: + if !errors.Is(err, pool.ErrClosed) { + t.Fatalf("waiter err = %v, want ErrClosed", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Close did not wake the waiter") + } +} + +func TestAcquire_ContextCancel(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 1, New: f.New}) + defer p.Close() + + _, _ = p.Acquire(context.Background()) + ctx, cancel := context.WithCancel(context.Background()) + errc := make(chan error, 1) + go func() { + _, err := p.Acquire(ctx) + errc <- err + }() + waitFor(t, func() bool { return p.Stats().Waiters == 1 }) + + cancel() + select { + case err := <-errc: + if !errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want context.Canceled", err) + } + case <-time.After(2 * time.Second): + t.Fatal("context cancellation not honored") + } +} + +func TestNew_DialErrorDoesNotLeakSlot(t *testing.T) { + f := &factory{newErr: errors.New("dial failed")} + p, _ := pool.New(pool.Config{MaxSize: 2, New: f.New}) + defer p.Close() + + if _, err := p.Acquire(context.Background()); err == nil || !strings.Contains(err.Error(), "dial failed") { + t.Fatalf("err = %v, want dial failed", err) + } + if s := p.Stats(); s.Open != 0 { + t.Errorf("open = %d after a failed dial, want 0", s.Open) + } +} + +func TestDiscard_FreesSlot(t *testing.T) { + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: 1, New: f.New}) + defer p.Close() + + c1, _ := p.Acquire(context.Background()) + p.Discard(c1) + if !c1.(*fakeConn).closed.Load() { + t.Error("discarded connection should be closed") + } + if s := p.Stats(); s.Open != 0 { + t.Errorf("open = %d after discard, want 0", s.Open) + } + // The freed slot allows a new connection. + if _, err := p.Acquire(context.Background()); err != nil { + t.Fatalf("acquire after discard: %v", err) + } + if f.count() != 2 { + t.Errorf("created %d, want 2", f.count()) + } +} + +func TestConcurrent_NeverExceedsMaxSize(t *testing.T) { + const maxSize = 4 + f := &factory{} + p, _ := pool.New(pool.Config{MaxSize: maxSize, AcquireTimeout: 5 * time.Second, New: f.New}) + defer p.Close() + + var inUse atomic.Int32 + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + c, err := p.Acquire(context.Background()) + if err != nil { + t.Errorf("acquire: %v", err) + return + } + if n := inUse.Add(1); n > maxSize { + t.Errorf("in-use %d exceeds MaxSize %d", n, maxSize) + } + inUse.Add(-1) + p.Release(c) + } + }() + } + wg.Wait() + + if s := p.Stats(); s.Open > maxSize { + t.Errorf("open %d exceeds MaxSize %d", s.Open, maxSize) + } + if f.count() > maxSize { + t.Errorf("created %d exceeds MaxSize %d", f.count(), maxSize) + } +} From 98276d6dbf9edd0daa692f5c56933e71c57c4766 Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 14:25:36 -0700 Subject: [PATCH 2/2] docs(changelog): note the connection pool primitive Record internal/pool (Phase 4a): bounded size, acquire/idle timeouts, health checks, and backpressure, under the Unreleased section. --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b9afcfe..5258faf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 bytes and tracks each session's transaction status from ReadyForQuery (`I`/`T`/`E`). `cmd/pgpilot` gains a `-log-level` flag. Rationale recorded in `docs/adr/0003-message-aware-relay.md`. +- Connection pool (Phase 4a): `internal/pool` is a bounded, health-checked pool + of backend connections — configurable max size, acquire timeout, idle timeout + with a background reaper, and per-connection health checks — that applies + backpressure (a bounded wait, or an immediate refusal once a waiter limit is + reached) instead of queueing acquirers without bound. It is the primitive the + session- and transaction-pooling layers will build on. ### Dependencies