iotlcm/internal/service/inbound/rate_limiter.go

package inbound

import (
	"context"
	"fmt"
	"math"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model"
	"github.com/gogf/gf/v2/container/gvar"
	"github.com/gogf/gf/v2/frame/g"
)

const (
	inboundUpstreamRateLimitConfigKey = "batch.inbound.upstreamRateLimit"

	defaultInboundUpstreamRateLimitMode       = "redis"
	defaultInboundUpstreamRateLimitRedisGroup = "default"
	defaultInboundUpstreamRateLimitKeyPrefix  = "iotlcm:worker:inbound:upstream_basic_info"
	defaultInboundUpstreamQPS                 = 20
	defaultInboundUpstreamBurst               = 1
	defaultInboundUpstreamRequestCost         = 1
	defaultInboundChannelTempKeyCacheTTL      = time.Minute

	inboundRateLimitFallbackLocal = "local"
	inboundRateLimitFallbackError = "error"

	inboundCarrierTelecom = "telecom"
	inboundCarrierMobile  = "mobile"
)

type upstreamBasicInfoRateLimiter interface {
	Wait(ctx context.Context, chanKey string, requestCount int) error
}

type redisEvaler interface {
	Eval(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error)
}

type inboundUpstreamRateLimitConfig struct {
	Enabled              *bool                                   `json:"enabled"`
	Mode                 string                                  `json:"mode"`
	RedisGroup           string                                  `json:"redisGroup"`
	KeyPrefix            string                                  `json:"keyPrefix"`
	RedisFailureFallback string                                  `json:"redisFailureFallback"`
	QPS                  float64                                 `json:"qps"`
	Burst                int                                     `json:"burst"`
	RequestCost          int                                     `json:"requestCost"`
	Default              inboundUpstreamRateLimitRule            `json:"default"`
	Carriers             map[string]inboundUpstreamRateLimitRule `json:"carriers"`
	Channels             map[string]inboundUpstreamRateLimitRule `json:"channels"`
}

type inboundUpstreamRateLimitRule struct {
	Enabled     *bool   `json:"enabled"`
	QPS         float64 `json:"qps"`
	Burst       int     `json:"burst"`
	RequestCost int     `json:"requestCost"`
}

type inboundUpstreamRateLimitSetting struct {
	enabled              bool
	mode                 string
	redisGroup           string
	keyPrefix            string
	redisFailureFallback string
	qps                  float64
	burst                int
	requestCost          int
	carrier              string
	tempKey              string
	skipReason           string
}

type inboundChannelMetadata struct {
	TempKey  string `orm:"temp_key"`
	Operator int    `orm:"operator"`
}

type inboundChannelTempKeyCacheEntry struct {
	metadata  inboundChannelMetadata
	expiresAt time.Time
}

type configUpstreamBasicInfoRateLimiter struct {
	redis           func(string) redisEvaler
	channelMetadata func(context.Context, string) (inboundChannelMetadata, error)
	channelTempKey  func(context.Context, string) (string, error)
	local           *localTokenBucketLimiter
	channelCache    sync.Map
	redisLogTimes   sync.Map
	skipLogTimes    sync.Map
}

func newUpstreamBasicInfoRateLimiter() *configUpstreamBasicInfoRateLimiter {
	return &configUpstreamBasicInfoRateLimiter{
		redis: func(name string) redisEvaler {
			return g.Redis(name)
		},
		channelMetadata: queryInboundChannelMetadata,
		local:           newLocalTokenBucketLimiter(),
	}
}

func (l *configUpstreamBasicInfoRateLimiter) Wait(ctx context.Context, chanKey string, requestCount int) error {
	if l == nil {
		return nil
	}
	cfg := loadInboundUpstreamRateLimitConfig(ctx)
	if !inboundUpstreamRateLimitGloballyEnabled(cfg) {
		return nil
	}
	metadata, err := l.lookupChannelMetadata(ctx, chanKey)
	if err != nil {
		return err
	}
	setting := resolveInboundUpstreamRateLimitFromChannelMetadata(cfg, chanKey, metadata.Operator, metadata.TempKey)
	if !setting.enabled || setting.qps <= 0 || setting.mode == "off" {
		if setting.skipReason != "" {
			l.logSkipped(ctx, chanKey, setting)
		}
		return nil
	}
	if requestCount <= 0 {
		requestCount = 1
	}
	permits := requestCount * setting.requestCost
	key := inboundUpstreamRateLimitKey(setting.keyPrefix, chanKey, setting.carrier)
	for i := 0; i < permits; i++ {
		if err := l.waitOne(ctx, setting, key); err != nil {
			return err
		}
	}
	return nil
}

