134 lines
3.7 KiB
Go
134 lines
3.7 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
const (
|
|
defaultRatePerSecond = 100
|
|
defaultKeyPrefix = "rate:gateway:channel:"
|
|
)
|
|
|
|
// Limiter reserves one distributed, per-channel submit slot and waits until it
|
|
// becomes available. The reservation is stored in Redis so multiple Gateway
|
|
// instances share one supplier TPS budget.
|
|
type Limiter interface {
|
|
Wait(context.Context, string, int) (time.Duration, error)
|
|
}
|
|
|
|
// Configurer stores the authoritative channel limit received with the channel
|
|
// connection command. Wait always applies the lower of this value and the
|
|
// message value, so a producer cannot raise the supplier TPS ceiling.
|
|
type Configurer interface {
|
|
Configure(context.Context, string, int) error
|
|
}
|
|
|
|
type RedisLimiter struct {
|
|
Redis *redis.Client
|
|
KeyPrefix string
|
|
}
|
|
|
|
var reserveScript = redis.NewScript(`
|
|
local now = redis.call('TIME')
|
|
local now_us = (tonumber(now[1]) * 1000000) + tonumber(now[2])
|
|
local requested_rate = tonumber(ARGV[1])
|
|
local configured_rate = tonumber(redis.call('GET', KEYS[2]))
|
|
local effective_rate = requested_rate
|
|
if configured_rate and configured_rate > 0 and configured_rate < effective_rate then
|
|
effective_rate = configured_rate
|
|
end
|
|
local interval_us = math.ceil(1000000 / effective_rate)
|
|
local next_us = tonumber(redis.call('GET', KEYS[1])) or now_us
|
|
if next_us < now_us then
|
|
next_us = now_us
|
|
end
|
|
local delay_us = next_us - now_us
|
|
local reserved_until_us = next_us + interval_us
|
|
local ttl_ms = math.ceil((reserved_until_us - now_us) / 1000) + 1000
|
|
redis.call('PSETEX', KEYS[1], ttl_ms, reserved_until_us)
|
|
return delay_us
|
|
`)
|
|
|
|
func New(redisURL string) (*RedisLimiter, error) {
|
|
if strings.TrimSpace(redisURL) == "" {
|
|
redisURL = "redis://127.0.0.1:6379"
|
|
}
|
|
options, err := redis.ParseURL(redisURL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return NewWithClient(redis.NewClient(options)), nil
|
|
}
|
|
|
|
func NewWithClient(client *redis.Client) *RedisLimiter {
|
|
return &RedisLimiter{Redis: client, KeyPrefix: defaultKeyPrefix}
|
|
}
|
|
|
|
func (l *RedisLimiter) Configure(ctx context.Context, channelID string, ratePerSecond int) error {
|
|
if l == nil || l.Redis == nil {
|
|
return fmt.Errorf("gateway rate limiter Redis client is required")
|
|
}
|
|
channelID = strings.TrimSpace(channelID)
|
|
if channelID == "" {
|
|
return fmt.Errorf("gateway rate limiter channelId is required")
|
|
}
|
|
if ratePerSecond <= 0 {
|
|
ratePerSecond = defaultRatePerSecond
|
|
}
|
|
if err := l.Redis.Set(ctx, l.configKey(channelID), ratePerSecond, 0).Err(); err != nil {
|
|
return fmt.Errorf("configure gateway channel rate: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (l *RedisLimiter) Wait(ctx context.Context, channelID string, ratePerSecond int) (time.Duration, error) {
|
|
if l == nil || l.Redis == nil {
|
|
return 0, fmt.Errorf("gateway rate limiter Redis client is required")
|
|
}
|
|
channelID = strings.TrimSpace(channelID)
|
|
if channelID == "" {
|
|
return 0, fmt.Errorf("gateway rate limiter channelId is required")
|
|
}
|
|
if ratePerSecond <= 0 {
|
|
ratePerSecond = defaultRatePerSecond
|
|
}
|
|
delayMicros, err := reserveScript.Run(
|
|
ctx,
|
|
l.Redis,
|
|
[]string{l.key(channelID), l.configKey(channelID)},
|
|
ratePerSecond,
|
|
).Int64()
|
|
if err != nil {
|
|
return 0, fmt.Errorf("reserve gateway channel rate slot: %w", err)
|
|
}
|
|
delay := time.Duration(delayMicros) * time.Microsecond
|
|
if delay <= 0 {
|
|
return 0, nil
|
|
}
|
|
timer := time.NewTimer(delay)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-ctx.Done():
|
|
return delay, ctx.Err()
|
|
case <-timer.C:
|
|
return delay, nil
|
|
}
|
|
}
|
|
|
|
func (l *RedisLimiter) configKey(channelID string) string {
|
|
return l.key("config:" + channelID)
|
|
}
|
|
|
|
func (l *RedisLimiter) key(channelID string) string {
|
|
prefix := l.KeyPrefix
|
|
if prefix == "" {
|
|
prefix = defaultKeyPrefix
|
|
}
|
|
return prefix + channelID
|
|
}
|