2fe7c1f8ac
- 四个轮询处理器拆为步骤函数 Step（返回结果分类），旧 Asynq Handle 只做并发控制与原有重入队映射，行为不变；carrier 模式下旧 handler 直接结束
- 网关限流增加等待模式（ctx 标记），NSQ 消费者等名额到单步超时，超时按限流处理不计失败；记录限流等待次数
- 卡任务执行器：核对 token 并抢执行权，按配置步骤顺序串行执行到期检查（单步超时、总超时、panic 恢复、流量已触发评估时跳过套餐），执行中续期并 Touch，续期失败立即停止；收尾写回到期、连续失败、退避、延迟确认与告警
- 分发器：每秒按账户抢分发锁领取（执行上限按消费能力收紧、优先占比、按限流组剩余 QPS 估算），发布到 polling.card.{账户}，发布失败释放；消费旧优先提示与手动触发；心跳与处理中标记到期计数
- NSQ 生产者与按账户动态订阅的消费者管理器，优雅退出宽限期
- 失败告警异步写 tb_polling_failure_alert（同卡同类型一条未处理，成功转已恢复），账户窗口内告警数超阈值产生账户级告警
- 模式监听：切 carrier 时旧分片调度停止出队，加锁初始化新结构后开始领取并订阅；回 legacy 时停止消费、重建旧分片队列；生命周期（API 与 Worker）在 carrier 模式写账户队列
239 lines
7.8 KiB
Go
239 lines
7.8 KiB
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)
|
||
}
|
||
|
||
// Budget 返回运营商账户各查询类操作所属限流组本秒剩余名额，供分发器按预计调用数领取。
|
||
// groups 为限流组 → 本秒剩余；opGroup 为操作 → 限流组。只统计模式下返回 limited=false（不按 QPS 限制领取）。
|
||
// 全局组只在账户组之外再收紧：剩余名额取两者较小值。
|
||
func (l *Limiter) Budget(ctx context.Context, carrierID uint, ops []string) (groups map[string]int, opGroup map[string]string, limited bool) {
|
||
if l == nil || l.redis == nil {
|
||
return nil, nil, false
|
||
}
|
||
settings := l.settings.Current(ctx)
|
||
if !settings.Enforce() {
|
||
return nil, nil, false
|
||
}
|
||
groups, opGroup = map[string]int{}, map[string]string{}
|
||
now := time.Now().UnixMilli()
|
||
since := now - constants.GatewayRateLimitWindow.Milliseconds()
|
||
for _, op := range ops {
|
||
account, accountOK, global, globalOK := l.groups.Resolve(ctx, carrierID, op)
|
||
if !accountOK {
|
||
account = GroupRule{GroupCode: defaultGroupCode + ":" + op, MaxQPS: settings.DefaultQPS}
|
||
}
|
||
code := account.GroupCode
|
||
opGroup[op] = code
|
||
if _, seen := groups[code]; seen {
|
||
continue
|
||
}
|
||
limit := account.MaxQPS
|
||
if !account.HighPriority {
|
||
limit -= account.ReservedQPS
|
||
}
|
||
used, err := l.redis.ZCount(ctx, constants.RedisGatewayQPSKey(carrierID, code), strconv.FormatInt(since, 10), "+inf").Result()
|
||
if err != nil {
|
||
return nil, nil, false
|
||
}
|
||
remaining := limit - int(used)
|
||
if globalOK {
|
||
globalUsed, gerr := l.redis.ZCount(ctx, constants.RedisGatewayGlobalQPSKey(global.GroupCode), strconv.FormatInt(since, 10), "+inf").Result()
|
||
if gerr == nil && global.MaxQPS-int(globalUsed) < remaining {
|
||
remaining = global.MaxQPS - int(globalUsed)
|
||
}
|
||
}
|
||
if remaining < 0 {
|
||
remaining = 0
|
||
}
|
||
groups[code] = remaining
|
||
}
|
||
return groups, opGroup, true
|
||
}
|