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}
}

func (c *circuitBreaker) allow(now time.Time) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.openUntil.IsZero() {
		return true
	}
	if now.Before(c.openUntil) || c.halfOpenRun {
		return false
	}
	c.halfOpenRun = true
	return true
}

func (c *circuitBreaker) success() {
	c.mu.Lock()
	c.failures = 0
	c.openUntil = time.Time{}
	c.halfOpenRun = false
	c.mu.Unlock()
}

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) {
	if !t.circuit.allow(time.Now()) {
		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 {
				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 retryableResult(response, err) {
		t.circuit.failure(time.Now())
	} else {
		t.circuit.success()
	}
	return response, err
}

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)
}

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
	}
}