fix: prevent connection storms and MaxConns violations in pgxpool (Fixes #3)

go.mod

@@ -0,0 +1,3 @@
+module github.com/laurentketterle-hub/pgx-1
+
+go 1.26.5

main.go

@@ -1,7 +1,17 @@
 package main
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/laurentketterle-hub/pgx-1/pgxpool"
+)
 
 func main() {
-	fmt.Println("Hello, Bounty Hunter!")
+	// Create a pool with MaxConns=5 and a simple dialer for demonstration.
+	pool := pgxpool.NewPool(5, pgxpool.SimpleDialFunc())
+	defer pool.Close()
+
+	stat := pool.Stat()
+	fmt.Printf("Pool: MaxConns=%d Active=%d Idle=%d InFlight=%d\n",
+		stat.MaxConns, stat.ActiveConns, stat.IdleConns, stat.InFlightConns)
 }

pgxpool/pool.go

@@ -2,35 +2,244 @@ package pgxpool
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"sync"
 	"sync/atomic"
 )
 
+// Conn represents a connection acquired from the pool.
+// Release() must be called to return the connection to the pool.
+type Conn struct {
+	pool *Pool
+	conn *conn
+}
+
+// Release returns the connection to the pool.
+func (c *Conn) Release() {
+	if c.pool == nil || c.conn == nil {
+		return
+	}
+	c.pool.release(c.conn)
+	c.pool = nil
+	c.conn = nil
+}
+
+// conn is an internal tracked connection.
+type conn struct {
+	id int32
+}
+
+// DialFunc is the signature for a function that creates a new connection.
+type DialFunc func(ctx context.Context) (*conn, error)
+
+// Pool is a concurrency-safe bounded connection pool.
+//
+// It tracks both active (acquired) connections and pending (in-flight)
+// connection attempts to prevent connection storms. The invariant
+//
+//	activeCount + inFlightConns + len(idleConns) <= maxConns
+//
+// is maintained under the pool mutex, ensuring MaxConns is never
+// breached even during database recovery scenarios.
 type Pool struct {
-	// ... existing fields
 	maxConns      int32
-	conns         []*conn
-	inFlightConns int32
+	idleConns     []*conn // idle connections ready for reuse
+	activeCount   int32   // currently acquired (checked-out) connections
+	inFlightConns int32   // connection attempts in progress (dial/handshake)
 	mu            sync.Mutex
-	// ...
+	cond          *sync.Cond // signaled when a connection is released or a dial fails
+	closed        bool
+	dialFunc      DialFunc
+	nextID        int32
+}
+
+// NewPool creates a new Pool that limits total connections to maxConns.
+// dialFunc is called to establish new physical connections.
+func NewPool(maxConns int32, dialFunc DialFunc) *Pool {
+	p := &Pool{
+		maxConns:  maxConns,
+		dialFunc:  dialFunc,
+		idleConns: make([]*conn, 0, maxConns),
+	}
+	p.cond = sync.NewCond(&p.mu)
+	return p
 }
 
+// Acquire obtains a connection from the pool, blocking until one is
+// available or ctx is cancelled.
+//
+// Connection storm prevention:
+//   - Before dialing, Acquire atomically checks that
+//     activeCount + inFlightConns < maxConns and increments
+//     inFlightConns — all under the pool mutex.
+//   - If the limit is reached, the caller blocks on a condition
+//     variable until a connection is released or a pending dial
+//     completes (success or failure).
+//   - On dial failure, inFlightConns is decremented and waiters
+//     are woken so another goroutine can attempt a connection.
+//   - On dial success, the connection is added to the pool's
+//     accounting before inFlightConns is decremented, preserving
+//     the invariant.
 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()
