finance-duck/internal/ratelimit/controller.go

// Package ratelimit coordinates bounded HTTP 429 retries for one provider client.
package ratelimit

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	maxRateAttempts    = 4
	maxRateWait        = 2 * time.Minute
	maxLearnedInterval = 30 * time.Second
	maxBackoff         = 15 * time.Minute
)

// Controller serializes provider calls and retains their HTTP 429 cooldown.
// Its zero value is ready for use. A Controller must not be copied after use.
// The active caller owns retries; other callers fail fast during known cooldowns.
type Controller struct {
	// Configure before first use; these fields must remain immutable afterward.
	// Zero values preserve unpaced requests and a one-second initial backoff.
	MinimumInterval time.Duration
	InitialBackoff  time.Duration

	once   sync.Once
	active chan struct{}
	mu     sync.Mutex
	limit  *RateLimitError

	// Only the acquired caller accesses pacing and consecutive-failure state.
	lastAttempt     time.Time
	learnedInterval time.Duration
	backoff         time.Duration
}

// RateLimitError is a safe provider HTTP 429 error. Its message contains only
// the status and retry timing, never provider response text or request details.
// Use errors.As with *RateLimitError to identify it through wrapped errors.
type RateLimitError struct {
	next      time.Time
	unbounded bool
}

// NewError records a trusted provider-specific retry deadline without exposing
// provider response text. A zero deadline disables automatic retries.
func NewError(retryAt time.Time) *RateLimitError {
	return &RateLimitError{next: retryAt, unbounded: retryAt.IsZero()}
}

func (r *RateLimitError) Error() string {
	if r.unbounded {
		return "provider rate limit (HTTP 429): retry time exceeds the supported range; automatic retry disabled"
	}
	// Second precision: this message is read by operators, not machines. Use
	// RetryAt for scheduling.
	return "provider rate limit (HTTP 429): automatic retry at " + r.next.UTC().Format(time.RFC3339)
}

// RetryAt returns the earliest allowed retry time. Zero means the provider's
// delay exceeded the supported range and automatic retries remain disabled.
func (r *RateLimitError) RetryAt() time.Time {
	return r.next
}

func (g *Controller) cooldown() error {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.limit != nil && (g.limit.unbounded || time.Now().Before(g.limit.next)) {
		return g.limit
	}
	return nil
}

// Acquire waits for the active call, unless canceled or a cooldown is known.
// A successful acquisition must be paired with Release, including on errors.
func (g *Controller) Acquire(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("provider request canceled: %w", err)
	}
	if err := g.cooldown(); err != nil {
		return err
	}
	g.once.Do(func() { g.active = make(chan struct{}, 1) })
	select {
	case g.active <- struct{}{}:
	case <-ctx.Done():
		return fmt.Errorf("provider request canceled: %w", ctx.Err())
	}
	if err := ctx.Err(); err != nil {
		g.Release()
		return fmt.Errorf("provider request canceled: %w", err)
	}
	// The preceding request may have established a cooldown while we queued.
	if err := g.cooldown(); err != nil {
		g.Release()
		return err
	}
	return nil
}

// Release relinquishes a successful acquisition.
func (g *Controller) Release() {
	<-g.active
}

// retryLimit never converts a positive overflowing delay into a short wait.
// Delays beyond time.Duration's range disable retries rather than truncate the
// provider's instruction. HTTP dates retain their absolute timestamp unchanged.
func retryLimit(header string, now time.Time, fallback time.Duration) *RateLimitError {
	limit := &RateLimitError{next: now.Add(fallback)}
	header = strings.TrimSpace(header)
	if header == "" {
		return limit
	}
	digits := true
	for _, c := range header {
		if c < '0' || c > '9' {
			digits = false
			break
		}
	}
	if digits {
		seconds, err := strconv.ParseUint(header, 10, 64)
		if err != nil || seconds > uint64((1<<63-1)/int64(time.Second)) {
			return &RateLimitError{unbounded: true}
		}
		if delay := time.Duration(seconds) * time.Second; delay > fallback {
			limit.next = now.Add(delay)
		}
		return limit
	}
	if date, err := http.ParseTime(header); err == nil && date.After(limit.next) {
		limit.next = date
	}
	return limit
}

// Do executes an attempt under an already-acquired Controller. Only HTTP 429
// responses are retried, and only when retry is true (safe/idempotent requests).
// Each 429 body is closed here; other response bodies remain caller-owned.
// The callback must honor ctx and return errors safe to expose to the caller.
// Its per-attempt timeout must not include this controller's retry waiting.
func (g *Controller) Do(ctx context.Context, attempt func(context.Context) (*http.Response, error), retry bool) (*http.Response, error) {
	remainingWait := maxRateWait
	var lastLimit *RateLimitError
	canceled := func(err error) error {
		if lastLimit != nil {
			return fmt.Errorf("%w: %w", lastLimit, err)
		}
		return fmt.Errorf("provider request canceled: %w", err)
	}
	for number := range maxRateAttempts {
		if err := ctx.Err(); err != nil {
			return nil, canceled(err)
		}
		// Pace all attempts, not just retries, including after successful calls.
		next := g.lastAttempt.Add(max(g.MinimumInterval, g.learnedInterval))
		if lastLimit != nil && lastLimit.next.After(next) {
			next = lastLimit.next
		}
		delay := time.Until(next)
		if lastLimit != nil && (lastLimit.unbounded || delay > remainingWait) {
			return nil, lastLimit
		}
		if delay > 0 {
			if lastLimit != nil {
				remainingWait -= delay
			}
			timer := time.NewTimer(delay)
			select {
			case <-ctx.Done():
				timer.Stop()
				return nil, canceled(ctx.Err())
			case <-timer.C:
			}
		}
		if err := ctx.Err(); err != nil {
			return nil, canceled(err)
		}
		g.lastAttempt = time.Now()
		resp, err := attempt(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return nil, canceled(ctx.Err())
			}
			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
				return nil, canceled(err)
			}
			return nil, err
		}
		if resp.StatusCode != http.StatusTooManyRequests {
			if resp.StatusCode < http.StatusBadRequest {
				// Recovery resets consecutive-failure escalation, but never
				// erases learned spacing: a single success is not a quota reset.
				g.backoff = 0
			}
			return resp, nil
		}
		if g.backoff <= 0 {
			g.backoff = g.InitialBackoff
			if g.backoff <= 0 {
				g.backoff = time.Second
			}
		} else if g.backoff >= maxBackoff/2 {
			g.backoff = max(g.backoff, maxBackoff)
		} else {
			g.backoff *= 2
		}
		fallback := max(g.backoff, g.MinimumInterval, g.learnedInterval)
		limit := retryLimit(resp.Header.Get("Retry-After"), time.Now(), fallback)
		g.mu.Lock()
		g.limit = limit
		g.mu.Unlock()
		// Keep the most conservative learned cadence for this controller's
		// lifetime, capped at 30 seconds. The actual provider deadline is never
		// capped; persistent failures separately escalate up to 15 minutes.
		learned := maxLearnedInterval
		if !limit.unbounded {
			learned = min(learned, time.Until(limit.next))
		}
		g.learnedInterval = max(g.learnedInterval, learned)
		// Never read or expose provider errors, and release each response before
		// any sleep or retry. Other responses are processed by the caller.
		resp.Body.Close()
		lastLimit = limit
		if err := ctx.Err(); err != nil {
			return nil, canceled(err)
		}
		if !retry || number == maxRateAttempts-1 {
			return nil, limit
		}
	}
	return nil, lastLimit
}