func (l *configUpstreamBasicInfoRateLimiter) waitOne(ctx context.Context, setting inboundUpstreamRateLimitSetting, key string) error {
	if setting.mode == "local" {
		if l.local == nil {
			return fmt.Errorf("local rate limiter is nil")
		}
		return l.local.Wait(ctx, key, setting.qps, setting.burst)
	}
	if l.redis == nil {
		return l.handleRedisError(ctx, setting, key, fmt.Errorf("redis client resolver is nil"))
	}
	redis := l.redis(setting.redisGroup)
	if redis == nil {
		return l.handleRedisError(ctx, setting, key, fmt.Errorf("redis client is nil"))
	}
	if err := (&redisTokenBucketLimiter{redis: redis}).Wait(ctx, key, setting.qps, setting.burst); err != nil {
		return l.handleRedisError(ctx, setting, key, err)
	}
	return nil
}

func (l *configUpstreamBasicInfoRateLimiter) handleRedisError(ctx context.Context, setting inboundUpstreamRateLimitSetting, key string, err error) error {
	if setting.redisFailureFallback == inboundRateLimitFallbackError {
		return err
	}
	if l.local == nil {
		return err
	}
	l.logRedisFallback(ctx, key, err)
	return l.local.Wait(ctx, key, setting.qps, setting.burst)
}

func (l *configUpstreamBasicInfoRateLimiter) logRedisFallback(ctx context.Context, key string, err error) {
	now := time.Now()
	if value, ok := l.redisLogTimes.Load(key); ok {
		if last, ok := value.(time.Time); ok && now.Sub(last) < 10*time.Second {
			return
		}
	}
	l.redisLogTimes.Store(key, now)
	g.Log().Errorf(ctx, "inbound upstream basic-info redis rate limiter failed, fallback local, key:%s, err:%+v", key, err)
}

func (l *configUpstreamBasicInfoRateLimiter) lookupChannelMetadata(ctx context.Context, chanKey string) (inboundChannelMetadata, error) {
	chanKey = strings.TrimSpace(chanKey)
	if chanKey == "" {
		return inboundChannelMetadata{}, nil
	}
	now := time.Now()
	if value, ok := l.channelCache.Load(chanKey); ok {
		if entry, ok := value.(inboundChannelTempKeyCacheEntry); ok && now.Before(entry.expiresAt) {
			return entry.metadata, nil
		}
	}
	if l.channelMetadata != nil {
		metadata, err := l.channelMetadata(ctx, chanKey)
		if err != nil {
			return inboundChannelMetadata{}, fmt.Errorf("query channel metadata failed, chanKey:%s: %w", chanKey, err)
		}
		metadata.TempKey = strings.TrimSpace(metadata.TempKey)
		l.channelCache.Store(chanKey, inboundChannelTempKeyCacheEntry{
			metadata:  metadata,
			expiresAt: now.Add(defaultInboundChannelTempKeyCacheTTL),
		})
		return metadata, nil
	}
	if l.channelTempKey == nil {
		return inboundChannelMetadata{}, fmt.Errorf("channel metadata resolver is nil")
	}
	tempKey, err := l.channelTempKey(ctx, chanKey)
	if err != nil {
		return inboundChannelMetadata{}, fmt.Errorf("query channel template failed, chanKey:%s: %w", chanKey, err)
	}
	metadata := inboundChannelMetadata{TempKey: strings.TrimSpace(tempKey)}
	l.channelCache.Store(chanKey, inboundChannelTempKeyCacheEntry{
		metadata:  metadata,
		expiresAt: now.Add(defaultInboundChannelTempKeyCacheTTL),
	})
	return metadata, nil
}

func (l *configUpstreamBasicInfoRateLimiter) logSkipped(ctx context.Context, chanKey string, setting inboundUpstreamRateLimitSetting) {
	key := strings.TrimSpace(chanKey) + ":" + setting.skipReason
	now := time.Now()
	if value, ok := l.skipLogTimes.Load(key); ok {
		if last, ok := value.(time.Time); ok && now.Sub(last) < 10*time.Second {
			return
		}
	}
	l.skipLogTimes.Store(key, now)
	g.Log().Warningf(ctx, "inbound upstream basic-info rate limit skipped, chanKey:%s, tempKey:%s, reason:%s",
		chanKey, setting.tempKey, setting.skipReason)
}

