AIRouter/internal/scheduler/circuit_breaker.go

package scheduler

import (
	"sync"
	"time"
)

// CircuitBreaker implements a sliding-window circuit breaker for adapter health.
type CircuitBreaker struct {
	mu                 sync.Mutex
	errorRateThreshold float64
	minRequests        int
	windowSeconds      int
	openDuration       time.Duration
	halfOpenMax        int

	// sliding window state
	requests []time.Time
	errors   []time.Time

	// breaker state
	state         breakerState
	openedAt      time.Time
	halfOpenCount int
}

type breakerState int

const (
	breakerClosed breakerState = iota
	breakerOpen
	breakerHalfOpen
)

// NewCircuitBreaker creates a new circuit breaker from config.
func NewCircuitBreaker(errorRate float64, minRequests, windowSec, openSec, halfOpenMax int) *CircuitBreaker {
	return &CircuitBreaker{
		errorRateThreshold: errorRate,
		minRequests:        minRequests,
		windowSeconds:      windowSec,
		openDuration:       time.Duration(openSec) * time.Second,
		halfOpenMax:        halfOpenMax,
		state:              breakerClosed,
	}
}

// AllowRequest checks if a request should be allowed through.
func (cb *CircuitBreaker) AllowRequest() bool {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	now := time.Now()
	cb.prune(now)

	switch cb.state {
	case breakerClosed:
		return true
	case breakerOpen:
		if now.Sub(cb.openedAt) >= cb.openDuration {
			cb.state = breakerHalfOpen
			cb.halfOpenCount = 0
			return true
		}
		return false
	case breakerHalfOpen:
		if cb.halfOpenCount < cb.halfOpenMax {
			cb.halfOpenCount++
			return true
		}
		return false
	}
	return true
}

// RecordSuccess records a successful request.
func (cb *CircuitBreaker) RecordSuccess() {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	now := time.Now()
	cb.requests = append(cb.requests, now)

	if cb.state == breakerHalfOpen {
		cb.state = breakerClosed
		cb.requests = nil
		cb.errors = nil
	}
}

// RecordError records a failed request and may trip the breaker.
func (cb *CircuitBreaker) RecordError() {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	now := time.Now()
	cb.requests = append(cb.requests, now)
	cb.errors = append(cb.errors, now)

	if cb.state == breakerHalfOpen {
		cb.state = breakerOpen
		cb.openedAt = now
		return
	}

	if cb.state == breakerClosed && len(cb.requests) >= cb.minRequests {
		errorRate := float64(len(cb.errors)) / float64(len(cb.requests))
		if errorRate >= cb.errorRateThreshold {
			cb.state = breakerOpen
			cb.openedAt = now
		}
	}
}

// State returns the current breaker state name.
func (cb *CircuitBreaker) State() string {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	switch cb.state {
	case breakerClosed:
		return "closed"
	case breakerOpen:
		return "open"
	case breakerHalfOpen:
		return "half_open"
	}
	return "unknown"
}

// BreakerStats 熔断器统计信息。
type BreakerStats struct {
	State         string  `json:"state"`
	TotalRequests int     `json:"total_requests"`
	TotalErrors   int     `json:"total_errors"`
	ErrorRate     float64 `json:"error_rate"`
	WindowSeconds int     `json:"window_seconds"`
	OpenDuration  string  `json:"open_duration"`
	Threshold     float64 `json:"error_rate_threshold"`
	MinRequests   int     `json:"min_requests"`
}

// Stats 返回熔断器的详细统计信息。
func (cb *CircuitBreaker) Stats() BreakerStats {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	now := time.Now()
	cb.prune(now)

	total := len(cb.requests)
	errors := len(cb.errors)
	var errorRate float64
	if total > 0 {
		errorRate = float64(errors) / float64(total)
	}

	stateName := "closed"
	switch cb.state {
	case breakerOpen:
		stateName = "open"
	case breakerHalfOpen:
		stateName = "half_open"
	}

	return BreakerStats{
		State:         stateName,
		TotalRequests: total,
		TotalErrors:   errors,
		ErrorRate:     errorRate,
		WindowSeconds: cb.windowSeconds,
		OpenDuration:  cb.openDuration.String(),
		Threshold:     cb.errorRateThreshold,
		MinRequests:   cb.minRequests,
	}
}

// prune removes entries outside the sliding window.
func (cb *CircuitBreaker) prune(now time.Time) {
	cutoff := now.Add(-time.Duration(cb.windowSeconds) * time.Second)
	cb.requests = pruneBefore(cb.requests, cutoff)
	cb.errors = pruneBefore(cb.errors, cutoff)
}

func pruneBefore(times []time.Time, cutoff time.Time) []time.Time {
	idx := 0
	for idx < len(times) && times[idx].Before(cutoff) {
		idx++
	}
	if idx > 0 {
		times = times[idx:]
	}
	return times
}