story/internal/ratelimit/limiter.go

package ratelimit

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
)

type Policy struct {
	RatePerSecond float64
	Burst         int
}

type Result struct {
	Allowed    bool
	Remaining  int
	RetryAfter time.Duration
}

type Limiter struct {
	client   *redis.Client
	prefix   string
	policies map[string]Policy
	script   *redis.Script
}

func New(client *redis.Client, prefix string, policy Policy) (*Limiter, error) {
	if client == nil || prefix == "" || policy.RatePerSecond <= 0 || policy.Burst <= 0 {
		return nil, errors.New("redis client, prefix and positive rate policy are required")
	}
	return NewPolicies(client, prefix, map[string]Policy{"*": policy})
}

// NewPolicies builds a limiter whose token-bucket policy is selected by scope.
// A "*" policy is used as the fallback for unrecognised scopes.
func NewPolicies(client *redis.Client, prefix string, policies map[string]Policy) (*Limiter, error) {
	if client == nil || prefix == "" || len(policies) == 0 {
		return nil, errors.New("redis client, prefix and rate policies are required")
	}
	cloned := make(map[string]Policy, len(policies))
	for scope, policy := range policies {
		if scope == "" || policy.RatePerSecond <= 0 || policy.Burst <= 0 {
			return nil, errors.New("rate policy scopes, rates and bursts must be positive")
		}
		cloned[scope] = policy
	}
	return &Limiter{client: client, prefix: prefix, policies: cloned, script: redis.NewScript(tokenBucketScript)}, nil
}

func (l *Limiter) Allow(ctx context.Context, scope, identifier string) (Result, error) {
	if scope == "" || identifier == "" {
		return Result{}, errors.New("rate limit scope and identifier are required")
	}
	policy, ok := l.policies[scope]
	if !ok {
		policy, ok = l.policies["*"]
	}
	if !ok {
		return Result{}, fmt.Errorf("rate limit policy for scope %q is not configured", scope)
	}
	key := fmt.Sprintf("%s:ratelimit:%s:%s", l.prefix, scope, IdentifierHash(identifier))
	values, err := l.script.Run(ctx, l.client, []string{key}, policy.RatePerSecond, policy.Burst).Slice()
	if err != nil {
		return Result{}, fmt.Errorf("execute rate limit: %w", err)
	}
	if len(values) != 3 {
		return Result{}, errors.New("unexpected rate limit response")
	}
	allowed, ok1 := values[0].(int64)
	remaining, ok2 := values[1].(int64)
	retryMilliseconds, ok3 := values[2].(int64)
	if !ok1 || !ok2 || !ok3 {
		return Result{}, errors.New("invalid rate limit response types")
	}
	return Result{
		Allowed: allowed == 1, Remaining: int(remaining),
		RetryAfter: time.Duration(retryMilliseconds) * time.Millisecond,
	}, nil
}

func IdentifierHash(identifier string) string {
	digest := sha256.Sum256([]byte(identifier))
	return hex.EncodeToString(digest[:16])
}

const tokenBucketScript = `
local current_time = redis.call('TIME')
local now = tonumber(current_time[1]) + tonumber(current_time[2]) / 1000000
local rate = tonumber(ARGV[1])
local burst = tonumber(ARGV[2])
local values = redis.call('HMGET', KEYS[1], 'tokens', 'updated_at')
local tokens = tonumber(values[1])
local updated_at = tonumber(values[2])
if tokens == nil then tokens = burst end
if updated_at == nil then updated_at = now end
if now > updated_at then
  tokens = math.min(burst, tokens + (now - updated_at) * rate)
end
local allowed = 0
local retry_ms = 0
if tokens >= 1 then
  allowed = 1
  tokens = tokens - 1
else
  retry_ms = math.ceil(((1 - tokens) / rate) * 1000)
end
redis.call('HSET', KEYS[1], 'tokens', tokens, 'updated_at', now)
redis.call('PEXPIRE', KEYS[1], math.ceil((burst / rate) * 2000) + 1000)
return {allowed, math.floor(tokens), retry_ms}
`