circuit_breaker.go

package pgpool

import (
	"sync/atomic"
	"time"

	"go.uber.org/zap"
)

// CircuitBreakerState represents the state of the circuit breaker
type CircuitBreakerState int32

const (
	// StateClosed means circuit is operating normally
	StateClosed CircuitBreakerState = iota
	// StateOpen means circuit is open due to failures
	StateOpen
	// StateHalfOpen means circuit is testing if backend recovered
	StateHalfOpen
)

// CircuitBreaker implements the circuit breaker pattern for backend health
type CircuitBreaker struct {
	// Configuration
	maxFailures int           // Consecutive failures to open circuit
	timeout     time.Duration // Time before attempting reset

	// State
	state        atomic.Int32 // CircuitBreakerState
	failures     atomic.Int32 // Current consecutive failure count
	lastFailTime atomic.Int64 // Unix timestamp of last failure

	// Logger
	log *zap.Logger
}

// NewCircuitBreaker creates a new circuit breaker
func NewCircuitBreaker(maxFailures int, timeout time.Duration, log *zap.Logger) *CircuitBreaker {
	cb := &CircuitBreaker{
		maxFailures: maxFailures,
		timeout:     timeout,
		log:         log,
	}
	cb.state.Store(int32(StateClosed))
	return cb
}

// Allow checks if operation should be allowed
func (cb *CircuitBreaker) Allow() bool {
	state := CircuitBreakerState(cb.state.Load())

	switch state {
	case StateClosed:
		return true

	case StateOpen:
		// Check if timeout has elapsed
		lastFail := time.Unix(0, cb.lastFailTime.Load())
		if time.Since(lastFail) >= cb.timeout {
			// Try half-open state
			if cb.state.CompareAndSwap(int32(StateOpen), int32(StateHalfOpen)) {
				cb.log.Info("circuit breaker entering half-open state")
				return true
			}
		}
		return false

	case StateHalfOpen:
		// Allow one request through to test
		return true

	default:
		return false
	}
}

// RecordSuccess records a successful operation
func (cb *CircuitBreaker) RecordSuccess() {
	state := CircuitBreakerState(cb.state.Load())

	switch state {
	case StateHalfOpen:
		// Successful test, close the circuit
		if cb.state.CompareAndSwap(int32(StateHalfOpen), int32(StateClosed)) {
			cb.failures.Store(0)
			cb.log.Info("circuit breaker closed after successful test")
		}

	case StateClosed:
		// Reset failure count on success
		cb.failures.Store(0)
	}
}

// RecordFailure records a failed operation
func (cb *CircuitBreaker) RecordFailure() {
	cb.lastFailTime.Store(time.Now().UnixNano())
	failures := cb.failures.Add(1)

	state := CircuitBreakerState(cb.state.Load())

	switch state {
	case StateClosed:
		if failures >= int32(cb.maxFailures) {
			if cb.state.CompareAndSwap(int32(StateClosed), int32(StateOpen)) {
				cb.log.Warn("circuit breaker opened due to failures",
					zap.Int32("failures", failures),
					zap.Int("max_failures", cb.maxFailures))
			}
		}

	case StateHalfOpen:
		// Failed test, open circuit again
		if cb.state.CompareAndSwap(int32(StateHalfOpen), int32(StateOpen)) {
			cb.log.Warn("circuit breaker reopened after failed test")
		}
	}
}

// GetState returns current circuit breaker state
func (cb *CircuitBreaker) GetState() CircuitBreakerState {
	return CircuitBreakerState(cb.state.Load())
}

// Reset manually resets the circuit breaker to closed state
func (cb *CircuitBreaker) Reset() {
	cb.state.Store(int32(StateClosed))
	cb.failures.Store(0)
	cb.log.Info("circuit breaker manually reset")
}