pool.go

package pgpool

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"
)

// DatabasePool manages a pool of PostgreSQL connections for a single database
type DatabasePool struct {
	// Pool name for identification
	name string

	// Primary connection pool
	primary *pgxpool.Pool

	// Replica pools for read scaling
	replicas []*pgxpool.Pool

	// Configuration
	config *BackendConfig

	// Logger
	log *zap.Logger

	// Statistics
	stats *PoolStats

	// Lifecycle management
	mu     sync.RWMutex
	closed atomic.Bool

	// Circuit breaker for backend health
	circuitBreaker *CircuitBreaker
}

// PoolStats tracks pool usage metrics
type PoolStats struct {
	// Connection statistics
	TotalConnections     int32
	IdleConnections      int32
	ActiveConnections    int32
	WaitingClients       int32
	TotalAcquired        int64
	TotalReleased        int64
	AcquireCount         int64
	AcquireDuration      time.Duration
	CanceledAcquireCount int64
	ConstructingConns    int32
	MaxConnections       int32

	// Query statistics
	QueryCount     int64
	QueryDuration  time.Duration
	ErrorCount     int64
	SlowQueryCount int64

	// Pool health
	HealthChecksPassed int64
	HealthChecksFailed int64
	LastHealthCheck    time.Time
	LastError          error

	mu sync.RWMutex
}

// NewDatabasePool creates a new connection pool for the specified database
func NewDatabasePool(
	name string,
	config *BackendConfig,
	log *zap.Logger,
) (*DatabasePool, error) {
	pool := &DatabasePool{
		name:   name,
		config: config,
		log:    log,
		stats:  &PoolStats{},
	}

	// Initialize circuit breaker
	pool.circuitBreaker = NewCircuitBreaker(
		5,              // max failures
		30*time.Second, // timeout
		log.Named("circuit_breaker"),
	)

	// Create primary connection pool
	if err := pool.initPrimaryPool(); err != nil {
		return nil, fmt.Errorf("failed to initialize primary pool: %w", err)
	}

	log.Info("database pool initialized",
		zap.String("name", name),
		zap.Int("max_connections", config.Pool.MaxConnections),
		zap.String("mode", config.Behavior.Mode))

	return pool, nil
}

// initPrimaryPool creates the pgxpool for primary database
func (p *DatabasePool) initPrimaryPool() error {
	// Parse DSN
	dsn := p.config.BuildDSN()

	// Build pool configuration
	poolConfig, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return fmt.Errorf("invalid DSN: %w", err)
	}

	// Apply pool settings
	poolConfig.MaxConns = int32(p.config.Pool.MaxConnections)
	poolConfig.MinConns = int32(p.config.Pool.MinConnections)
	poolConfig.MaxConnLifetime = p.config.Pool.MaxLifetime
	poolConfig.MaxConnIdleTime = p.config.Pool.MaxIdleTime
	poolConfig.HealthCheckPeriod = p.config.Pool.HealthCheckPeriod

	// Configure connection behavior
	poolConfig.BeforeAcquire = p.beforeAcquire
	poolConfig.AfterRelease = p.afterRelease
	poolConfig.BeforeClose = p.beforeClose

	// Set query timeout if configured
	if p.config.Behavior.QueryTimeout > 0 {
		poolConfig.ConnConfig.ConnectTimeout = p.config.Behavior.QueryTimeout
	}

	// Create pool with context
	ctx := context.Background()
	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
	if err != nil {
		return fmt.Errorf("failed to create pool: %w", err)
	}

	// Lazy connect or test connection immediately
	if !p.config.Behavior.LazyConnect {
		if err := pool.Ping(ctx); err != nil {
			pool.Close()
			return fmt.Errorf("failed to ping database: %w", err)
		}
		p.log.Info("database connection verified")
	}

	p.primary = pool
	return nil
}

// beforeAcquire is called before a connection is acquired from the pool
func (p *DatabasePool) beforeAcquire(ctx context.Context, conn *pgx.Conn) bool {
	// Check circuit breaker state
	if !p.circuitBreaker.Allow() {
		p.log.Warn("circuit breaker is open, rejecting connection request")
		return false
	}

	// Update statistics
	atomic.AddInt64(&p.stats.AcquireCount, 1)
	return true
}

