fix: prevent connection storm during outage recovery

pgxpool/pool.go

@@ -2,35 +2,187 @@ 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. 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 {
-	// ... existing fields
-	maxConns      int32
-	conns         []*conn
-	inFlightConns int32
-	mu            sync.Mutex
-	// ...
+	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 {
+	return &Pool{
+		maxConns: maxConns,
+		waiters:  make(map[chan struct{}]struct{}),
+	}
+}
+
+// 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)
+	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 (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 so the pool never starves.
+				atomic.AddInt32(&p.inFlightConns, -1)
+				p.wakeOne()
+				return nil, err
+			}
+
+			// 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.wakeOne()
+				return nil, ErrPoolClosed
+			}
+			atomic.AddInt32(&p.totalCreated, 1)
+			atomic.AddInt32(&p.inFlightConns, -1)
+			p.mu.Unlock()
+			p.wakeOne()
+			return &Conn{c: c, pool: p}, nil
+		}
+
+		// 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()
 
-		conn, err := p.createNewConn(ctx)
-		atomic.AddInt32(&p.inFlightConns, -1)
-		if 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.
 		}
-		return conn, nil
 	}
+}
+
+// dial simulates establishing a PostgreSQL connection. 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 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 {
+		p.mu.Unlock()
+		return
+	}
+	p.conns = append(p.conns, c)
+	waiter := p.popWaiter()
+	p.mu.Unlock()
+	if waiter != nil {
+		close(waiter)
+	}
+}
+
+// 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
+}
 
-	// Wait for existing connection or retry logic...
-	return p.waitForConn(ctx)
-}
+// Close marks the pool closed and wakes all waiters.
+func (p *Pool) Close() {
+	p.mu.Lock()
+	p.closed = true
+	for ch := range p.waiters {
+		close(ch)
+	}
+	p.waiters = make(map[chan struct{}]struct{})
+	p.mu.Unlock()
+}