1 files changed, 316 insertions, 0 deletions
diff --git a/internal/limits/limits.go b/internal/limits/limits.go
new file mode 100644
index 0000000..6e0bf74
--- /dev/null
+++ b/internal/limits/limits.go
@@ -0,0 +1,316 @@
+package limits
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"log/slog"
+	"math"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"aigw/internal/domain"
+
+	"github.com/redis/go-redis/v9"
+)
+
+var (
+	ErrRequestsExceeded = errors.New("requests per minute limit exceeded")
+	ErrTokensExceeded   = errors.New("tokens per minute limit exceeded")
+	ErrConcurrencyLimit = errors.New("concurrent request limit exceeded")
+)
+
+type Lease interface{ Release() }
+
+type Limiter struct {
+	redis            *redis.Client
+	prefix           string
+	defaultMaxOutput int64
+	logger           *slog.Logger
+	policies         atomic.Pointer[policySnapshot]
+	localMu          sync.Mutex
+	local            map[string]*localWindow
+	redisUp          atomic.Bool
+	redisFailureSeen atomic.Bool
+	redisRetryAt     atomic.Int64
+}
+
+type policySnapshot struct{ policies map[string]domain.LimitPolicy }
+
+type localWindow struct {
+	minute     int64
+	requests   int64
+	tokens     int64
+	concurrent int64
+}
+
+type localLease struct {
+	limiter  *Limiter
+	project  string
+	released atomic.Bool
+}
+type redisLease struct {
+	limiter  *Limiter
+	key      string
+	released atomic.Bool
+}
+
+const (
+	redisCommandTimeout = 250 * time.Millisecond
+	redisRetryCooldown  = 5 * time.Second
+)
+
+const acquireScript = `
+local req = tonumber(ARGV[1])
+local tok = tonumber(ARGV[2])
+local conc = tonumber(ARGV[3])
+local estimate = tonumber(ARGV[4])
+local ttl = tonumber(ARGV[5])
+local r = 0
+local t = 0
+local c = 0
+if req > 0 then r = redis.call('INCR', KEYS[1]); if r == 1 then redis.call('PEXPIRE', KEYS[1], ttl) end end
+if tok > 0 then t = redis.call('INCRBY', KEYS[2], estimate); if t == estimate then redis.call('PEXPIRE', KEYS[2], ttl) end end
+if conc > 0 then c = redis.call('INCR', KEYS[3]); redis.call('PEXPIRE', KEYS[3], 3600000) end
+if (req > 0 and r > req) then if req > 0 then redis.call('DECR', KEYS[1]) end; if tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end; if conc > 0 then redis.call('DECR', KEYS[3]) end; return {0,1} end
+if (tok > 0 and t > tok) then if req > 0 then redis.call('DECR', KEYS[1]) end; if tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end; if conc > 0 then redis.call('DECR', KEYS[3]) end; return {0,2} end
+if (conc > 0 and c > conc) then if req > 0 then redis.call('DECR', KEYS[1]) end; if tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end; if conc > 0 then redis.call('DECR', KEYS[3]) end; return {0,3} end
+return {1,0}
+`
+
+const releaseScript = `
+local value = tonumber(redis.call('GET', KEYS[1]) or '0')
+if value > 0 then redis.call('DECR', KEYS[1]) end
+return value
+`
+
+func New(redisURL, prefix string, defaultMaxOutput int64, logger *slog.Logger) *Limiter {
+	if logger == nil {
+		logger = slog.Default()
+	}
+	if prefix == "" {
+		prefix = "aigw:limits"
+	}
+	l := &Limiter{prefix: strings.TrimRight(prefix, ":"), defaultMaxOutput: defaultMaxOutput, logger: logger, local: make(map[string]*localWindow)}
+	if strings.TrimSpace(redisURL) != "" {
+		if options, err := redis.ParseURL(redisURL); err == nil {
+			options.MaxRetries = -1
+			options.DialerRetries = 1
+			options.DialTimeout = redisCommandTimeout
+			options.ReadTimeout = redisCommandTimeout
+			options.WriteTimeout = redisCommandTimeout
+			options.PoolTimeout = redisCommandTimeout
+			l.redis = redis.NewClient(options)
+		} else {
+			logger.Warn("limits_redis_config_invalid", "error", err, "fallback", "local")
+		}
+	}
+	l.ReplacePolicies(nil)
+	return l
+}
+
+func (l *Limiter) Close() error {
+	if l.redis == nil {
+		return nil
+	}
+	return l.redis.Close()
+}
+
+func (l *Limiter) ReplacePolicies(policies []domain.LimitPolicy) {
+	copyMap := make(map[string]domain.LimitPolicy, len(policies))
+	for _, policy := range policies {
+		if policy.ProjectID != "" {
+			copyMap[policy.ProjectID] = policy
+		}
+	}
+	l.policies.Store(&policySnapshot{policies: copyMap})
+}
+
+func (l *Limiter) Policy(projectID string) (domain.LimitPolicy, bool) {
+	snapshot := l.policies.Load()
+	if snapshot == nil {
+		return domain.LimitPolicy{}, false
+	}
+	policy, ok := snapshot.policies[projectID]
+	return policy, ok
+}
+
+func (l *Limiter) Acquire(ctx context.Context, principal domain.Principal, body []byte) (Lease, error) {
+	policy, ok := l.Policy(principal.ProjectID)
+	if !ok || (policy.RequestsPerMinute == 0 && policy.TokensPerMinute == 0 && policy.Concurrent == 0) {
+		return noopLease{}, nil
+	}
+	estimate := EstimateTokens(body, l.defaultMaxOutput)
+	minute := time.Now().Unix() / 60
+	if l.redis != nil && time.Now().UnixNano() >= l.redisRetryAt.Load() {
+		redisContext, cancel := context.WithTimeout(ctx, redisCommandTimeout)
+		lease, err := l.acquireRedis(redisContext, principal.ProjectID, minute, policy, estimate)
+		cancel()
+		if err == nil {
+			return lease, nil
+		}
+		if errors.Is(err, ErrRequestsExceeded) || errors.Is(err, ErrTokensExceeded) || errors.Is(err, ErrConcurrencyLimit) {
+			l.markRedis(true, nil)
+			return nil, err
+		}
+		l.markRedis(false, err)
+	}
+	return l.acquireLocal(principal.ProjectID, minute, policy, estimate)
+}
+
+func EstimateTokens(body []byte, defaultMaxOutput int64) int64 {
+	input := int64((len(body) + 3) / 4)
+	if input < 1 {
+		input = 1
+	}
+	maxOutput := defaultMaxOutput
+	var limits struct {
+		MaxTokens           int64 `json:"max_tokens"`
+		MaxCompletionTokens int64 `json:"max_completion_tokens"`
+		MaxOutputTokens     int64 `json:"max_output_tokens"`
+	}
+	if json.Unmarshal(body, &limits) == nil {
+		explicit := int64(0)
+		for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} {
+			if value > explicit {
+				explicit = value
+			}
+		}
+		if explicit > 0 {
+			maxOutput = explicit
+		}
+	}
+	if maxOutput < 0 {
+		maxOutput = 0
+	}
+	if input > math.MaxInt64-maxOutput {
+		return math.MaxInt64
+	}
+	return input + maxOutput
+}
+
+func (l *Limiter) acquireRedis(ctx context.Context, project string, minute int64, policy domain.LimitPolicy, estimate int64) (Lease, error) {
+	base := l.prefix + ":" + project + ":" + strconv.FormatInt(minute, 10)
+	keys := []string{base + ":requests", base + ":tokens", l.prefix + ":" + project + ":concurrent"}
+	values, err := l.redis.Eval(ctx, acquireScript, keys, policy.RequestsPerMinute, policy.TokensPerMinute, policy.Concurrent, estimate, 125000).Result()
+	if err != nil {
+		return nil, err
+	}
+	items, ok := values.([]any)
+	if !ok || len(items) < 2 {
+		return nil, errors.New("invalid limits Redis response")
+	}
+	allowed, _ := toInt64(items[0])
+	reason, _ := toInt64(items[1])
+	if allowed == 0 {
+		switch reason {
+		case 1:
+			return nil, ErrRequestsExceeded
+		case 2:
+			return nil, ErrTokensExceeded
+		default:
+			return nil, ErrConcurrencyLimit
+		}
+	}
+	l.markRedis(true, nil)
+	return &redisLease{limiter: l, key: keys[2]}, nil
+}
+
+func (l *Limiter) acquireLocal(project string, minute int64, policy domain.LimitPolicy, estimate int64) (Lease, error) {
+	l.localMu.Lock()
+	defer l.localMu.Unlock()
+	for key, value := range l.local {
+		if value.minute < minute-2 && value.concurrent == 0 {
+			delete(l.local, key)
+		}
+	}
+	window := l.local[project]
+	if window == nil {
+		window = &localWindow{minute: minute}
+		l.local[project] = window
+	} else if window.minute != minute {
+		window.minute = minute
+		window.requests = 0
+		window.tokens = 0
+	}
+	if policy.RequestsPerMinute > 0 && window.requests >= policy.RequestsPerMinute {
+		return nil, ErrRequestsExceeded
+	}
+	if policy.TokensPerMinute > 0 && (estimate > policy.TokensPerMinute || window.tokens > policy.TokensPerMinute-estimate) {
+		return nil, ErrTokensExceeded
+	}
+	if policy.Concurrent > 0 && window.concurrent >= policy.Concurrent {
+		return nil, ErrConcurrencyLimit
+	}
+	window.requests++
+	window.tokens += estimate
+	window.concurrent++
+	return &localLease{limiter: l, project: project}, nil
+}
+
+func (l *Limiter) releaseLocal(project string) {
+	l.localMu.Lock()
+	defer l.localMu.Unlock()
+	if value := l.local[project]; value != nil && value.concurrent > 0 {
+		value.concurrent--
+	}
+}
+
+func (l *Limiter) releaseRedis(ctx context.Context, key string) {
+	if l.redis == nil {
+		return
+	}
+	if _, err := l.redis.Eval(ctx, releaseScript, []string{key}).Result(); err != nil {
+		l.markRedis(false, err)
+	}
+}
+
+func (x *localLease) Release() {
+	if x.released.CompareAndSwap(false, true) {
+		x.limiter.releaseLocal(x.project)
+	}
+}
+func (x *redisLease) Release() {
+	if x.released.CompareAndSwap(false, true) {
+		ctx, cancel := context.WithTimeout(context.Background(), redisCommandTimeout)
+		defer cancel()
+		x.limiter.releaseRedis(ctx, x.key)
+	}
+}
+
+type noopLease struct{}
+
+func (noopLease) Release() {}
+
+func (l *Limiter) markRedis(healthy bool, err error) {
+	l.redisUp.Store(healthy)
+	if healthy {
+		l.redisRetryAt.Store(0)
+		if l.redisFailureSeen.Swap(false) {
+			l.logger.Info("limits_redis_recovered")
+		}
+		return
+	}
+	l.redisRetryAt.Store(time.Now().Add(redisRetryCooldown).UnixNano())
+	if l.redisFailureSeen.CompareAndSwap(false, true) {
+		l.logger.Warn("limits_redis_unavailable", "error", err, "fallback", "local", "retry_after", redisRetryCooldown)
+	}
+}
+
+func toInt64(value any) (int64, bool) {
+	switch v := value.(type) {
+	case int64:
+		return v, true
+	case string:
+		n, e := strconv.ParseInt(v, 10, 64)
+		return n, e == nil
+	case []byte:
+		n, e := strconv.ParseInt(string(v), 10, 64)
+		return n, e == nil
+	default:
+		return 0, false
+	}
+}