[#3] Fix connection storm: prevent MaxConns violations during DB recovery

go.mod

@@ -0,0 +1,3 @@
+module github.com/lincai505011-ops/pgx
+
+go 1.21

pgxpool/pool.go

@@ -2,35 +2,200 @@ package pgxpool
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
+// ErrPoolClosed is returned when acquiring from a closed pool.
+var ErrPoolClosed = errors.New("pgxpool: pool is closed")
+
+// conn represents a database connection.
+type conn struct {
+	id       int64
+	created  time.Time
+	released bool
+}
+
+// Pool manages a fixed-size connection pool.
 type Pool struct {
-	// ... existing fields
 	maxConns      int32
+	totalConns    int32  // total connections created (in-pool + checked-out)
+	inFlightConns int32  // connections currently being established
 	conns         []*conn
-	inFlightConns int32
+	closed        int32  // atomic: 0=open, 1=closed
 	mu            sync.Mutex
-	// ...
+	cond          *sync.Cond
+
+	// For testing: allows injecting a slow dialer
+	dialFunc func(ctx context.Context) (*conn, error)
+}
+
+// Config holds pool configuration.
+type Config struct {
+	MaxConns int32
+}
+
+// New creates a new connection pool.
+func New(cfg *Config) *Pool {
+	if cfg.MaxConns <= 0 {
+		cfg.MaxConns = 4
+	}
+	p := &Pool{
+		maxConns:  cfg.MaxConns,
+		conns:     make([]*conn, 0),
+		dialFunc:  defaultDial,
+	}
+	p.cond = sync.NewCond(&p.mu)
+	return p
 }
 
+func defaultDial(ctx context.Context) (*conn, error) {
+	// Check for cancellation before dialing
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+	return &conn{
+		id:      time.Now().UnixNano(),
+		created: time.Now(),
+	}, nil
+}
+
+// Acquire gets a connection from the pool, creating a new one if possible.
+// It blocks until a connection is available or the context is cancelled.
 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()
+	if atomic.LoadInt32(&p.closed) == 1 {
+		return nil, ErrPoolClosed
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+		}
+
+		p.mu.Lock()
+
+		// Check if pool is closed while we were waiting
+		if atomic.LoadInt32(&p.closed) == 1 {
+			p.mu.Unlock()
+			return nil, ErrPoolClosed
+		}
 
-		conn, err := p.createNewConn(ctx)
-		atomic.AddInt32(&p.inFlightConns, -1)
-		if err != nil {
-			return nil, err
+		// Try to get an idle connection from the pool first
+		for i, c := range p.conns {
+			if !c.released {
+				// Remove from idle pool
+				p.conns = append(p.conns[:i], p.conns[i+1:]...)
+				p.mu.Unlock()
+				return &Conn{conn: c, pool: p}, nil
+			}
 		}
-		return conn, nil
+
+		// No idle connection available — can we create a new one?
+		total := atomic.LoadInt32(&p.totalConns)
+		inFlight := atomic.LoadInt32(&p.inFlightConns)
+
+		if total+inFlight < p.maxConns {
+			// Reserve a slot before releasing the lock
+			atomic.AddInt32(&p.inFlightConns, 1)
+			p.mu.Unlock()
+
+			// Establish the connection outside the lock
+			c, err := p.dialFunc(ctx)
+
+			// Always decrement in-flight counter
+			atomic.AddInt32(&p.inFlightConns, -1)
+
+			if err != nil {
+				// Connection failed — slot is freed (inFlightConns decremented)
+				// Wake up waiters so someone else can try
+				p.mu.Lock()
+				p.cond.Broadcast()
+				p.mu.Unlock()
+				return nil, fmt.Errorf("pgxpool: dial failed: %w", err)
+			}
+
+			// Connection succeeded — increment total
+			atomic.AddInt32(&p.totalConns, 1)
+
+			return &Conn{conn: c, pool: p}, nil
+		}
+
+		// Pool is at capacity — block until a connection is returned or context done
+		p.cond.Wait()
+		p.mu.Unlock()
 	}
+}
+
+// release returns a connection to the pool.
+func (p *Pool) release(c *conn) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	c.released = false
+	p.conns = append(p.conns, c)
+	// Signal one waiting goroutine
+	p.cond.Signal()
+}
+
+// removeConnection decrements totalConns when a connection is permanently removed.
+func (p *Pool) removeConnection() {
+	atomic.AddInt32(&p.totalConns, -1)
+}
+
+// Close shuts down the pool and releases all connections.
+func (p *Pool) Close() {
+	atomic.StoreInt32(&p.closed, 1)
+	p.mu.Lock()
+	p.conns = nil
+	atomic.StoreInt32(&p.totalConns, 0)
+	atomic.StoreInt32(&p.inFlightConns, 0)
+	p.cond.Broadcast()
 	p.mu.Unlock()
+}
+
+// Stats returns current pool statistics.
+func (p *Pool) Stats() Stats {
+	return Stats{
+		TotalConns:    atomic.LoadInt32(&p.totalConns),
+		InFlightConns: atomic.LoadInt32(&p.inFlightConns),
+		MaxConns:      p.maxConns,
+		IdleConns:     int32(len(p.conns)),
+	}
+}
 
-	// Wait for existing connection or retry logic...
-	return p.waitForConn(ctx)
-}
+// Stats holds pool statistics.
+type Stats struct {
+	TotalConns    int32
+	InFlightConns int32
+	MaxConns      int32
+	IdleConns     int32
+}
+
+// Conn wraps a connection with pool bookkeeping.
+type Conn struct {
+	conn   *conn
+	pool   *Pool
+	closed int32
+}
+
+// Release returns the connection to the pool.
+func (c *Conn) Release() {
+	if atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
+		c.pool.release(c.conn)
+	}
+}
+
+// Close permanently removes the connection.
+func (c *Conn) Close() error {
+	if atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
+		c.pool.removeConnection()
+	}
+	return nil
+}