func queryInboundChannelMetadata(ctx context.Context, chanKey string) (inboundChannelMetadata, error) {
	var row inboundChannelMetadata
	if err := g.DB().Model("t_sim_channels").Ctx(ctx).
		Fields("temp_key,operator").
		Where("chan_key", strings.TrimSpace(chanKey)).
		Limit(1).
		Scan(&row); err != nil {
		return inboundChannelMetadata{}, err
	}
	row.TempKey = strings.TrimSpace(row.TempKey)
	return row, nil
}

func queryInboundChannelTempKey(ctx context.Context, chanKey string) (string, error) {
	metadata, err := queryInboundChannelMetadata(ctx, chanKey)
	if err != nil {
		return "", err
	}
	return metadata.TempKey, nil
}

func resolveInboundUpstreamRateLimit(ctx context.Context, chanKey string) inboundUpstreamRateLimitSetting {
	return resolveInboundUpstreamRateLimitWithTempKey(ctx, chanKey, "")
}

func resolveInboundUpstreamRateLimitWithTempKey(ctx context.Context, chanKey, tempKey string) inboundUpstreamRateLimitSetting {
	cfg := loadInboundUpstreamRateLimitConfig(ctx)
	return resolveInboundUpstreamRateLimitFromConfig(cfg, chanKey, tempKey)
}

func resolveInboundUpstreamRateLimitWithChannelMetadata(ctx context.Context, chanKey string, operator int, tempKey string) inboundUpstreamRateLimitSetting {
	cfg := loadInboundUpstreamRateLimitConfig(ctx)
	return resolveInboundUpstreamRateLimitFromChannelMetadata(cfg, chanKey, operator, tempKey)
}

func loadInboundUpstreamRateLimitConfig(ctx context.Context) inboundUpstreamRateLimitConfig {
	cfg := defaultInboundUpstreamRateLimitConfig()
	value, err := g.Cfg().Get(ctx, inboundUpstreamRateLimitConfigKey)
	if err != nil {
		g.Log().Warningf(ctx, "load inbound upstream rate limit config failed: %+v", err)
		return normalizeInboundUpstreamRateLimitConfig(cfg)
	}
	if value != nil && !value.IsNil() {
		if err = value.Scan(&cfg); err != nil {
			g.Log().Warningf(ctx, "parse inbound upstream rate limit config failed: %+v", err)
			cfg = defaultInboundUpstreamRateLimitConfig()
		}
	}
	return normalizeInboundUpstreamRateLimitConfig(cfg)
}

func defaultInboundUpstreamRateLimitConfig() inboundUpstreamRateLimitConfig {
	return inboundUpstreamRateLimitConfig{
		Mode:                 defaultInboundUpstreamRateLimitMode,
		RedisGroup:           defaultInboundUpstreamRateLimitRedisGroup,
		KeyPrefix:            defaultInboundUpstreamRateLimitKeyPrefix,
		RedisFailureFallback: inboundRateLimitFallbackError,
		Default: inboundUpstreamRateLimitRule{
			Enabled: inboundBool(false),
		},
		Carriers: map[string]inboundUpstreamRateLimitRule{
			inboundCarrierTelecom: {
				Enabled:     inboundBool(true),
				QPS:         defaultInboundUpstreamQPS,
				Burst:       defaultInboundUpstreamBurst,
				RequestCost: defaultInboundUpstreamRequestCost,
			},
			inboundCarrierMobile: {
				Enabled: inboundBool(false),
			},
		},
	}
}