// afterRelease is called after a connection is released back to the pool
func (p *DatabasePool) afterRelease(conn *pgx.Conn) bool {
	// Update statistics
	atomic.AddInt64(&p.stats.TotalReleased, 1)

	// Reset connection state if configured
	if p.config.Behavior.ResetOnReturn {
		// Execute DISCARD ALL to reset session state
		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
		defer cancel()

		if _, err := conn.Exec(ctx, "DISCARD ALL"); err != nil {
			p.log.Warn("failed to reset connection state", zap.Error(err))
			return false // Discard this connection
		}
	}

	return true
}

// beforeClose is called before a connection is closed
func (p *DatabasePool) beforeClose(conn *pgx.Conn) {
	// Connection is being removed from pool
	p.log.Debug("connection closed")
}

// Acquire obtains a connection from the pool
// The connection must be released back to the pool when done
func (p *DatabasePool) Acquire(ctx context.Context) (*pgxpool.Conn, error) {
	if p.closed.Load() {
		return nil, fmt.Errorf("pool is closed")
	}

	// Apply acquisition timeout from config
	if p.config.Pool.AcquireTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, p.config.Pool.AcquireTimeout)
		defer cancel()
	}

	start := time.Now()
	conn, err := p.primary.Acquire(ctx)
	duration := time.Since(start)

	// Update statistics
	p.stats.mu.Lock()
	p.stats.AcquireDuration += duration
	atomic.AddInt64(&p.stats.TotalAcquired, 1)
	p.stats.mu.Unlock()

	if err != nil {
		atomic.AddInt64(&p.stats.ErrorCount, 1)
		p.circuitBreaker.RecordFailure()
		return nil, fmt.Errorf("failed to acquire connection: %w", err)
	}

	p.circuitBreaker.RecordSuccess()
	return conn, nil
}

// Query executes a query and returns the result rows
func (p *DatabasePool) Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error) {
	if p.closed.Load() {
		return nil, fmt.Errorf("pool is closed")
	}

	start := time.Now()
	rows, err := p.primary.Query(ctx, sql, args...)
	duration := time.Since(start)

	// Update statistics
	atomic.AddInt64(&p.stats.QueryCount, 1)
	p.stats.mu.Lock()
	p.stats.QueryDuration += duration
	p.stats.mu.Unlock()

	// Log slow queries
	if p.config.Behavior.QueryTimeout > 0 && duration > p.config.Behavior.QueryTimeout {
		atomic.AddInt64(&p.stats.SlowQueryCount, 1)
		p.log.Warn("slow query detected",
			zap.Duration("duration", duration),
			zap.String("query", sql))
	}

	if err != nil {
		atomic.AddInt64(&p.stats.ErrorCount, 1)
		p.circuitBreaker.RecordFailure()
		return nil, err
	}

	p.circuitBreaker.RecordSuccess()
	return rows, nil
}

// QueryRow executes a query that returns at most one row
func (p *DatabasePool) QueryRow(ctx context.Context, sql string, args ...interface{}) pgx.Row {
	if p.closed.Load() {
		// Return an error row
		return &errorRow{err: fmt.Errorf("pool is closed")}
	}

	start := time.Now()
	row := p.primary.QueryRow(ctx, sql, args...)
	duration := time.Since(start)

	// Update statistics
	atomic.AddInt64(&p.stats.QueryCount, 1)
	p.stats.mu.Lock()
	p.stats.QueryDuration += duration
	p.stats.mu.Unlock()

	// Log slow queries
	if p.config.Behavior.QueryTimeout > 0 && duration > p.config.Behavior.QueryTimeout {
		atomic.AddInt64(&p.stats.SlowQueryCount, 1)
		p.log.Warn("slow query detected",
			zap.Duration("duration", duration),
			zap.String("query", sql))
	}

	return row
}

