// Package ratelimit implements the per-source and per-company token // bucket from SPEC §22 layers 3 & 4. // // We use the INCR + EXPIRE pattern (works on stock Redis, no // redis-cell module needed). The window is 1 second; rate is // cap tokens / second. Tokens are integer; the bucket refills // atomically on the next request. // // This is an *approximation* of a token bucket — a leaky bucket // would be more accurate. Good enough for protecting the tier. package ratelimit import ( "context" "fmt" "time" "github.com/redis/go-redis/v9" ) const windowSeconds = 1 type Limiter struct { rdb *redis.Client } func New(rdb *redis.Client) *Limiter { return &Limiter{rdb: rdb} } // Allow consumes one token from the bucket `key` and returns: // // (true, 0, nil) – allowed // (false, ttl, nil) – denied, ttl is the seconds until reset // (false, 0, err) – redis error // // The cap is the per-second limit. If cap <= 0 the limiter is a // no-op (always allow) — useful for sources we explicitly disable. func (l *Limiter) Allow(ctx context.Context, key string, cap int) (bool, time.Duration, error) { if cap <= 0 { return true, 0, nil } full := fmt.Sprintf("rl:%s:%d", key, time.Now().Unix()) pipe := l.rdb.Pipeline() incr := pipe.Incr(ctx, full) pipe.Expire(ctx, full, windowSeconds*2*time.Second) if _, err := pipe.Exec(ctx); err != nil { return false, 0, fmt.Errorf("ratelimit: %w", err) } n := incr.Val() if n > int64(cap) { ttl, _ := l.rdb.TTL(ctx, full).Result() if ttl < 0 { ttl = time.Second } return false, ttl, nil } return true, 0, nil }