func normalizeInboundUpstreamRateLimitConfig(cfg inboundUpstreamRateLimitConfig) inboundUpstreamRateLimitConfig {
	cfg.Mode = normalizeInboundRateLimitMode(cfg.Mode)
	cfg.RedisGroup = firstNonEmpty(strings.TrimSpace(cfg.RedisGroup), defaultInboundUpstreamRateLimitRedisGroup)
	cfg.KeyPrefix = firstNonEmpty(strings.TrimSpace(cfg.KeyPrefix), defaultInboundUpstreamRateLimitKeyPrefix)
	cfg.RedisFailureFallback = normalizeInboundRateLimitFallback(cfg.RedisFailureFallback)
	if cfg.Default.Enabled == nil {
		cfg.Default.Enabled = inboundBool(false)
	}
	if len(cfg.Carriers) == 0 {
		cfg.Carriers = make(map[string]inboundUpstreamRateLimitRule, 2)
	} else {
		normalized := make(map[string]inboundUpstreamRateLimitRule, len(cfg.Carriers))
		for key, rule := range cfg.Carriers {
			normalized[normalizeInboundCarrier(key)] = rule
		}
		cfg.Carriers = normalized
	}
	telecomRule, ok := cfg.Carriers[inboundCarrierTelecom]
	if !ok {
		telecomRule = inboundUpstreamRateLimitRule{}
	}
	if telecomRule.Enabled == nil {
		telecomRule.Enabled = inboundBool(true)
	}
	if telecomRule.QPS <= 0 {
		telecomRule.QPS = defaultInboundUpstreamQPS
	}
	if telecomRule.Burst <= 0 {
		telecomRule.Burst = defaultInboundUpstreamBurst
	}
	if telecomRule.RequestCost <= 0 {
		telecomRule.RequestCost = defaultInboundUpstreamRequestCost
	}
	cfg.Carriers[inboundCarrierTelecom] = telecomRule

	mobileRule, ok := cfg.Carriers[inboundCarrierMobile]
	if !ok {
		mobileRule = inboundUpstreamRateLimitRule{}
	}
	if mobileRule.Enabled == nil {
		mobileRule.Enabled = inboundBool(false)
	}
	cfg.Carriers[inboundCarrierMobile] = mobileRule

	return cfg
}

func resolveInboundUpstreamRateLimitFromConfig(cfg inboundUpstreamRateLimitConfig, chanKey, tempKey string) inboundUpstreamRateLimitSetting {
	return resolveInboundUpstreamRateLimitFromChannelMetadata(cfg, chanKey, 0, tempKey)
}

func resolveInboundUpstreamRateLimitFromChannelMetadata(cfg inboundUpstreamRateLimitConfig, chanKey string, operator int, tempKey string) inboundUpstreamRateLimitSetting {
	cfg = normalizeInboundUpstreamRateLimitConfig(cfg)
	setting := inboundUpstreamRateLimitSetting{
		enabled:              true,
		mode:                 cfg.Mode,
		redisGroup:           cfg.RedisGroup,
		keyPrefix:            cfg.KeyPrefix,
		redisFailureFallback: cfg.RedisFailureFallback,
		qps:                  defaultInboundUpstreamQPS,
		burst:                defaultInboundUpstreamBurst,
		requestCost:          defaultInboundUpstreamRequestCost,
		tempKey:              strings.TrimSpace(tempKey),
	}
	if cfg.Enabled != nil {
		setting.enabled = *cfg.Enabled
	}
	applyInboundRateLimitRule(&setting, inboundUpstreamRateLimitRule{
		QPS:         cfg.QPS,
		Burst:       cfg.Burst,
		RequestCost: cfg.RequestCost,
	})
	applyInboundRateLimitRule(&setting, cfg.Default)
	carrier, known := inboundCarrierForOperator(operator)
	if model.IsTelecomTemplate(tempKey) {
		carrier, known = inboundCarrierTelecom, true
	}
	if known {
		setting.carrier = carrier
		if rule, ok := inboundCarrierRateLimitRule(cfg.Carriers, carrier); ok {
			applyInboundRateLimitRule(&setting, rule)
		}
	} else if operator == 0 {
		setting.skipReason = "missing_channel_operator"
	} else {
		setting.skipReason = "unsupported_channel_operator"
	}
	if rule, ok := inboundChannelRateLimitRule(cfg.Channels, chanKey); ok {
		applyInboundRateLimitRule(&setting, rule)
	}
	return normalizeInboundRateLimitSetting(setting)
}

func inboundUpstreamRateLimitGloballyEnabled(cfg inboundUpstreamRateLimitConfig) bool {
	cfg = normalizeInboundUpstreamRateLimitConfig(cfg)
	if cfg.Mode == "off" {
		return false
	}
	if cfg.Enabled == nil {
		return true
	}
	return *cfg.Enabled
}

func applyInboundRateLimitRule(setting *inboundUpstreamRateLimitSetting, rule inboundUpstreamRateLimitRule) {
	if setting == nil {
		return
	}
	if rule.Enabled != nil {
		setting.enabled = *rule.Enabled
	}
	if rule.QPS > 0 {
		setting.qps = rule.QPS
	}
	if rule.Burst > 0 {
		setting.burst = rule.Burst
	}
	if rule.RequestCost > 0 {
		setting.requestCost = rule.RequestCost
	}
}