-
-		conn, err := p.createNewConn(ctx)
-		atomic.AddInt32(&p.inFlightConns, -1)
-		if err != nil {
-			return nil, err
+	defer p.mu.Unlock()
+
+	for {
+		if p.closed {
+			return nil, errors.New("pgxpool: pool is closed")
+		}
+
+		// Check context before blocking
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+		}
+
+		// 1. Try to return an idle connection immediately.
+		if len(p.idleConns) > 0 {
+			c := p.idleConns[len(p.idleConns)-1]
+			p.idleConns = p.idleConns[:len(p.idleConns)-1]
+			p.activeCount++
+			return &Conn{pool: p, conn: c}, nil
 		}
-		return conn, nil
+
+		// 2. Check if we're allowed to start a new dial.
+		//    activeCount + inFlightConns < maxConns means there is
+		//    at least one free slot (neither idle nor pending).
+		totalInUse := p.activeCount + atomic.LoadInt32(&p.inFlightConns)
+		if totalInUse < p.maxConns {
+			// Reserve a pending slot before releasing the lock.
+			atomic.AddInt32(&p.inFlightConns, 1)
+			p.mu.Unlock()
+
+			// Perform the slow I/O outside the lock.
+			c, dialErr := p.dialFunc(ctx)
+
+			// Reacquire the lock to atomically update pool state.
+			p.mu.Lock()
+
+			if dialErr != nil {
+				// Dial failed: release the reserved slot and wake
+				// waiters so another caller can try.
+				atomic.AddInt32(&p.inFlightConns, -1)
+				p.cond.Broadcast()
+				return nil, fmt.Errorf("pgxpool: dial failed: %w", dialErr)
+			}
+
+			// Dial succeeded: add to active set and release the
+			// pending reservation — both under the lock so the
+			// invariant
+			//   activeCount + inFlightConns + len(idle) <= maxConns
+			// is never violated.
+			p.activeCount++
+			atomic.AddInt32(&p.inFlightConns, -1)
+
+			// Wake any waiters — another slot may have opened up.
+			p.cond.Broadcast()
+			return &Conn{pool: p, conn: c}, nil
+		}
+
+		// 3. Pool is fully saturated (active + inFlight >= maxConns).
+		//    Wait for a connection release or a dial completion, but
+		//    respect context cancellation by waking when ctx is done.
+		done := p.waitWithContext(ctx)
+		if done {
+			// Context was cancelled while waiting; the deferred unlock
+			// will run, and the loop will catch ctx.Done() on the next
+			// iteration.
+			continue
+		}
+		// p.cond.Wait() was woken by a signal; loop again to
+		// re-check conditions.
 	}
-	p.mu.Unlock()
+}
 
-	// Wait for existing connection or retry logic...
-	return p.waitForConn(ctx)
-}
+// waitWithContext waits on the pool's condition variable, but also
+// respects ctx cancellation. It returns true if the context was
+// cancelled (caller should re-check and return the error), false
+// if the cond was signaled normally.
+//
+// Must be called with p.mu held.
+func (p *Pool) waitWithContext(ctx context.Context) (cancelled bool) {
+	// Use a channel to detect context cancellation.
+	// We can't select on cond.Wait, so we spawn a goroutine
+	// that broadcasts the cond when ctx is done.
+	done := make(chan struct{})
+	defer close(done)
+
+	go func() {
+		select {
+		case <-ctx.Done():
+			// Context cancelled: wake up the cond waiter.
+			p.mu.Lock()
+			p.cond.Broadcast()
+			p.mu.Unlock()
+		case <-done:
+			// Normal exit: context wasn't cancelled before
+			// cond.Wait returned.
+		}
+	}()
+
+	p.cond.Wait()
+
+	// Check if we woke up because of context cancellation.
+	select {
+	case <-ctx.Done():
+		return true
+	default:
+		return false
+	}
+}
+
+// release returns a connection to the pool's idle set and wakes a waiter.
+func (p *Pool) release(c *conn) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	p.activeCount--
+	p.idleConns = append(p.idleConns, c)
+	p.cond.Signal()
+}
+
+// Close shuts down the pool. Subsequent Acquire calls return an error.
+func (p *Pool) Close() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.closed {
+		return
+	}
+	p.closed = true
+	p.idleConns = nil
+	p.cond.Broadcast()
+}
+
+// Stat returns a snapshot of the pool's current state.
+type Stat struct {
+	MaxConns      int32
+	ActiveConns   int32
+	IdleConns     int32
+	InFlightConns int32
+}
+
+// Stat returns current pool statistics.
+func (p *Pool) Stat() Stat {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	return Stat{
+		MaxConns:      p.maxConns,
+		ActiveConns:   p.activeCount,
+		IdleConns:     int32(len(p.idleConns)),
+		InFlightConns: atomic.LoadInt32(&p.inFlightConns),
+	}
+}
+
+// Len returns the number of idle connections in the pool.
+func (p *Pool) Len() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.idleConns)
+}