junhong_cmp_fiber/internal/infrastructure/gatewaylimit/limiter.go

package gatewaylimit

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"strconv"
	"time"

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

	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// slidingWindowScript 在一次原子执行内判断账户组窗口与可选的全局组窗口。
//
// KEYS[1] 账户组窗口 ZSET；KEYS[2] 账户组高优先级等待计数；KEYS[3] 全局组窗口 ZSET（无全局组时传空串占位）。
// ARGV：now_ms、window_ms、max_qps、reserved_qps、is_high(0/1)、global_max(0 表示无全局组)、member、dry_run(0/1)。
// 返回 {allowed(1/0), wait_ms}：dry_run=1（只统计）时总是记录本次请求，allowed=0 表示"本应被限流"。
const slidingWindowLua = `
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local maxQPS = tonumber(ARGV[3])
local reserved = tonumber(ARGV[4])
local isHigh = tonumber(ARGV[5])
local globalMax = tonumber(ARGV[6])
local member = ARGV[7]
local dryRun = tonumber(ARGV[8])
local useGlobal = globalMax > 0 and KEYS[3] ~= ''

redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - window)
if useGlobal then
  redis.call('ZREMRANGEBYSCORE', KEYS[3], '-inf', now - window)
end

local limit = maxQPS
if isHigh == 0 then
  limit = maxQPS - reserved
  if reserved > 0 then
    local waiting = tonumber(redis.call('GET', KEYS[2]) or '0')
    if waiting > 0 then
      limit = 0
    end
  end
end

local allowed = 1
local wait = 0
local count = redis.call('ZCARD', KEYS[1])
if count >= limit then
  allowed = 0
  local oldest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
  if oldest[2] then
    wait = tonumber(oldest[2]) + window - now
  else
    wait = window
  end
end
if allowed == 1 and useGlobal then
  local gcount = redis.call('ZCARD', KEYS[3])
  if gcount >= globalMax then
    allowed = 0
    local goldest = redis.call('ZRANGE', KEYS[3], 0, 0, 'WITHSCORES')
    if goldest[2] then
      wait = tonumber(goldest[2]) + window - now
    else
      wait = window
    end
  end
end

if allowed == 1 or dryRun == 1 then
  redis.call('ZADD', KEYS[1], now, member)
  redis.call('PEXPIRE', KEYS[1], window * 2)
  if useGlobal then
    redis.call('ZADD', KEYS[3], now, member)
    redis.call('PEXPIRE', KEYS[3], window * 2)
  end
end
if wait < 1 and allowed == 0 then
  wait = 1
end
return {allowed, wait}
`

var slidingWindowScript = redis.NewScript(slidingWindowLua)

// Decision 是一次限流判断的结果。
type Decision struct {
	Allowed bool          // 是否放行（只统计模式下恒为 true）
	Exceed  bool          // 是否超过上限（只统计模式下表示"本应被限流"）
	Wait    time.Duration // 拦截时建议的等待时长
}

// Limiter 是基于 Redis 滑动窗口的运营商账户限流器，跨全部进程共享。
type Limiter struct {
	redis    *redis.Client
	groups   *GroupProvider
	settings *SettingsProvider
}

// NewLimiter 创建限流器。
func NewLimiter(redisClient *redis.Client, groups *GroupProvider, settings *SettingsProvider) *Limiter {
	return &Limiter{redis: redisClient, groups: groups, settings: settings}
}

// Check 立即判断一次请求能否放行，不等待。拦截模式下超过上限返回 Allowed=false；
// 只统计模式下总是放行，Exceed 表示本应被限流。Redis 故障时放行（失败开放），避免限流故障阻断业务。
func (l *Limiter) Check(ctx context.Context, carrierID uint, op string) (Decision, error) {
	settings := l.settings.Current(ctx)
	account, accountOK, global, globalOK := l.groups.Resolve(ctx, carrierID, op)
	if !accountOK {
		account = GroupRule{GroupCode: defaultGroupCode + ":" + op, MaxQPS: settings.DefaultQPS}
	}
	globalKey, globalMax := "", 0
	if globalOK {
		globalKey, globalMax = constants.RedisGatewayGlobalQPSKey(global.GroupCode), global.MaxQPS
	}
	isHigh, dryRun := 0, 1
	if account.HighPriority {
		isHigh = 1
	}
	if settings.Enforce() {
		dryRun = 0
	}
	now := time.Now().UnixMilli()
	result, err := slidingWindowScript.Run(ctx, l.redis,
		[]string{constants.RedisGatewayQPSKey(carrierID, account.GroupCode), constants.RedisGatewayQPSHighWaitingKey(carrierID, account.GroupCode), globalKey},
		now, constants.GatewayRateLimitWindow.Milliseconds(), account.MaxQPS, account.ReservedQPS, isHigh, globalMax, windowMember(now), dryRun,
	).Int64Slice()
	if err != nil || len(result) != 2 {
		return Decision{Allowed: true}, err
	}
	exceed := result[0] == 0
	decision := Decision{Allowed: !exceed || dryRun == 1, Exceed: exceed}
	if exceed && dryRun == 0 {
		decision.Wait = time.Duration(result[1]) * time.Millisecond
	}
	return decision, nil
}

// Wait 在拦截模式下按建议时长等待名额，总等待不超过 maxWait；超时返回 false。
// 高优先级操作等待期间登记等待计数，使同组普通操作拿不到新名额。本 Change 中 Asynq 任务不使用该方法。
func (l *Limiter) Wait(ctx context.Context, carrierID uint, op string, maxWait time.Duration) (bool, error) {
	deadline := time.Now().Add(maxWait)
	registered := false
	var waitingKey string
	defer func() {
		if registered {
			l.redis.Decr(context.WithoutCancel(ctx), waitingKey)
		}
	}()
	for {
		decision, err := l.Check(ctx, carrierID, op)
		if err != nil || decision.Allowed {
			return true, err
		}
		if !registered {
			account, accountOK, _, _ := l.groups.Resolve(ctx, carrierID, op)
			if accountOK && account.HighPriority && account.ReservedQPS > 0 {
				waitingKey = constants.RedisGatewayQPSHighWaitingKey(carrierID, account.GroupCode)
				pipe := l.redis.TxPipeline()
				pipe.Incr(ctx, waitingKey)
				pipe.Expire(ctx, waitingKey, 10*time.Second)
				if _, err := pipe.Exec(ctx); err == nil {
					registered = true
				}
			}
		}
		sleep := decision.Wait
		if remaining := time.Until(deadline); remaining <= 0 {
			return false, nil
		} else if sleep > remaining {
			sleep = remaining
		}
		timer := time.NewTimer(sleep)
		select {
		case <-ctx.Done():
			timer.Stop()
			return false, ctx.Err()
		case <-timer.C:
		}
	}
}

// windowMember 生成窗口内唯一成员，避免同毫秒请求互相覆盖。
func windowMember(nowMS int64) string {
	buf := make([]byte, 6)
	_, _ = rand.Read(buf)
	return strconv.FormatInt(nowMS, 10) + "-" + hex.EncodeToString(buf)
}