func inboundChannelRateLimitRule(configs map[string]inboundUpstreamRateLimitRule, chanKey string) (inboundUpstreamRateLimitRule, bool) {
	chanKey = strings.TrimSpace(chanKey)
	if chanKey == "" || len(configs) == 0 {
		return inboundUpstreamRateLimitRule{}, false
	}
	if cfg, ok := configs[chanKey]; ok {
		return cfg, true
	}
	lower := strings.ToLower(chanKey)
	for key, cfg := range configs {
		if strings.ToLower(strings.TrimSpace(key)) == lower {
			return cfg, true
		}
	}
	return inboundUpstreamRateLimitRule{}, false
}

func inboundCarrierRateLimitRule(configs map[string]inboundUpstreamRateLimitRule, carrier string) (inboundUpstreamRateLimitRule, bool) {
	carrier = normalizeInboundCarrier(carrier)
	if carrier == "" || len(configs) == 0 {
		return inboundUpstreamRateLimitRule{}, false
	}
	if cfg, ok := configs[carrier]; ok {
		return cfg, true
	}
	for key, cfg := range configs {
		if normalizeInboundCarrier(key) == carrier {
			return cfg, true
		}
	}
	return inboundUpstreamRateLimitRule{}, false
}

func inboundCarrierForOperator(operator int) (string, bool) {
	switch operator {
	case 1:
		return inboundCarrierMobile, true
	case 2:
		return inboundCarrierTelecom, true
	default:
		return "", false
	}
}

func normalizeInboundCarrier(value string) string {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "", "unknown":
		return ""
	case inboundCarrierTelecom, "dx", "ctcc", "china_telecom", "chinatelecom":
		return inboundCarrierTelecom
	case inboundCarrierMobile, "yd", "cmcc", "china_mobile", "chinamobile":
		return inboundCarrierMobile
	default:
		return strings.ToLower(strings.TrimSpace(value))
	}
}

func normalizeInboundRateLimitSetting(setting inboundUpstreamRateLimitSetting) inboundUpstreamRateLimitSetting {
	if setting.qps <= 0 || math.IsNaN(setting.qps) || math.IsInf(setting.qps, 0) {
		setting.qps = defaultInboundUpstreamQPS
	}
	if setting.burst <= 0 {
		setting.burst = defaultInboundUpstreamBurst
	}
	if setting.requestCost <= 0 {
		setting.requestCost = defaultInboundUpstreamRequestCost
	}
	return setting
}

func normalizeInboundRateLimitMode(mode string) string {
	switch strings.ToLower(strings.TrimSpace(mode)) {
	case "", "redis", "distributed":
		return defaultInboundUpstreamRateLimitMode
	case "local", "memory":
		return "local"
	case "off", "disable", "disabled", "none":
		return "off"
	default:
		return defaultInboundUpstreamRateLimitMode
	}
}

func normalizeInboundRateLimitFallback(value string) string {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "", inboundRateLimitFallbackError:
		return inboundRateLimitFallbackError
	case inboundRateLimitFallbackLocal:
		return inboundRateLimitFallbackLocal
	default:
		return inboundRateLimitFallbackError
	}
}

func inboundUpstreamRateLimitKey(prefix, chanKey, carrier string) string {
	prefix = strings.TrimRight(strings.TrimSpace(prefix), ":")
	if prefix == "" {
		prefix = defaultInboundUpstreamRateLimitKeyPrefix
	}
	chanKey = strings.TrimLeft(strings.TrimSpace(chanKey), ":")
	carrier = strings.TrimLeft(normalizeInboundCarrier(carrier), ":")
	if carrier != "" {
		return prefix + ":carrier:" + carrier + ":chan:" + chanKey
	}
	return prefix + ":chan:" + chanKey
}

func inboundBool(value bool) *bool {
	return &value
}

type localTokenBucketLimiter struct {
	mu      sync.Mutex
	buckets map[string]*tokenBucket
}

func newLocalTokenBucketLimiter() *localTokenBucketLimiter {
	return &localTokenBucketLimiter{buckets: make(map[string]*tokenBucket)}
}

