467 lines
17 KiB
Go
467 lines
17 KiB
Go
// Package giftlimit 为 gateway 送礼入口提供跨节点容量保护。
|
||
//
|
||
// 送礼请求的最坏下游成本与 gift_count*去重 target_count 成正比：幸运礼物需要逐份开奖，
|
||
// 每个目标也会形成独立账务和事实。因此本包不按 HTTP 次数计费，而按服务端已经校验过
|
||
// 的 total_gift_units 在 App、房间、用户三层扣令牌，并额外约束单用户在途请求数。
|
||
package giftlimit
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"errors"
|
||
"fmt"
|
||
"math"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
const (
|
||
defaultKeyPrefix = "gateway:gift:capacity:"
|
||
|
||
// 一个合法请求最多 5000 gift units，因此用户桶必须至少容纳一个完整请求；500/s
|
||
// 让极端 999 数量请求形成自然冷却，同时普通 300ms 连击批次不会感到逐击等待。
|
||
defaultUserBucketCapacity = 5_000
|
||
defaultUserRefillTokensPerSecond = 500
|
||
// 房间和 App 层允许多个真实用户同时连击，但把 5 人叠加造成的无界增长转换成 429
|
||
// 背压。数值均为可灰度配置的保护起点，不冒充未经压测的 Broker 容量结论。
|
||
defaultRoomBucketCapacity = 20_000
|
||
defaultRoomRefillTokensPerSecond = 5_000
|
||
defaultAppBucketCapacity = 100_000
|
||
defaultAppRefillTokensPerSecond = 25_000
|
||
|
||
defaultMaxInFlight = 2
|
||
defaultInFlightLease = 10 * time.Second
|
||
defaultRedisOperationTimeout = 500 * time.Millisecond
|
||
redisReleaseTimeout = time.Second
|
||
|
||
redisDialTimeout = time.Second
|
||
redisReadTimeout = 300 * time.Millisecond
|
||
redisWriteTimeout = 300 * time.Millisecond
|
||
redisPoolTimeout = 300 * time.Millisecond
|
||
redisMinRetryBackoff = 20 * time.Millisecond
|
||
redisMaxRetryBackoff = 50 * time.Millisecond
|
||
)
|
||
|
||
// Config 描述三层令牌桶、单用户在途租约和 Redis 操作预算。
|
||
type Config struct {
|
||
Enabled bool
|
||
KeyPrefix string
|
||
|
||
UserBucketCapacity int
|
||
UserRefillTokensPerSecond float64
|
||
RoomBucketCapacity int
|
||
RoomRefillTokensPerSecond float64
|
||
AppBucketCapacity int
|
||
AppRefillTokensPerSecond float64
|
||
|
||
MaxInFlight int
|
||
InFlightLease time.Duration
|
||
RedisOperationTimeout time.Duration
|
||
}
|
||
|
||
// Input 是一次送礼请求在容量维度上的稳定身份。
|
||
type Input struct {
|
||
AppCode string
|
||
UserID int64
|
||
RoomID string
|
||
RequestID string
|
||
Weight int
|
||
}
|
||
|
||
// Lease 标识已占用的分布式在途槽。字段只供 Limiter.Release 使用。
|
||
type Lease struct {
|
||
inFlightKey string
|
||
member string
|
||
}
|
||
|
||
// Decision 返回是否获准以及客户端下一次可重试的最短等待时间。
|
||
type Decision struct {
|
||
Allowed bool
|
||
RetryAfter time.Duration
|
||
Lease Lease
|
||
}
|
||
|
||
// Limiter 是 room HTTP adapter 所需的最小容量保护边界。
|
||
type Limiter interface {
|
||
Acquire(ctx context.Context, input Input) (Decision, error)
|
||
Release(ctx context.Context, lease Lease) error
|
||
}
|
||
|
||
type redisScriptRunner interface {
|
||
evalSlice(ctx context.Context, script string, keys []string, args ...any) ([]any, error)
|
||
}
|
||
|
||
type goRedisScriptRunner struct {
|
||
client *redis.Client
|
||
}
|
||
|
||
type bucketPolicy struct {
|
||
capacity int
|
||
refillPerSecond float64
|
||
}
|
||
|
||
type redisScopeKeys struct {
|
||
appBucket string
|
||
roomBucket string
|
||
userBucket string
|
||
inFlight string
|
||
leaseMember string
|
||
}
|
||
|
||
// RedisLimiter 按 user→room→app 顺序执行三个独立 Lua。
|
||
//
|
||
// Redis Cluster 只允许一次脚本访问同一 hash slot。App、房间、用户是不同聚合域，不能
|
||
// 伪装成一个跨 slot 原子事务。用户层先检查 in-flight，可防止已占满并发槽的账号用必拒
|
||
// 请求持续消耗房间/App 额度；后层拒绝时释放用户 lease，但保留用户令牌作为请求方成本。
|
||
// App 拒绝时不回滚房间令牌，少量保守损耗比并发回滚造成超发或绕过容量上限更安全。
|
||
type RedisLimiter struct {
|
||
runner redisScriptRunner
|
||
config Config
|
||
}
|
||
|
||
// redisTokenBucketScript 在一个聚合维度内原子补充并扣减令牌，返回
|
||
// {allowed, retry_after_ms}。Redis TIME 是多 gateway 节点的统一时间源。
|
||
const redisTokenBucketScript = `
|
||
local redis_time = redis.call("TIME")
|
||
local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000)
|
||
local capacity = tonumber(ARGV[1])
|
||
local refill_per_ms = tonumber(ARGV[2])
|
||
local weight = tonumber(ARGV[3])
|
||
local bucket_ttl_ms = tonumber(ARGV[4])
|
||
|
||
local tokens = tonumber(redis.call("HGET", KEYS[1], "tokens"))
|
||
local last_ms = tonumber(redis.call("HGET", KEYS[1], "last_ms"))
|
||
if not tokens or not last_ms then
|
||
tokens = capacity
|
||
last_ms = now_ms
|
||
else
|
||
tokens = math.min(capacity, tokens + math.max(0, now_ms - last_ms) * refill_per_ms)
|
||
end
|
||
|
||
if tokens < weight then
|
||
redis.call("HSET", KEYS[1], "tokens", tokens, "last_ms", now_ms)
|
||
redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms)
|
||
local retry_ms = math.ceil((weight - tokens) / refill_per_ms)
|
||
return {0, math.max(1, retry_ms)}
|
||
end
|
||
|
||
redis.call("HSET", KEYS[1], "tokens", tokens - weight, "last_ms", now_ms)
|
||
redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms)
|
||
return {1, 0}
|
||
`
|
||
|
||
// redisUserAcquireScript 把用户桶和 app:user 在途 ZSET 放进同一个 hash slot，因此用户
|
||
// 额度检查、并发检查和租约写入在单个 Redis 节点内保持原子。
|
||
const redisUserAcquireScript = `
|
||
local redis_time = redis.call("TIME")
|
||
local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000)
|
||
local capacity = tonumber(ARGV[1])
|
||
local refill_per_ms = tonumber(ARGV[2])
|
||
local weight = tonumber(ARGV[3])
|
||
local max_in_flight = tonumber(ARGV[4])
|
||
local lease_ms = tonumber(ARGV[5])
|
||
local member = ARGV[6]
|
||
local bucket_ttl_ms = tonumber(ARGV[7])
|
||
|
||
-- gateway 异常退出不会执行 Release；每次申请前清理过期租约，确保 lease_ms 后自动恢复。
|
||
redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", now_ms)
|
||
local in_flight = redis.call("ZCARD", KEYS[2])
|
||
if in_flight >= max_in_flight then
|
||
local earliest = redis.call("ZRANGE", KEYS[2], 0, 0, "WITHSCORES")
|
||
local retry_ms = 1
|
||
if earliest[2] then
|
||
retry_ms = math.max(1, tonumber(earliest[2]) - now_ms)
|
||
end
|
||
return {0, retry_ms}
|
||
end
|
||
|
||
local tokens = tonumber(redis.call("HGET", KEYS[1], "tokens"))
|
||
local last_ms = tonumber(redis.call("HGET", KEYS[1], "last_ms"))
|
||
if not tokens or not last_ms then
|
||
tokens = capacity
|
||
last_ms = now_ms
|
||
else
|
||
tokens = math.min(capacity, tokens + math.max(0, now_ms - last_ms) * refill_per_ms)
|
||
end
|
||
|
||
if tokens < weight then
|
||
redis.call("HSET", KEYS[1], "tokens", tokens, "last_ms", now_ms)
|
||
redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms)
|
||
local retry_ms = math.ceil((weight - tokens) / refill_per_ms)
|
||
return {0, math.max(1, retry_ms)}
|
||
end
|
||
|
||
redis.call("HSET", KEYS[1], "tokens", tokens - weight, "last_ms", now_ms)
|
||
redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms)
|
||
redis.call("ZADD", KEYS[2], now_ms + lease_ms, member)
|
||
redis.call("PEXPIRE", KEYS[2], lease_ms * 2)
|
||
return {1, 0}
|
||
`
|
||
|
||
const redisReleaseScript = `
|
||
redis.call("ZREM", KEYS[1], ARGV[1])
|
||
if redis.call("ZCARD", KEYS[1]) == 0 then
|
||
redis.call("DEL", KEYS[1])
|
||
end
|
||
return {1}
|
||
`
|
||
|
||
// Normalize 收敛与 10k MQ TPS 和当前送礼行为匹配的安全默认值。
|
||
func Normalize(config Config) Config {
|
||
config.KeyPrefix = strings.TrimSpace(config.KeyPrefix)
|
||
if config.KeyPrefix == "" {
|
||
config.KeyPrefix = defaultKeyPrefix
|
||
}
|
||
if config.UserBucketCapacity <= 0 {
|
||
config.UserBucketCapacity = defaultUserBucketCapacity
|
||
}
|
||
if config.UserRefillTokensPerSecond <= 0 {
|
||
config.UserRefillTokensPerSecond = defaultUserRefillTokensPerSecond
|
||
}
|
||
if config.RoomBucketCapacity <= 0 {
|
||
config.RoomBucketCapacity = defaultRoomBucketCapacity
|
||
}
|
||
if config.RoomRefillTokensPerSecond <= 0 {
|
||
config.RoomRefillTokensPerSecond = defaultRoomRefillTokensPerSecond
|
||
}
|
||
if config.AppBucketCapacity <= 0 {
|
||
config.AppBucketCapacity = defaultAppBucketCapacity
|
||
}
|
||
if config.AppRefillTokensPerSecond <= 0 {
|
||
config.AppRefillTokensPerSecond = defaultAppRefillTokensPerSecond
|
||
}
|
||
if config.MaxInFlight <= 0 {
|
||
config.MaxInFlight = defaultMaxInFlight
|
||
}
|
||
if config.InFlightLease <= 0 {
|
||
config.InFlightLease = defaultInFlightLease
|
||
}
|
||
if config.RedisOperationTimeout <= 0 {
|
||
config.RedisOperationTimeout = defaultRedisOperationTimeout
|
||
}
|
||
return config
|
||
}
|
||
|
||
// NewRedisClient 建立并探测容量保护使用的 Redis 连接。
|
||
// ContextTimeoutEnabled 必须开启，否则 go-redis 可能只依赖较长 socket timeout，使 handler
|
||
// 的 500ms fail-closed 预算和 Release 的 1s 回收预算失效。
|
||
func NewRedisClient(ctx context.Context, addr string, password string, db int) (*redis.Client, error) {
|
||
client := redis.NewClient(redisClientOptions(addr, password, db))
|
||
if err := client.Ping(ctx).Err(); err != nil {
|
||
_ = client.Close()
|
||
return nil, err
|
||
}
|
||
return client, nil
|
||
}
|
||
|
||
func redisClientOptions(addr string, password string, db int) *redis.Options {
|
||
return &redis.Options{
|
||
Addr: strings.TrimSpace(addr),
|
||
Password: password,
|
||
DB: db,
|
||
MaxRetries: 1,
|
||
MinRetryBackoff: redisMinRetryBackoff,
|
||
MaxRetryBackoff: redisMaxRetryBackoff,
|
||
DialTimeout: redisDialTimeout,
|
||
DialerRetries: 1,
|
||
ReadTimeout: redisReadTimeout,
|
||
WriteTimeout: redisWriteTimeout,
|
||
PoolTimeout: redisPoolTimeout,
|
||
ContextTimeoutEnabled: true,
|
||
}
|
||
}
|
||
|
||
// NewRedisLimiter 创建生产使用的分布式 limiter。
|
||
func NewRedisLimiter(client *redis.Client, config Config) *RedisLimiter {
|
||
return newRedisLimiter(goRedisScriptRunner{client: client}, config)
|
||
}
|
||
|
||
func newRedisLimiter(runner redisScriptRunner, config Config) *RedisLimiter {
|
||
return &RedisLimiter{runner: runner, config: Normalize(config)}
|
||
}
|
||
|
||
// Acquire 在一个独立短 deadline 内按 user→room→app 顺序保守扣减。
|
||
// 用户并发先行可挡住低成本全局 DoS；后层失败会立即释放用户 lease，但不会返还已扣令牌。
|
||
func (l *RedisLimiter) Acquire(ctx context.Context, input Input) (Decision, error) {
|
||
if l == nil || l.runner == nil {
|
||
return Decision{}, errors.New("gift capacity limiter is not configured")
|
||
}
|
||
appCode := strings.ToLower(strings.TrimSpace(input.AppCode))
|
||
if appCode == "" || input.UserID <= 0 || strings.TrimSpace(input.RoomID) == "" || strings.TrimSpace(input.RequestID) == "" || input.Weight <= 0 {
|
||
return Decision{}, errors.New("gift capacity limiter input is invalid")
|
||
}
|
||
config := Normalize(l.config)
|
||
if input.Weight > config.UserBucketCapacity || input.Weight > config.RoomBucketCapacity || input.Weight > config.AppBucketCapacity {
|
||
// handler 会先用 max_targets_per_request 返回 4xx；到达这里说明容量配置与协议上限漂移，
|
||
// 必须 fail-closed，而不是返回一个永远无法成功的 429 重试循环。
|
||
return Decision{}, errors.New("gift target weight exceeds configured bucket capacity")
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
operationCtx, cancel := context.WithTimeout(ctx, config.RedisOperationTimeout)
|
||
defer cancel()
|
||
|
||
keys := buildRedisScopeKeys(config.KeyPrefix, appCode, input)
|
||
userPolicy := bucketPolicy{capacity: config.UserBucketCapacity, refillPerSecond: config.UserRefillTokensPerSecond}
|
||
result, err := l.runner.evalSlice(operationCtx, redisUserAcquireScript, []string{keys.userBucket, keys.inFlight},
|
||
userPolicy.capacity,
|
||
refillPerMillisecond(userPolicy.refillPerSecond),
|
||
input.Weight,
|
||
config.MaxInFlight,
|
||
config.InFlightLease.Milliseconds(),
|
||
keys.leaseMember,
|
||
bucketTTL(userPolicy, config.InFlightLease).Milliseconds(),
|
||
)
|
||
if err != nil {
|
||
return Decision{}, err
|
||
}
|
||
userDecision, err := decodeDecision(result)
|
||
if err != nil {
|
||
// Lua 已返回但协议解码失败时执行幂等 ZREM；脚本可能已经写入 lease，不能等到 TTL 才恢复。
|
||
l.releaseAcquiredLease(operationCtx, Lease{inFlightKey: keys.inFlight, member: keys.leaseMember})
|
||
return Decision{}, err
|
||
}
|
||
if !userDecision.Allowed {
|
||
// 并发已满或用户桶不足时尚未触达 room/app key，因此攻击者不能靠必拒请求耗尽共享额度。
|
||
return userDecision, nil
|
||
}
|
||
lease := Lease{inFlightKey: keys.inFlight, member: keys.leaseMember}
|
||
|
||
layers := []struct {
|
||
key string
|
||
policy bucketPolicy
|
||
}{
|
||
{key: keys.roomBucket, policy: bucketPolicy{capacity: config.RoomBucketCapacity, refillPerSecond: config.RoomRefillTokensPerSecond}},
|
||
{key: keys.appBucket, policy: bucketPolicy{capacity: config.AppBucketCapacity, refillPerSecond: config.AppRefillTokensPerSecond}},
|
||
}
|
||
for _, layer := range layers {
|
||
decision, err := l.consumeBucket(operationCtx, layer.key, layer.policy, input.Weight)
|
||
if err != nil || !decision.Allowed {
|
||
// 使用同一个短 deadline 立即释放用户 lease；deadline 已耗尽时释放可能失败，
|
||
// 但 ZSET score 的 InFlightLease 仍会兜底回收，绝不会永久锁死用户。
|
||
l.releaseAcquiredLease(operationCtx, lease)
|
||
return decision, err
|
||
}
|
||
}
|
||
userDecision.Lease = lease
|
||
return userDecision, nil
|
||
}
|
||
|
||
func (l *RedisLimiter) consumeBucket(ctx context.Context, key string, policy bucketPolicy, weight int) (Decision, error) {
|
||
result, err := l.runner.evalSlice(ctx, redisTokenBucketScript, []string{key},
|
||
policy.capacity,
|
||
refillPerMillisecond(policy.refillPerSecond),
|
||
weight,
|
||
bucketTTL(policy, 0).Milliseconds(),
|
||
)
|
||
if err != nil {
|
||
return Decision{}, err
|
||
}
|
||
return decodeDecision(result)
|
||
}
|
||
|
||
func (l *RedisLimiter) releaseAcquiredLease(ctx context.Context, lease Lease) {
|
||
if lease.inFlightKey == "" || lease.member == "" {
|
||
return
|
||
}
|
||
// 不调用公开 Release，避免给一次 Acquire 额外再增加 1s；补偿动作必须服从 Acquire 的
|
||
// 共同短预算。失败时租约 TTL 是安全兜底，用户 token 则有意不返还。
|
||
_, _ = l.runner.evalSlice(ctx, redisReleaseScript, []string{lease.inFlightKey}, lease.member)
|
||
}
|
||
|
||
// Release 始终再包一层 1s deadline。handler 使用与客户端取消解耦的 background context，
|
||
// go-redis 又启用 ContextTimeoutEnabled，因此慢 Redis 不会把请求收尾永久挂住。
|
||
func (l *RedisLimiter) Release(ctx context.Context, lease Lease) error {
|
||
if l == nil || l.runner == nil || lease.inFlightKey == "" || lease.member == "" {
|
||
return nil
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
releaseCtx, cancel := context.WithTimeout(ctx, redisReleaseTimeout)
|
||
defer cancel()
|
||
_, err := l.runner.evalSlice(releaseCtx, redisReleaseScript, []string{lease.inFlightKey}, lease.member)
|
||
return err
|
||
}
|
||
|
||
func (r goRedisScriptRunner) evalSlice(ctx context.Context, script string, keys []string, args ...any) ([]any, error) {
|
||
return r.client.Eval(ctx, script, keys, args...).Slice()
|
||
}
|
||
|
||
func buildRedisScopeKeys(prefix string, appCode string, input Input) redisScopeKeys {
|
||
appPart := boundedHash(appCode)
|
||
roomPart := boundedHash(appCode + ":" + strings.TrimSpace(input.RoomID))
|
||
userPart := boundedHash(appCode + ":" + strconv.FormatInt(input.UserID, 10))
|
||
member := boundedHash(input.RequestID)
|
||
// 每个聚合域拥有自己的 hash tag；只有用户桶与 in-flight ZSET 共享 user tag，因而只有
|
||
// 这两个 key 会放入同一 Lua。App/room/user 三层绝不跨 slot 伪原子。
|
||
return redisScopeKeys{
|
||
appBucket: prefix + "{app:" + appPart + "}:bucket",
|
||
roomBucket: prefix + "{room:" + roomPart + "}:bucket",
|
||
userBucket: prefix + "{user:" + userPart + "}:bucket",
|
||
inFlight: prefix + "{user:" + userPart + "}:inflight",
|
||
leaseMember: member,
|
||
}
|
||
}
|
||
|
||
func bucketTTL(policy bucketPolicy, minimum time.Duration) time.Duration {
|
||
ttl := time.Duration(math.Ceil(float64(policy.capacity)/policy.refillPerSecond*2)) * time.Second
|
||
if ttl < time.Second {
|
||
ttl = time.Second
|
||
}
|
||
if ttl < minimum {
|
||
ttl = minimum
|
||
}
|
||
return ttl
|
||
}
|
||
|
||
func refillPerMillisecond(tokensPerSecond float64) string {
|
||
return strconv.FormatFloat(tokensPerSecond/1000, 'f', -1, 64)
|
||
}
|
||
|
||
func decodeDecision(result []any) (Decision, error) {
|
||
if len(result) != 2 {
|
||
return Decision{}, fmt.Errorf("gift capacity limiter returned %d values", len(result))
|
||
}
|
||
allowed, err := redisInt64(result[0])
|
||
if err != nil {
|
||
return Decision{}, fmt.Errorf("decode gift capacity allowed: %w", err)
|
||
}
|
||
retryAfterMS, err := redisInt64(result[1])
|
||
if err != nil {
|
||
return Decision{}, fmt.Errorf("decode gift capacity retry_after_ms: %w", err)
|
||
}
|
||
if allowed == 1 {
|
||
return Decision{Allowed: true}, nil
|
||
}
|
||
if retryAfterMS < 1 {
|
||
retryAfterMS = 1
|
||
}
|
||
return Decision{RetryAfter: time.Duration(retryAfterMS) * time.Millisecond}, nil
|
||
}
|
||
|
||
func boundedHash(value string) string {
|
||
sum := sha256.Sum256([]byte(strings.TrimSpace(value)))
|
||
return hex.EncodeToString(sum[:16])
|
||
}
|
||
|
||
func redisInt64(value any) (int64, error) {
|
||
switch typed := value.(type) {
|
||
case int64:
|
||
return typed, nil
|
||
case int:
|
||
return int64(typed), nil
|
||
case string:
|
||
return strconv.ParseInt(typed, 10, 64)
|
||
case []byte:
|
||
return strconv.ParseInt(string(typed), 10, 64)
|
||
default:
|
||
return 0, fmt.Errorf("unsupported redis integer type %T", value)
|
||
}
|
||
}
|