path: root/internal/limits/limits.go

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