func (l *localTokenBucketLimiter) Wait(ctx context.Context, key string, qps float64, burst int) error {
	if qps <= 0 {
		return nil
	}
	bucket := l.bucket(key)
	delay := bucket.reserve(time.Now(), qps, burst)
	if delay <= 0 {
		return nil
	}
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func (l *localTokenBucketLimiter) bucket(key string) *tokenBucket {
	l.mu.Lock()
	defer l.mu.Unlock()
	bucket := l.buckets[key]
	if bucket == nil {
		bucket = &tokenBucket{}
		l.buckets[key] = bucket
	}
	return bucket
}

type tokenBucket struct {
	mu     sync.Mutex
	tokens float64
	last   time.Time
}

func (b *tokenBucket) reserve(now time.Time, qps float64, burst int) time.Duration {
	if qps <= 0 {
		return 0
	}
	capacity := float64(burst)
	if capacity <= 0 {
		capacity = defaultInboundUpstreamBurst
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.last.IsZero() {
		b.last = now
		b.tokens = capacity
	}
	if now.After(b.last) {
		elapsed := now.Sub(b.last).Seconds()
		b.tokens += elapsed * qps
		if b.tokens > capacity {
			b.tokens = capacity
		}
		b.last = now
	}
	if b.tokens >= 1 {
		b.tokens--
		return 0
	}
	waitSeconds := (1 - b.tokens) / qps
	wait := time.Duration(waitSeconds * float64(time.Second))
	if wait <= 0 {
		wait = time.Nanosecond
	}
	b.tokens = 0
	b.last = b.last.Add(wait)
	return b.last.Sub(now)
}

type redisTokenBucketLimiter struct {
	redis redisEvaler
}

const redisTokenBucketScript = `
local now = tonumber(ARGV[1])
local qps = tonumber(ARGV[2])
local burst = tonumber(ARGV[3])
local ttl = tonumber(ARGV[4])

if qps == nil or qps <= 0 then
  return {1, 0}
end
if burst == nil or burst <= 0 then
  burst = 1
end
if ttl == nil or ttl <= 0 then
  ttl = 60000
end

local values = redis.call("HMGET", KEYS[1], "tokens", "updated_at")
local tokens = tonumber(values[1])
local updated_at = tonumber(values[2])
local capacity = burst

if tokens == nil or updated_at == nil then
  tokens = capacity
  updated_at = now
end
if tokens > capacity then
  tokens = capacity
end
if now > updated_at then
  local elapsed = (now - updated_at) / 1000
  tokens = math.min(capacity, tokens + elapsed * qps)
  updated_at = now
end

local wait = 0
if tokens >= 1 then
  tokens = tokens - 1
else
  wait = math.ceil((1 - tokens) / qps * 1000)
  if wait < 1 then
    wait = 1
  end
  tokens = 0
  updated_at = updated_at + wait
  wait = updated_at - now
  if wait < 0 then
    wait = 0
  end
end

redis.call("HSET", KEYS[1], "tokens", tostring(tokens), "updated_at", tostring(updated_at))
redis.call("PEXPIRE", KEYS[1], ttl)
return {1, wait}
`

func (l *redisTokenBucketLimiter) Wait(ctx context.Context, key string, qps float64, burst int) error {
	wait, err := l.reserve(ctx, key, qps, burst)
	if err != nil {
		return err
	}
	if wait <= 0 {
		return nil
	}
	timer := time.NewTimer(wait)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func (l *redisTokenBucketLimiter) reserve(ctx context.Context, key string, qps float64, burst int) (time.Duration, error) {
	if l == nil || l.redis == nil || qps <= 0 {
		return 0, nil
	}
	if burst <= 0 {
		burst = defaultInboundUpstreamBurst
	}
	reply, err := l.redis.Eval(ctx, redisTokenBucketScript, 1, []string{key}, []any{
		time.Now().UnixMilli(),
		strconv.FormatFloat(qps, 'f', -1, 64),
		burst,
		rateLimitRedisKeyTTLMillis(qps, burst),
	})
	if err != nil {
		return 0, err
	}
	values := reply.Int64s()
	if len(values) < 2 || values[0] != 1 {
		return 0, fmt.Errorf("unexpected redis rate limiter reply: %v", reply.Val())
	}
	waitMillis := values[1]
	if waitMillis <= 0 {
		return 0, nil
	}
	return time.Duration(waitMillis) * time.Millisecond, nil
}

func rateLimitRedisKeyTTLMillis(qps float64, burst int) int64 {
	if qps <= 0 {
		return int64(time.Minute / time.Millisecond)
	}
	if burst <= 0 {
		burst = defaultInboundUpstreamBurst
	}
	idleMillis := int64(math.Ceil(float64(burst) / qps * 1000 * 2))
	minMillis := int64(time.Minute / time.Millisecond)
	if idleMillis < minMillis {
		return minMillis
	}
	return idleMillis
}