// Exec executes a command (INSERT, UPDATE, DELETE, etc.)
func (p *DatabasePool) Exec(ctx context.Context, sql string, args ...interface{}) (pgconn.CommandTag, error) {
	if p.closed.Load() {
		return pgconn.CommandTag{}, fmt.Errorf("pool is closed")
	}

	start := time.Now()
	tag, err := p.primary.Exec(ctx, sql, args...)
	duration := time.Since(start)

	// Update statistics
	atomic.AddInt64(&p.stats.QueryCount, 1)
	p.stats.mu.Lock()
	p.stats.QueryDuration += duration
	p.stats.mu.Unlock()

	if err != nil {
		atomic.AddInt64(&p.stats.ErrorCount, 1)
		p.circuitBreaker.RecordFailure()
		return tag, err
	}

	p.circuitBreaker.RecordSuccess()
	return tag, nil
}

// Begin starts a new transaction
func (p *DatabasePool) Begin(ctx context.Context) (pgx.Tx, error) {
	if p.closed.Load() {
		return nil, fmt.Errorf("pool is closed")
	}

	tx, err := p.primary.Begin(ctx)
	if err != nil {
		atomic.AddInt64(&p.stats.ErrorCount, 1)
		p.circuitBreaker.RecordFailure()
		return nil, err
	}

	p.circuitBreaker.RecordSuccess()
	return tx, nil
}

// HealthCheck verifies pool connectivity and health
func (p *DatabasePool) HealthCheck(ctx context.Context) error {
	if p.closed.Load() {
		return fmt.Errorf("pool is closed")
	}

	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	// Ping the database
	err := p.primary.Ping(ctx)

	// Update health check statistics
	p.stats.mu.Lock()
	p.stats.LastHealthCheck = time.Now()
	if err != nil {
		p.stats.LastError = err
		atomic.AddInt64(&p.stats.HealthChecksFailed, 1)
		p.circuitBreaker.RecordFailure()
	} else {
		atomic.AddInt64(&p.stats.HealthChecksPassed, 1)
		p.circuitBreaker.RecordSuccess()
	}
	p.stats.mu.Unlock()

	return err
}

// Stats returns current pool statistics
func (p *DatabasePool) Stats() *PoolStats {
	if p.primary == nil {
		return p.stats
	}

	// Get pgxpool stats
	pgxStats := p.primary.Stat()

	// Copy statistics atomically
	stats := &PoolStats{
		TotalConnections:     pgxStats.TotalConns(),
		IdleConnections:      pgxStats.IdleConns(),
		ActiveConnections:    pgxStats.AcquiredConns(),
		ConstructingConns:    pgxStats.ConstructingConns(),
		MaxConnections:       pgxStats.MaxConns(),
		TotalAcquired:        atomic.LoadInt64(&p.stats.TotalAcquired),
		TotalReleased:        atomic.LoadInt64(&p.stats.TotalReleased),
		AcquireCount:         atomic.LoadInt64(&p.stats.AcquireCount),
		CanceledAcquireCount: atomic.LoadInt64(&p.stats.CanceledAcquireCount),
		QueryCount:           atomic.LoadInt64(&p.stats.QueryCount),
		ErrorCount:           atomic.LoadInt64(&p.stats.ErrorCount),
		SlowQueryCount:       atomic.LoadInt64(&p.stats.SlowQueryCount),
		HealthChecksPassed:   atomic.LoadInt64(&p.stats.HealthChecksPassed),
		HealthChecksFailed:   atomic.LoadInt64(&p.stats.HealthChecksFailed),
	}

	p.stats.mu.RLock()
	stats.AcquireDuration = p.stats.AcquireDuration
	stats.QueryDuration = p.stats.QueryDuration
	stats.LastHealthCheck = p.stats.LastHealthCheck
	stats.LastError = p.stats.LastError
	p.stats.mu.RUnlock()

	return stats
}

// Close gracefully shuts down the pool
func (p *DatabasePool) Close(ctx context.Context) error {
	if !p.closed.CompareAndSwap(false, true) {
		return fmt.Errorf("pool already closed")
	}

	p.log.Info("closing database pool", zap.String("name", p.name))

	// Close primary pool
	if p.primary != nil {
		p.primary.Close()
	}

	// Close replica pools
	for _, replica := range p.replicas {
		if replica != nil {
			replica.Close()
		}
	}

	p.log.Info("database pool closed successfully", zap.String("name", p.name))
	return nil
}

// errorRow is a pgx.Row that always returns an error
type errorRow struct {
	err error
}

func (e *errorRow) Scan(dest ...interface{}) error {
	return e.err
}