feat(llm): add configurable request/token rate limiting for LLM client

Introduce process-local and Redis-backed global rate limiting for the shared LLM client, plus a dedicated generation consumer prefetch knob. - Add request (RPM) and token (TPM) rate limits with burst controls, process or global scope, and fail-closed/process-degraded failover - Wire the LLM client through NewWithRedis so limits can be shared across replicas via a Redis token bucket - Add generation_consumer_prefetch to tune RabbitMQ generation consumers independently of the default prefetch - Bump generation worker_concurrency to 24 and promote golang.org/x/time to a direct dependency - Drop t.Parallel() from a few middleware/bootstrap tests Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-10 00:14:40 +08:00
parent 0acd37918e
commit 347cea1296
16 changed files with 1016 additions and 30 deletions
@@ -0,0 +1,472 @@
+package llm
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math"
+	"regexp"
+	"strings"
+	"sync/atomic"
+	"time"
+	"unicode/utf8"
+
+	goredis "github.com/redis/go-redis/v9"
+	"golang.org/x/time/rate"
+
+	"github.com/geo-platform/tenant-api/internal/shared/config"
+)
+
+const (
+	defaultRateLimitMaxOutputTokens int64 = 4000
+	redisRateLimitFailoverCooldown        = 5 * time.Second
+)
+
+var redisTokenBucketScript = goredis.NewScript(`
+local tokens_key = KEYS[1]
+local timestamp_key = KEYS[2]
+local cost = tonumber(ARGV[1])
+local rate = tonumber(ARGV[2])
+local capacity = tonumber(ARGV[3])
+local ttl = tonumber(ARGV[4])
+
+if cost <= 0 or rate <= 0 or capacity <= 0 then
+  return 0
+end
+
+local current_time = redis.call("TIME")
+local now = tonumber(current_time[1]) * 1000 + math.floor(tonumber(current_time[2]) / 1000)
+local tokens = tonumber(redis.call("GET", tokens_key))
+if tokens == nil then
+  tokens = capacity
+end
+local timestamp = tonumber(redis.call("GET", timestamp_key))
+if timestamp == nil then
+  timestamp = now
+end
+if now < timestamp then
+  timestamp = now
+end
+
+local elapsed = now - timestamp
+if elapsed > 0 then
+  tokens = math.min(capacity, tokens + elapsed * rate / 1000)
+end
+
+if tokens >= cost then
+  tokens = tokens - cost
+  redis.call("SET", tokens_key, tokens, "PX", ttl)
+  redis.call("SET", timestamp_key, now, "PX", ttl)
+  return 0
+end
+
+local wait = math.ceil((cost - tokens) / rate * 1000)
+redis.call("SET", tokens_key, tokens, "PX", ttl)
+redis.call("SET", timestamp_key, now, "PX", ttl)
+return wait
+`)
+
+type rateLimitedClient struct {
+	next                   Client
+	requestLimiter         *rate.Limiter
+	tokenLimiter           *rate.Limiter
+	redisLimiter           *redisRateLimiter
+	redisFailoverRequests  *rate.Limiter
+	redisFailoverTokens    *rate.Limiter
+	redisUnavailableUntil  int64
+	rateLimitKey           string
+	requestRateLimitRPM    int
+	requestRateLimitBurst  int
+	tokenRateLimitTPM      int64
+	tokenRateLimitBurst    int
+	maxOutputTokens        int64
+	estimatedCharsPerToken float64
+}
+
+func NewRateLimitedClient(next Client, cfg config.LLMConfig) Client {
+	return NewRateLimitedClientWithRedis(next, cfg, nil)
+}
+
+func NewRateLimitedClientWithRedis(next Client, cfg config.LLMConfig, redisClient goredis.Cmdable) Client {
+	if next == nil {
+		next = disabledClient{reason: "llm client is nil"}
+	}
+
+	config.NormalizeLLMConfig(&cfg)
+	if cfg.RequestRateLimitRPM <= 0 && cfg.TokenRateLimitTPM <= 0 {
+		return next
+	}
+
+	wrapped := &rateLimitedClient{
+		next:                   next,
+		rateLimitKey:           llmRateLimitKey(cfg),
+		requestRateLimitRPM:    cfg.RequestRateLimitRPM,
+		requestRateLimitBurst:  cfg.RequestRateLimitBurst,
+		tokenRateLimitTPM:      cfg.TokenRateLimitTPM,
+		tokenRateLimitBurst:    cfg.TokenRateLimitBurst,
+		maxOutputTokens:        cfg.MaxOutputTokens,
+		estimatedCharsPerToken: cfg.EstimatedCharsPerToken,
+	}
+	if wrapped.maxOutputTokens <= 0 {
+		wrapped.maxOutputTokens = defaultRateLimitMaxOutputTokens
+	}
+
+	if cfg.RateLimitScope == "global" && redisClient != nil {
+		wrapped.redisLimiter = &redisRateLimiter{redis: redisClient}
+		wrapped.configureRedisFailover(cfg)
+		return wrapped
+	}
+
+	if wrapped.requestRateLimitRPM > 0 {
+		wrapped.requestLimiter = rate.NewLimiter(
+			rate.Limit(float64(wrapped.requestRateLimitRPM)/60.0),
+			rateLimitBurst(wrapped.requestRateLimitRPM, wrapped.requestRateLimitBurst),
+		)
+	}
+	if wrapped.tokenRateLimitTPM > 0 {
+		wrapped.tokenLimiter = rate.NewLimiter(
+			rate.Limit(float64(wrapped.tokenRateLimitTPM)/60.0),
+			rateLimitInt64Burst(wrapped.tokenRateLimitTPM, wrapped.tokenRateLimitBurst),
+		)
+	}
+
+	return wrapped
+}
+
+func (c *rateLimitedClient) configureRedisFailover(cfg config.LLMConfig) {
+	if c == nil || cfg.RateLimitFailoverMode != "process_degraded" {
+		return
+	}
+	replicas := cfg.RateLimitFailoverReplicas
+	if replicas <= 0 {
+		replicas = 1
+	}
+	if c.requestRateLimitRPM > 0 {
+		limit := divideLimit(c.requestRateLimitRPM, replicas)
+		burst := divideLimit(rateLimitBurst(c.requestRateLimitRPM, c.requestRateLimitBurst), replicas)
+		c.redisFailoverRequests = rate.NewLimiter(rate.Limit(float64(limit)/60.0), burst)
+	}
+	if c.tokenRateLimitTPM > 0 {
+		limit := divideLimit64(c.tokenRateLimitTPM, replicas)
+		burst := divideLimit(rateLimitInt64Burst(c.tokenRateLimitTPM, c.tokenRateLimitBurst), replicas)
+		c.redisFailoverTokens = rate.NewLimiter(rate.Limit(float64(limit)/60.0), burst)
+	}
+}
+
+func (c *rateLimitedClient) Validate() error {
+	if c == nil || c.next == nil {
+		return disabledClient{reason: "llm client is nil"}.Validate()
+	}
+	return c.next.Validate()
+}
+
+func (c *rateLimitedClient) Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error) {
+	if err := c.Validate(); err != nil {
+		return nil, err
+	}
+	tokens := estimateRateLimitTokens(req, c.maxOutputTokens, c.estimatedCharsPerToken)
+	if c.redisLimiter != nil {
+		if err := c.waitGlobalRateLimit(ctx, "requests", 1, int64(c.requestRateLimitRPM), c.requestRateLimitBurst); err != nil {
+			return nil, err
+		}
+		if err := c.waitGlobalRateLimit(ctx, "tokens", tokens, c.tokenRateLimitTPM, c.tokenRateLimitBurst); err != nil {
+			return nil, err
+		}
+		return c.next.Generate(ctx, req, onDelta)
+	}
+	if c.requestLimiter != nil {
+		if err := c.requestLimiter.Wait(ctx); err != nil {
+			return nil, fmt.Errorf("llm request rate limit wait: %w", err)
+		}
+	}
+	if c.tokenLimiter != nil {
+		if err := waitLimiterN(ctx, c.tokenLimiter, tokens); err != nil {
+			return nil, fmt.Errorf("llm token rate limit wait: %w", err)
+		}
+	}
+	return c.next.Generate(ctx, req, onDelta)
+}
+
+func (c *rateLimitedClient) waitGlobalRateLimit(ctx context.Context, bucket string, cost int, limitPerMinute int64, burst int) error {
+	if c == nil || c.redisLimiter == nil || cost <= 0 || limitPerMinute <= 0 {
+		return nil
+	}
+	if c.redisCircuitOpen(time.Now()) {
+		return c.waitRedisFailoverRateLimit(ctx, bucket, cost)
+	}
+	if err := c.waitRedisRateLimit(ctx, bucket, cost, limitPerMinute, burst); err != nil {
+		if !c.canUseRedisFailover(err) {
+			return err
+		}
+		c.markRedisUnavailable(time.Now())
+		if fallbackErr := c.waitRedisFailoverRateLimit(ctx, bucket, cost); fallbackErr != nil {
+			return fmt.Errorf("llm %s global rate limit failed and process-degraded failover also failed: %w", bucket, errors.Join(err, fallbackErr))
+		}
+		return nil
+	}
+	return nil
+}
+
+func (c *rateLimitedClient) canUseRedisFailover(err error) bool {
+	if c == nil || err == nil {
+		return false
+	}
+	if c.redisFailoverRequests == nil && c.redisFailoverTokens == nil {
+		return false
+	}
+	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+		return false
+	}
+	return true
+}
+
+func (c *rateLimitedClient) waitRedisFailoverRateLimit(ctx context.Context, bucket string, cost int) error {
+	switch bucket {
+	case "requests":
+		if c.redisFailoverRequests == nil {
+			return nil
+		}
+		if err := c.redisFailoverRequests.Wait(ctx); err != nil {
+			return fmt.Errorf("llm %s process-degraded rate limit wait: %w", bucket, err)
+		}
+	case "tokens":
+		if c.redisFailoverTokens == nil {
+			return nil
+		}
+		if err := waitLimiterN(ctx, c.redisFailoverTokens, cost); err != nil {
+			return fmt.Errorf("llm %s process-degraded rate limit wait: %w", bucket, err)
+		}
+	}
+	return nil
+}
+
+func (c *rateLimitedClient) redisCircuitOpen(now time.Time) bool {
+	if c == nil {
+		return false
+	}
+	until := atomic.LoadInt64(&c.redisUnavailableUntil)
+	return until > now.UnixNano()
+}
+
+func (c *rateLimitedClient) markRedisUnavailable(now time.Time) {
+	if c == nil {
+		return
+	}
+	atomic.StoreInt64(&c.redisUnavailableUntil, now.Add(redisRateLimitFailoverCooldown).UnixNano())
+}
+
+func (c *rateLimitedClient) waitRedisRateLimit(ctx context.Context, bucket string, cost int, limitPerMinute int64, burst int) error {
+	if c == nil || c.redisLimiter == nil || cost <= 0 || limitPerMinute <= 0 {
+		return nil
+	}
+	capacity := rateLimitInt64Burst(limitPerMinute, burst)
+	if capacity <= 0 {
+		return nil
+	}
+	remaining := cost
+	for remaining > 0 {
+		chunk := remaining
+		if chunk > capacity {
+			chunk = capacity
+		}
+		if err := c.redisLimiter.Wait(ctx, c.rateLimitKey, bucket, chunk, limitPerMinute, capacity); err != nil {
+			return fmt.Errorf("llm %s global rate limit wait: %w", bucket, err)
+		}
+		remaining -= chunk
+	}
+	return nil
+}
+
+func estimateRateLimitTokens(req GenerateRequest, fallbackMaxOutputTokens int64, charsPerToken float64) int {
+	if charsPerToken <= 0 {
+		charsPerToken = 1
+	}
+	promptChars := utf8.RuneCountInString(strings.TrimSpace(req.Prompt))
+	promptTokens := int64(math.Ceil(float64(promptChars) / charsPerToken))
+	outputTokens := req.MaxOutputTokens
+	if outputTokens <= 0 {
+		outputTokens = fallbackMaxOutputTokens
+	}
+	if outputTokens <= 0 {
+		outputTokens = defaultRateLimitMaxOutputTokens
+	}
+
+	total := promptTokens + outputTokens
+	if total <= 0 {
+		return 1
+	}
+	if total > int64(maxInt()) {
+		return maxInt()
+	}
+	return int(total)
+}
+
+func waitLimiterN(ctx context.Context, limiter *rate.Limiter, n int) error {
+	if limiter == nil || n <= 0 {
+		return nil
+	}
+	burst := limiter.Burst()
+	if burst <= 0 {
+		return nil
+	}
+	for n > 0 {
+		chunk := n
+		if chunk > burst {
+			chunk = burst
+		}
+		if err := limiter.WaitN(ctx, chunk); err != nil {
+			return err
+		}
+		n -= chunk
+	}
+	return nil
+}
+
+func rateLimitBurst(limit, configured int) int {
+	if configured > 0 {
+		return configured
+	}
+	burst := limit / 10
+	if burst < 1 {
+		return 1
+	}
+	return burst
+}
+
+func rateLimitInt64Burst(limit int64, configured int) int {
+	if configured > 0 {
+		return configured
+	}
+	burst := limit / 10
+	if burst < 1 {
+		return 1
+	}
+	if burst > int64(maxInt()) {
+		return maxInt()
+	}
+	return int(burst)
+}
+
+func divideLimit(limit, divisor int) int {
+	if limit <= 0 {
+		return 0
+	}
+	if divisor <= 1 {
+		return limit
+	}
+	quotient := limit / divisor
+	if quotient < 1 {
+		return 1
+	}
+	return quotient
+}
+
+func divideLimit64(limit int64, divisor int) int64 {
+	if limit <= 0 {
+		return 0
+	}
+	if divisor <= 1 {
+		return limit
+	}
+	quotient := limit / int64(divisor)
+	if quotient < 1 {
+		return 1
+	}
+	return quotient
+}
+
+func maxInt() int {
+	return int(^uint(0) >> 1)
+}
+
+type redisRateLimiter struct {
+	redis goredis.Cmdable
+}
+
+func (l *redisRateLimiter) Wait(ctx context.Context, group, bucket string, cost int, limitPerMinute int64, capacity int) error {
+	if l == nil || l.redis == nil || cost <= 0 || limitPerMinute <= 0 || capacity <= 0 {
+		return nil
+	}
+
+	ratePerSecond := float64(limitPerMinute) / 60.0
+	if ratePerSecond <= 0 {
+		return nil
+	}
+	ttl := redisBucketTTL(ratePerSecond, capacity)
+	keys := redisBucketKeys(group, bucket)
+
+	for {
+		waitMillis, err := redisTokenBucketScript.Run(
+			ctx,
+			l.redis,
+			keys,
+			cost,
+			ratePerSecond,
+			capacity,
+			ttl.Milliseconds(),
+		).Int64()
+		if err != nil {
+			return err
+		}
+		if waitMillis <= 0 {
+			return nil
+		}
+
+		timer := time.NewTimer(time.Duration(waitMillis) * time.Millisecond)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return ctx.Err()
+		case <-timer.C:
+		}
+	}
+}
+
+func redisBucketTTL(ratePerSecond float64, capacity int) time.Duration {
+	if ratePerSecond <= 0 || capacity <= 0 {
+		return 2 * time.Minute
+	}
+	refill := time.Duration(math.Ceil(float64(capacity)/ratePerSecond)) * time.Second
+	if refill < time.Minute {
+		return 2 * time.Minute
+	}
+	return refill * 2
+}
+
+func redisBucketKeys(group, bucket string) []string {
+	escapedGroup := sanitizeRateLimitKey(group)
+	escapedBucket := sanitizeRateLimitKey(bucket)
+	prefix := fmt.Sprintf("geo:llm:rate:{%s}:%s", escapedGroup, escapedBucket)
+	return []string{prefix + ":tokens", prefix + ":ts"}
+}
+
+func llmRateLimitKey(cfg config.LLMConfig) string {
+	if key := strings.TrimSpace(cfg.RateLimitKey); key != "" {
+		return key
+	}
+	parts := []string{
+		strings.TrimSpace(cfg.Provider),
+		strings.TrimSpace(cfg.BaseURL),
+		strings.TrimSpace(cfg.Model),
+	}
+	key := strings.Join(parts, ":")
+	if strings.Trim(key, ": ") == "" {
+		return "default"
+	}
+	return key
+}
+
+var rateLimitKeySanitizer = regexp.MustCompile(`[^a-zA-Z0-9_.:-]+`)
+
+func sanitizeRateLimitKey(value string) string {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return "default"
+	}
+	value = rateLimitKeySanitizer.ReplaceAllString(value, "_")
+	value = strings.Trim(value, "_")
+	if value == "" {
+		return "default"
+	}
+	return value
+}