broad-announce/internal/ratelimit/limiter.go

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