From 0916380709c5f8d68bbd2c20f3df949c430f3c94 Mon Sep 17 00:00:00 2001 From: jearthliu Date: Fri, 7 Aug 2026 10:03:13 +0800 Subject: [PATCH 1/2] fix: prevent connection storm during outage recovery --- pgxpool/pool.go | 137 ++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 122 insertions(+), 15 deletions(-) diff --git a/pgxpool/pool.go b/pgxpool/pool.go index 00afc75..bc3dc67 100644 --- a/pgxpool/pool.go +++ b/pgxpool/pool.go @@ -2,35 +2,142 @@ package pgxpool import ( "context" + "errors" "sync" "sync/atomic" ) +// conn represents an established PostgreSQL connection. +type conn struct { + id int +} + +// Conn is the handle returned by Acquire. It wraps an underlying connection +// and a reference back to the pool so Release can return it. +type Conn struct { + c *conn + pool *Pool +} + +// Release returns the connection to the pool and wakes a waiting acquirer. +func (c *Conn) Release() { + c.pool.release(c.c) +} + +// pool is a connection pool that strictly respects MaxConns even during +// recovery from an outage: pending dial attempts count toward capacity, so +// a burst of blocked Acquire() calls cannot each start its own dial. type Pool struct { - // ... existing fields maxConns int32 conns []*conn - inFlightConns int32 + inFlightConns int32 // pending dials — counted toward capacity + idCounter int32 mu sync.Mutex - // ... + cond *sync.Cond + closed bool +} + +// New creates a pool with the given maximum number of connections. +func New(maxConns int32) *Pool { + p := &Pool{maxConns: maxConns} + p.cond = sync.NewCond(&p.mu) + return p } +// ErrPoolClosed is returned by Acquire after Close. +var ErrPoolClosed = errors.New("pgxpool: pool closed") + +// Acquire returns a connection, dialing a new one if capacity allows and +// blocking otherwise. The check-and-increment of inFlightConns happens under +// the pool mutex, so concurrent callers cannot collectively overshoot +// MaxConns during recovery. func (p *Pool) Acquire(ctx context.Context) (*Conn, error) { - p.mu.Lock() - // Check if we can create a new connection - if len(p.conns) + int(atomic.LoadInt32(&p.inFlightConns)) < int(p.maxConns) { - atomic.AddInt32(&p.inFlightConns, 1) - p.mu.Unlock() + for { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return nil, ErrPoolClosed + } + + // Fast path: reuse an idle connection. + if n := len(p.conns); n > 0 { + c := p.conns[n-1] + p.conns = p.conns[:n-1] + p.mu.Unlock() + return &Conn{c: c, pool: p}, nil + } + + // Capacity check: established + pending must stay under MaxConns. + if int32(len(p.conns))+atomic.LoadInt32(&p.inFlightConns) < p.maxConns { + atomic.AddInt32(&p.inFlightConns, 1) + p.mu.Unlock() + + // Dial outside the lock; the pending counter holds capacity. + c, err := p.dial(ctx) + if err != nil { + // Decrement on every failure path (dial timeout, ctx + // cancellation, auth failure) so the pool never starves. + atomic.AddInt32(&p.inFlightConns, -1) + p.cond.Signal() + return nil, err + } - conn, err := p.createNewConn(ctx) - atomic.AddInt32(&p.inFlightConns, -1) - if err != nil { + // Register the connection BEFORE releasing capacity: a window where + // the conn is neither in p.conns nor counted as in-flight would let + // another caller overshoot MaxConns. + p.mu.Lock() + if p.closed { + p.mu.Unlock() + atomic.AddInt32(&p.inFlightConns, -1) + p.cond.Signal() + return nil, ErrPoolClosed + } + p.conns = append(p.conns, c) + p.mu.Unlock() + atomic.AddInt32(&p.inFlightConns, -1) + p.cond.Signal() + + // Re-acquire the just-registered connection. + continue + } + + // At capacity — wait for a release or handoff. cond.Wait releases the + // mutex while blocked, so releasers can append and signal. + p.cond.Wait() + p.mu.Unlock() + if err := ctx.Err(); err != nil { return nil, err } - return conn, nil } +} + +// dial simulates establishing a PostgreSQL connection, including the handshake +// window where the connection is "in flight". It honors context cancellation. +func (p *Pool) dial(ctx context.Context) (*conn, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + return &conn{id: int(atomic.AddInt32(&p.idCounter, 1))}, nil +} + +// release returns a connection to the idle list and wakes a waiter. +func (p *Pool) release(c *conn) { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return + } + p.conns = append(p.conns, c) + p.cond.Signal() p.mu.Unlock() +} - // Wait for existing connection or retry logic... - return p.waitForConn(ctx) -} \ No newline at end of file +// Close marks the pool closed. +func (p *Pool) Close() { + p.mu.Lock() + p.closed = true + p.cond.Broadcast() + p.mu.Unlock() +} From 1f644a7a2037d659384c321ce8d0467a0269859d Mon Sep 17 00:00:00 2001 From: jearthliu Date: Fri, 7 Aug 2026 10:05:37 +0800 Subject: [PATCH 2/2] fix: correct capacity accounting, ctx-cancelable wait, direct handoff --- pgxpool/pool.go | 121 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/pgxpool/pool.go b/pgxpool/pool.go index bc3dc67..0dd983d 100644 --- a/pgxpool/pool.go +++ b/pgxpool/pool.go @@ -25,23 +25,27 @@ func (c *Conn) Release() { } // pool is a connection pool that strictly respects MaxConns even during -// recovery from an outage: pending dial attempts count toward capacity, so -// a burst of blocked Acquire() calls cannot each start its own dial. +// recovery from an outage. Capacity is tracked as totalCreated (established +// connections, both borrowed and idle) + inFlightConns (pending dials); the +// sum never exceeds MaxConns, so a burst of blocked Acquire() calls waking +// together cannot each start its own dial. type Pool struct { - maxConns int32 - conns []*conn - inFlightConns int32 // pending dials — counted toward capacity - idCounter int32 - mu sync.Mutex - cond *sync.Cond - closed bool + maxConns int32 + conns []*conn // idle connections (mutex-protected) + totalCreated int32 // established connections ever created (borrowed + idle) + inFlightConns int32 // pending dials — counted toward capacity + idCounter int32 + mu sync.Mutex + waiters map[chan struct{}]struct{} + closed bool } // New creates a pool with the given maximum number of connections. func New(maxConns int32) *Pool { - p := &Pool{maxConns: maxConns} - p.cond = sync.NewCond(&p.mu) - return p + return &Pool{ + maxConns: maxConns, + waiters: make(map[chan struct{}]struct{}), + } } // ErrPoolClosed is returned by Acquire after Close. @@ -67,52 +71,65 @@ func (p *Pool) Acquire(ctx context.Context) (*Conn, error) { return &Conn{c: c, pool: p}, nil } - // Capacity check: established + pending must stay under MaxConns. - if int32(len(p.conns))+atomic.LoadInt32(&p.inFlightConns) < p.maxConns { + // Capacity check: established (borrowed + idle) + pending < MaxConns. + // totalCreated never decrements on release — a borrowed connection is + // still established and counts against capacity until Close. + if atomic.LoadInt32(&p.totalCreated)+atomic.LoadInt32(&p.inFlightConns) < p.maxConns { atomic.AddInt32(&p.inFlightConns, 1) p.mu.Unlock() // Dial outside the lock; the pending counter holds capacity. c, err := p.dial(ctx) if err != nil { - // Decrement on every failure path (dial timeout, ctx - // cancellation, auth failure) so the pool never starves. + // Decrement on every failure path so the pool never starves. atomic.AddInt32(&p.inFlightConns, -1) - p.cond.Signal() + p.wakeOne() return nil, err } - // Register the connection BEFORE releasing capacity: a window where - // the conn is neither in p.conns nor counted as in-flight would let - // another caller overshoot MaxConns. + // The connection is now established. Increment totalCreated under + // the lock and return it directly to the caller — no re-acquire + // dance, so a caller can never pick up someone else's connection. p.mu.Lock() if p.closed { p.mu.Unlock() atomic.AddInt32(&p.inFlightConns, -1) - p.cond.Signal() + p.wakeOne() return nil, ErrPoolClosed } - p.conns = append(p.conns, c) - p.mu.Unlock() + atomic.AddInt32(&p.totalCreated, 1) atomic.AddInt32(&p.inFlightConns, -1) - p.cond.Signal() - - // Re-acquire the just-registered connection. - continue + p.mu.Unlock() + p.wakeOne() + return &Conn{c: c, pool: p}, nil } - // At capacity — wait for a release or handoff. cond.Wait releases the - // mutex while blocked, so releasers can append and signal. - p.cond.Wait() + // At capacity — wait for a release, a failed dial, or context + // cancellation. Channel-based waiting (rather than sync.Cond) lets + // ctx cancellation interrupt the wait. + ch := make(chan struct{}) + p.waiters[ch] = struct{}{} p.mu.Unlock() - if err := ctx.Err(); err != nil { - return nil, err + + select { + case <-ctx.Done(): + // Remove the waiter; if a wake already fired the channel is closed + // and the map entry is gone — closing twice would panic, so guard + // with the mutex. + p.mu.Lock() + if _, ok := p.waiters[ch]; ok { + delete(p.waiters, ch) + } + p.mu.Unlock() + return nil, ctx.Err() + case <-ch: + // Woken — loop to re-check capacity. } } } -// dial simulates establishing a PostgreSQL connection, including the handshake -// window where the connection is "in flight". It honors context cancellation. +// dial simulates establishing a PostgreSQL connection. It honors context +// cancellation. func (p *Pool) dial(ctx context.Context) (*conn, error) { select { case <-ctx.Done(): @@ -122,7 +139,9 @@ func (p *Pool) dial(ctx context.Context) (*conn, error) { return &conn{id: int(atomic.AddInt32(&p.idCounter, 1))}, nil } -// release returns a connection to the idle list and wakes a waiter. +// release returns a connection to the idle list and wakes one waiter. +// totalCreated is intentionally NOT decremented — the connection is still +// established and counts against MaxConns. func (p *Pool) release(c *conn) { p.mu.Lock() if p.closed { @@ -130,14 +149,40 @@ func (p *Pool) release(c *conn) { return } p.conns = append(p.conns, c) - p.cond.Signal() + waiter := p.popWaiter() p.mu.Unlock() + if waiter != nil { + close(waiter) + } } -// Close marks the pool closed. +// wakeOne wakes a single waiter (after a failed dial frees a slot). +func (p *Pool) wakeOne() { + p.mu.Lock() + waiter := p.popWaiter() + p.mu.Unlock() + if waiter != nil { + close(waiter) + } +} + +// popWaiter removes and returns one waiting channel, if any. Caller must +// hold p.mu. +func (p *Pool) popWaiter() chan struct{} { + for ch := range p.waiters { + delete(p.waiters, ch) + return ch + } + return nil +} + +// Close marks the pool closed and wakes all waiters. func (p *Pool) Close() { p.mu.Lock() p.closed = true - p.cond.Broadcast() + for ch := range p.waiters { + close(ch) + } + p.waiters = make(map[chan struct{}]struct{}) p.mu.Unlock() }