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
)

// 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 {
	once   sync.Once
	active chan struct{}
	mu     sync.Mutex
	limit  *RateLimitError
}

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

func (r *RateLimitError) Error() string {
	if r.unbounded {
		return "provider rate limit (HTTP 429): retry time exceeds the supported range; automatic retry disabled"
	}
	return "provider rate limit (HTTP 429): retry allowed at " + r.next.UTC().Format(time.RFC3339Nano)
}

// 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 error
	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)
		}
		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 {
			return resp, nil
		}
		limit := retryLimit(resp.Header.Get("Retry-After"), time.Now(), time.Second<<number)
		g.mu.Lock()
		g.limit = limit
		g.mu.Unlock()
		// 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)
		}
		delay := time.Until(limit.next)
		if !retry || number == maxRateAttempts-1 || limit.unbounded || delay > remainingWait {
			return nil, limit
		}
		if delay <= 0 {
			continue
		}
		remainingWait -= delay
		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, canceled(ctx.Err())
		case <-timer.C:
		}
	}
	return nil, lastLimit
}