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

go.mod

@@ -0,0 +1,3 @@
+module github.com/colmev080/pgx
+
+go 1.18

pgxpool/pool.go

@@ -2,35 +2,136 @@ package pgxpool
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
+var (
+	ErrMaxConnsReached = errors.New("max connections reached")
+	ErrConnClosed      = errors.New("connection closed")
+)
+
+type Conn struct {
+	id int
+}
+
 type Pool struct {
-	// ... existing fields
 	maxConns      int32
-	conns         []*conn
+	conns         []*Conn
 	inFlightConns int32
 	mu            sync.Mutex
-	// ...
+	cond          *sync.Cond
+	dialFunc      func(ctx context.Context) (*Conn, error)
+	connCounter   int32
+}
+
+func NewPool(maxConns int32, dialFunc func(ctx context.Context) (*Conn, error)) *Pool {
+	p := &Pool{
+		maxConns: maxConns,
+		conns:    make([]*Conn, 0, maxConns),
+		dialFunc: dialFunc,
+	}
+	p.cond = sync.NewCond(&p.mu)
+	return p
 }
 
 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 {
+
+	for {
+		// 1. Tentar pegar conexão já estabelecida e ociosa no pool
+		if len(p.conns) > 0 {
+			conn := p.conns[len(p.conns)-1]
+			p.conns = p.conns[:len(p.conns)-1]
+			p.mu.Unlock()
+			return conn, nil
+		}
+
+		// 2. Se a soma de (estabelecidas + em progresso/in-flight) for menor que MaxConns, abre nova vaga
+		currentInFlight := atomic.LoadInt32(&p.inFlightConns)
+		if int32(len(p.conns))+currentInFlight < p.maxConns {
+			atomic.AddInt32(&p.inFlightConns, 1)
+			p.mu.Unlock()
+
+			// Dispara discagem física (I/O) fora da trava principal
+			conn, err := p.createNewConn(ctx)
+			atomic.AddInt32(&p.inFlightConns, -1)
+
+			p.mu.Lock()
+			// Acorda goroutines aguardando por alteração de estado no pool
+			p.cond.Broadcast()
+
+			if err != nil {
+				p.mu.Unlock()
+				return nil, err
+			}
+
+			p.mu.Unlock()
+			return conn, nil
+		}
+
+		// 3. Se atingiu o limite de capacidade (estabelecidas + in-flight >= MaxConns), aguarda
+		select {
+		case <-ctx.Done():
+			p.mu.Unlock()
+			return nil, ctx.Err()
+		default:
+		}
+
+		// Aguarda sinal de liberação ou finalização de conexão física
+		done := make(chan struct{})
+		go func() {
+			select {
+			case <-ctx.Done():
+				p.mu.Lock()
+				p.cond.Broadcast()
+				p.mu.Unlock()
+			case <-done:
+			}
+		}()
+
+		p.cond.Wait()
+		close(done)
+
+		if err := ctx.Err(); err != nil {
+			p.mu.Unlock()
 			return nil, err
 		}
-		return conn, nil
 	}
+}
+
+func (p *Pool) Release(conn *Conn) {
+	if conn == nil {
+		return
+	}
+	p.mu.Lock()
+	p.conns = append(p.conns, conn)
+	p.cond.Signal()
 	p.mu.Unlock()
+}
+
+func (p *Pool) createNewConn(ctx context.Context) (*Conn, error) {
+	if p.dialFunc != nil {
+		return p.dialFunc(ctx)
+	}
+	// Fallback/Default Dial
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-time.After(10 * time.Millisecond):
+		id := atomic.AddInt32(&p.connCounter, 1)
+		return &Conn{id: int(id)}, nil
+	}
+}
+
+func (p *Pool) TotalEstablished() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.conns)
+}
 
-	// Wait for existing connection or retry logic...
-	return p.waitForConn(ctx)
+func (p *Pool) InFlightCount() int32 {
+	return atomic.LoadInt32(&p.inFlightConns)
 }

pgxpool/pool_test.go

@@ -0,0 +1,109 @@
+package pgxpool_test
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/colmev080/pgx/pgxpool"
+)
+
+// TestMaxConnsNeverExceededUnderRecovery verifies that during a DB recovery
+// with high contention (100 concurrent goroutines), the sum of active/established
+// connections and in-flight dials NEVER exceeds MaxConns.
+func TestMaxConnsNeverExceededUnderRecovery(t *testing.T) {
+	const maxConns int32 = 5
+	const numGoroutines = 100
+
+	var currentDials int32
+	var maxObservedDials int32
+
+	dialFunc := func(ctx context.Context) (*pgxpool.Conn, error) {
+		active := atomic.AddInt32(&currentDials, 1)
+		defer atomic.AddInt32(&currentDials, -1)
+
+		for {
+			max := atomic.LoadInt32(&maxObservedDials)
+			if active > max {
+				if atomic.CompareAndSwapInt32(&maxObservedDials, max, active) {
+					break
+				}
+			} else {
+				break
+			}
+		}
+
+		// Simula um atraso na conexão física (ex: handshake / I/O de rede)
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(20 * time.Millisecond):
+		}
+
+		return &pgxpool.Conn{}, nil
+	}
+
+	pool := pgxpool.NewPool(maxConns, dialFunc)
+
+	var wg sync.WaitGroup
+	errCh := make(chan error, numGoroutines)
+
+	for i := 0; i < numGoroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+			defer cancel()
+
+			conn, err := pool.Acquire(ctx)
+			if err != nil {
+				errCh <- err
+				return
+			}
+			// Retém a conexão brevemente e libera
+			time.Sleep(10 * time.Millisecond)
+			pool.Release(conn)
+		}()
+	}
+
+	wg.Wait()
+	close(errCh)
+
+	for err := range errCh {
+		t.Fatalf("Erro durante Acquire em concorrência: %v", err)
+	}
+
+	if maxObservedDials > maxConns {
+		t.Fatalf("VIOLAÇÃO DE MAXCONNS! Máximo de conexões simultâneas discando foi %d, limite era %d", maxObservedDials, maxConns)
+	}
+
+	t.Logf("✅ SUCESSO: 100 Goroutines executadas. Máximo de conexões simultâneas físicas registradas foi %d (limite = %d)", maxObservedDials, maxConns)
+}
+
+// TestInFlightCounterDecrementedOnError verifies that if connection creation fails,
+// the inFlightConns counter is decremented correctly to avoid pool starvation.
+func TestInFlightCounterDecrementedOnError(t *testing.T) {
+	const maxConns int32 = 2
+
+	dialErr := errors.New("network outage")
+	dialFunc := func(ctx context.Context) (*pgxpool.Conn, error) {
+		return nil, dialErr
+	}
+
+	pool := pgxpool.NewPool(maxConns, dialFunc)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	_, err := pool.Acquire(ctx)
+	if !errors.Is(err, dialErr) {
+		t.Fatalf("Esperava erro %v, veio: %v", dialErr, err)
+	}
+
+	if inFlight := pool.InFlightCount(); inFlight != 0 {
+		t.Fatalf("Esperava inFlightCount = 0 após erro, mas veio: %d", inFlight)
+	}
+}