ai-gateway-go/internal/gateway/resilience.go

package gateway

import (
	"context"
	"errors"
	"io"
	"net/http"
	"sync"
	"time"
)

var ErrCircuitOpen = errors.New("provider circuit is open")

type ResiliencePolicy struct {
	ResponseHeaderTimeout time.Duration
	MaxRetries            int
	RetryBackoff          time.Duration
	CircuitThreshold      int
	CircuitOpenDuration   time.Duration
}

func DefaultResiliencePolicy() ResiliencePolicy {
	return ResiliencePolicy{
		ResponseHeaderTimeout: 60 * time.Second, MaxRetries: 2, RetryBackoff: 50 * time.Millisecond,
		CircuitThreshold: 5, CircuitOpenDuration: 30 * time.Second,
	}
}

type circuitBreaker struct {
	mu          sync.Mutex
	failures    int
	threshold   int
	openFor     time.Duration
	openUntil   time.Time
	halfOpenRun bool
}

func newCircuitBreaker(policy ...ResiliencePolicy) *circuitBreaker {
	settings := DefaultResiliencePolicy()
	if len(policy) > 0 {
		settings = policy[0]
	}
	return &circuitBreaker{threshold: settings.CircuitThreshold, openFor: settings.CircuitOpenDuration}
}

// allow returns whether the request may proceed and whether it is the
// half-open probe (the single request allowed through an open circuit to
// test recovery).
func (c *circuitBreaker) allow(now time.Time) (allowed bool, probe bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.openUntil.IsZero() {
		return true, false
	}
	if now.Before(c.openUntil) || c.halfOpenRun {
		return false, false
	}
	c.halfOpenRun = true
	return true, true
}

func (c *circuitBreaker) success(probe bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.openUntil.IsZero() {
		// 关闭状态下普通成功:仅清零失败计数。
		c.failures = 0
		return
	}
	if probe && c.halfOpenRun {
		// 半开探针成功:关闭电路,恢复正常流量。
		c.failures = 0
		c.openUntil = time.Time{}
		c.halfOpenRun = false
		return
	}
	// 电路已打开而请求在打开前就通过 allow():陈旧成功不得关闭电路,
	// 否则刚触发熔断的上游被一个在途成功立即放行。
}

// abortProbe 在探针请求被客户端取消(而非上游失败)时调用:
// 既无成功也无失败的证据,释放探针名额但不改变电路状态,让下一次
// allow() 重新发起探针。
func (c *circuitBreaker) abortProbe() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.halfOpenRun = false
}

func (c *circuitBreaker) failure(now time.Time) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.halfOpenRun = false
	c.failures++
	if c.failures >= c.threshold || !c.openUntil.IsZero() {
		c.openUntil = now.Add(c.openFor)
	}
}

type resilientTransport struct {
	base       http.RoundTripper
	circuit    *circuitBreaker
	maxRetries int
	backoff    time.Duration
}

func (t *resilientTransport) RoundTrip(request *http.Request) (*http.Response, error) {
	allowed, probe := t.circuit.allow(time.Now())
	if !allowed {
		return nil, ErrCircuitOpen
	}
	replayable := request.Method == http.MethodGet || request.Method == http.MethodHead ||
		request.Header.Get("Idempotency-Key") != "" && request.GetBody != nil
	attempts := 1
	if replayable {
		attempts += t.maxRetries
	}
	var response *http.Response
	var err error
	for attempt := 0; attempt < attempts; attempt++ {
		current := request
		if attempt > 0 {
			if waitErr := waitBackoff(request.Context(), t.backoff*time.Duration(attempt)); waitErr != nil {
				err = waitErr
				break
			}
			current = request.Clone(request.Context())
			if request.Body != nil && request.Body != http.NoBody {
				if request.GetBody == nil {
					// 请求体不可重放(如 GET/HEAD 携带未设置 GetBody 的 body):
					// 放弃重试,避免把已消费的空 body 重发或调用 nil 方法。
					break
				}
				body, bodyErr := request.GetBody()
				if bodyErr != nil {
					err = bodyErr
					break
				}
				current.Body = body
			}
		}
		response, err = t.base.RoundTrip(current)
		if !retryableResult(response, err) || attempt == attempts-1 {
			break
		}
		if response != nil {
			_, _ = io.CopyN(io.Discard, response.Body, 4096)
			_ = response.Body.Close()
		}
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		// 客户端取消/超时不是上游故障证据:不计成败。若本请求是半开
		// 探针,释放探针名额让电路保持打开,由下一次请求重新探测。
		if probe {
			t.circuit.abortProbe()
		}
		return response, err
	}
	if failureResult(response, err) {
		t.circuit.failure(time.Now())
	} else {
		t.circuit.success(probe)
	}
	return response, err
}

// retryableResult 决定是否值得重试:仅传输错误与 502/503/504 会重试,
// 500 等其余 5xx 不做自动重试(响应可能已被上游处理,重试有副作用)。
func retryableResult(response *http.Response, err error) bool {
	if err != nil {
		return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)
	}
	return response != nil && (response.StatusCode == http.StatusBadGateway || response.StatusCode == http.StatusServiceUnavailable || response.StatusCode == http.StatusGatewayTimeout)
}

// failureResult 决定是否计入熔断失败:所有 5xx 都视为上游故障。否则持续返回
// 500 的上游永远不会触发熔断,而 success() 还会不断清零失败计数,熔断保护失效。
func failureResult(response *http.Response, err error) bool {
	if err != nil {
		return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)
	}
	return response != nil && response.StatusCode >= http.StatusInternalServerError
}

func waitBackoff(ctx context.Context, duration time.Duration) error {
	if duration <= 0 {
		return nil
	}
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}