Merge pull request #3762 from jianjianai/fix/add-chinese-comments

优化 Redis SCAN 清理架构
This commit is contained in:
Wesley Liddick
2026-07-07 09:36:20 +08:00
committed by GitHub
8 changed files with 1387 additions and 222 deletions
+568
View File
@@ -0,0 +1,568 @@
# Redis SCAN 架构优化执行文档
本文是架构级执行文档,只覆盖三个目标:
1. 账号并发活跃负载查询不再通过 Redis keyspace `SCAN` 发现账号。
2. 账号/用户并发槽过期清理、启动遗留槽清理不再通过 Redis keyspace `SCAN` 发现 key。
3. 用户消息队列孤儿锁清理不再通过 Redis keyspace `SCAN` 发现 lock key。
不覆盖旁路录制、OpenAI failover、日志量、业务限流配置调参。不要把本文扩展成短期止血方案。
## 成功标准
实现完成后必须同时满足:
- `backend/internal/repository/concurrency_cache.go` 中不得再调用 `c.rdb.Scan(...)`。
- `backend/internal/repository/user_msg_queue_cache.go` 中不得再调用 `c.rdb.Scan(...)`。
- `backend/internal/service/user_msg_queue_service.go` 中不得再出现 `ScanLockKeys` 接口调用。
- `GetActiveAccountLoadMap` 只读显式维护的 Redis 索引,不扫描 Redis keyspace。
- `CleanupExpiredAccountSlotKeys` 只处理显式索引中的候选账号,不扫描 Redis keyspace。
- `CleanupStaleProcessSlots` 不扫描 Redis keyspace;它必须基于显式索引清理,或只依赖 TTL/score 自然过期。
- UMQ cleanup worker 只读 `umq:lock:index` 候选项,不扫描 `umq:{*}:lock`。
- 主业务并发限制仍以原账号/用户 slot key 为准,不能以索引为准。索引只能用于发现候选对象、监控和清理。
执行完必须用下面命令确认生产代码没有遗留扫描:
```powershell
rg -n "\.Scan\(" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go
rg -n "ScanLockKeys|scanAccountIDsByPrefix|cleanupSlotsByPattern|deleteKeysByPattern|umqScanPattern" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go
```
第一条必须无输出。第二条必须无生产函数残留;测试文件不在此检查范围。
## 不可违反的约束
- 不要用 `KEYS` 替代 `SCAN`。
- 不要把全量 Redis keyspace 扫描移动到另一个函数、goroutine、启动流程或管理接口里。
- 不要在请求路径、后台 worker、启动流程中做 Redis keyspace pattern enumeration。
- 不要在 Redis Lua 脚本里同时操作“全局索引 key”和“账号/用户局部 key”。项目代码当前有 Redis Cluster 兼容要求,这种写法会在 Cluster 下触发 CROSSSLOT。
- 索引更新失败不得改变主业务 acquire/release 的成功结果。索引是 best-effort discovery structure,不是并发正确性的来源。
- 不能因为索引缺失而拒绝用户请求。索引缺失最多影响 Ops 实时视图和后台提前清理;原 slot/wait key 的 TTL 必须保证最终自愈。
## 新增 Redis Key
### 并发索引
在 `backend/internal/repository/concurrency_cache.go` 增加常量:
```go
const (
accountActiveIndexKey = "concurrency:account:active_index" // ZSET member=accountID, score=expireAtUnixSeconds
userActiveIndexKey = "concurrency:user:active_index" // ZSET member=userID, score=expireAtUnixSeconds
activeIndexCleanupBatchSize = 1000
activeIndexPipelineChunkSize = 500
)
```
语义:
- `accountActiveIndexKey` 记录“可能有账号槽位或账号等待计数”的账号 ID。
- `userActiveIndexKey` 记录“可能有用户槽位或用户等待计数”的用户 ID。
- ZSET score 是候选对象的保守过期时间,单位为 Unix 秒。
- member 必须是十进制 ID 字符串,不要存完整 Redis key。
- 索引允许短暂 stale;读索引后必须二次查询真实 slot/wait key。
score 规则:
- 成功获取账号槽位:score = Redis 当前秒 + `slotTTLSeconds`。
- 成功增加账号等待计数:score = Redis 当前秒 + `waitQueueTTLSeconds`。
- 成功获取用户槽位:score = Redis 当前秒 + `slotTTLSeconds`。
- 成功增加用户等待计数:score = Redis 当前秒 + `waitQueueTTLSeconds`。
- release/decrement 后如果真实 slot count 和 wait count 都为 0,则从索引 `ZREM`。
- release/decrement 后如果仍有 slot 或 wait,则重新 `ZADD` 一个新的保守过期时间。
### UMQ 锁索引
在 `backend/internal/repository/user_msg_queue_cache.go` 增加常量:
```go
const (
umqLockIndexKey = "umq:lock:index" // ZSET member=accountID, score=lockExpireAtUnixMs
umqLockIndexCleanupBatchSize = 1000
)
```
语义:
- `umqLockIndexKey` 记录“可能存在 UMQ lock”的账号 ID。
- ZSET score 是 lock 的预计过期时间,单位为 Unix 毫秒。
- member 必须是十进制 accountID 字符串。
- 索引只用于 cleanup worker 找候选 lock。锁是否存在、是否孤儿,必须再查 `umq:{accountID}:lock`。
## 第一部分:并发活跃索引
修改文件:`backend/internal/repository/concurrency_cache.go`。
### 1.1 增加 Redis 时间 helper
新增 helper,所有索引 score 使用 Redis server time,不用本机时间:
```go
func (c *concurrencyCache) redisUnixSeconds(ctx context.Context) (int64, error) {
now, err := c.rdb.Time(ctx).Result()
if err != nil {
return 0, fmt.Errorf("redis TIME: %w", err)
}
return now.Unix(), nil
}
```
不要在 Lua 脚本里写全局索引,避免 CROSSSLOT。
### 1.2 增加索引 touch/remove/refresh helper
新增以下 helper。名字可以微调,但行为不能改。
```go
func (c *concurrencyCache) touchAccountActiveIndex(ctx context.Context, accountID int64, ttlSeconds int) {
c.touchActiveIndex(ctx, accountActiveIndexKey, accountID, ttlSeconds)
}
func (c *concurrencyCache) touchUserActiveIndex(ctx context.Context, userID int64, ttlSeconds int) {
c.touchActiveIndex(ctx, userActiveIndexKey, userID, ttlSeconds)
}
func (c *concurrencyCache) touchActiveIndex(ctx context.Context, indexKey string, id int64, ttlSeconds int) {
if c == nil || c.rdb == nil || id <= 0 || ttlSeconds <= 0 {
return
}
now, err := c.redisUnixSeconds(ctx)
if err != nil {
return
}
_ = c.rdb.ZAdd(ctx, indexKey, redis.Z{
Score: float64(now + int64(ttlSeconds)),
Member: strconv.FormatInt(id, 10),
}).Err()
}
```
索引维护是 best-effort,所以 helper 内部吞掉错误。不要把索引错误返回给 acquire/release 调用方。
再新增 refresh helper:
```go
func (c *concurrencyCache) refreshAccountActiveIndex(ctx context.Context, accountID int64) {
// 真实状态以 accountSlotKey(accountID) 和 accountWaitKey(accountID) 为准。
// 先清理该账号 slot 中过期成员,再读 ZCARD 和 GET wait。
// 如果 slotCount == 0 && waitCount <= 0:ZREM accountActiveIndexKey accountID。
// 否则:ZADD accountActiveIndexKey accountID,score = now + maxRelevantTTL。
}
func (c *concurrencyCache) refreshUserActiveIndex(ctx context.Context, userID int64) {
// 真实状态以 userSlotKey(userID) 和 waitQueueKey(userID) 为准。
// 行为同 refreshAccountActiveIndex。
}
```
实现要求:
- `refresh*` 必须 best-effort,不能向 release/decrement 返回索引错误。
- `waitCount` 读取 `redis.Nil` 时按 0 处理。
- `waitCount < 0` 必须按 0 处理。
- `slotCount > 0` 时 score 至少延长 `slotTTLSeconds`。
- `waitCount > 0` 时 score 至少延长 `waitQueueTTLSeconds`。
- 两者都存在时使用更大的 TTL。
### 1.3 修改账号写路径
修改 `AcquireAccountSlot`:
```go
result, err := acquireScript.Run(...).Int()
if err != nil { return false, err }
if result == 1 {
c.touchAccountActiveIndex(ctx, accountID, c.slotTTLSeconds)
}
return result == 1, nil
```
修改 `ReleaseAccountSlot`:
```go
if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil {
return err
}
c.refreshAccountActiveIndex(ctx, accountID)
return nil
```
修改 `IncrementAccountWaitCount`:
```go
result, err := incrementAccountWaitScript.Run(...).Int()
if err != nil { return false, err }
if result == 1 {
c.touchAccountActiveIndex(ctx, accountID, c.waitQueueTTLSeconds)
}
return result == 1, nil
```
修改 `DecrementAccountWaitCount`:
```go
_, err := decrementWaitScript.Run(...).Result()
if err == nil {
c.refreshAccountActiveIndex(ctx, accountID)
}
return err
```
### 1.4 修改用户写路径
同账号路径,修改:
- `AcquireUserSlot`
- `ReleaseUserSlot`
- `IncrementWaitCount`
- `DecrementWaitCount`
用户索引使用 `userActiveIndexKey`。
### 1.5 重写 GetActiveAccountLoadMap
删除 `scanAccountIDsByPrefix` 和 `parseAccountIDFromPrefixedKey` 的生产调用。`GetActiveAccountLoadMap` 必须改成:
1. 获取 Redis 当前秒。
2. `ZRemRangeByScore(accountActiveIndexKey, "-inf", strconv.FormatInt(now, 10))` 删除过期候选。
3. `ZRangeByScore(accountActiveIndexKey, &redis.ZRangeBy{Min: strconv.FormatInt(now+1, 10), Max: "+inf"})` 获取候选账号 ID。
4. 解析 member 为 `int64`,非法 member 记录到待删除列表。
5. 分块 pipeline,块大小 `activeIndexPipelineChunkSize`。
6. 对每个候选账号执行:
- `ZRemRangeByScore(accountSlotKey(id), "-inf", cutoffUnixSeconds)`
- `ZCard(accountSlotKey(id))`
- `Get(accountWaitKey(id))`
7. 构造结果时只返回 `currentConcurrency > 0 || waitingCount > 0` 的账号。
8. 对真实状态为空或 member 非法的账号执行 `ZREM accountActiveIndexKey member`。
9. 对真实状态仍活跃但 index score 已接近过期的账号,调用 `touchAccountActiveIndex` 刷新。
禁止:
- 禁止再扫 `concurrency:account:*`。
- 禁止再扫 `wait:account:*`。
- 禁止用索引里的 score 直接判断并发数。
## 第二部分:并发槽清理和启动清理
修改文件:`backend/internal/repository/concurrency_cache.go`。
### 2.1 重写 CleanupExpiredAccountSlotKeys
当前实现调用 `cleanupExpiredSlotKeysByPattern(ctx, accountSlotKeyPrefix+"*")`,必须删除。
新行为:
1. 获取 Redis 当前秒 `now`。
2. 从 `accountActiveIndexKey` 读取过期候选:
```go
ids, err := c.rdb.ZRangeByScore(ctx, accountActiveIndexKey, &redis.ZRangeBy{
Min: "-inf",
Max: strconv.FormatInt(now, 10),
Count: activeIndexCleanupBatchSize,
}).Result()
```
3. 对每个候选账号清理该账号 slot 过期成员并读真实状态。
4. 如果真实 `slotCount == 0 && waitCount <= 0`,从 `accountActiveIndexKey` 删除该账号。
5. 如果真实仍活跃,刷新 `accountActiveIndexKey` score。
6. 不需要处理不在索引中的账号;其 slot key 自身有 `EXPIRE`,并且 acquire/get-load 会惰性清理过期成员。
这个函数不再表示“遍历所有账号槽位 key”,而是“处理索引中到期的账号候选”。保留原函数名是为了少改接口。
### 2.2 重写 CleanupStaleProcessSlots
当前实现会扫描:
- `concurrency:account:*`
- `concurrency:user:*`
- `wait:account:*`
- `concurrency:wait:*`
必须去掉这些扫描。
新行为必须基于索引:
1. 从 `accountActiveIndexKey` 读取所有未过期候选账号。
2. 对每个账号:
- 对 `accountSlotKey(id)` 运行“单 key 清理脚本”,删除 requestID 前缀不是当前 `activeRequestPrefix` 的成员。
- 删除 `accountWaitKey(id)`,因为等待者属于旧进程,重启后不能继续等待。
- 调用 `refreshAccountActiveIndex(ctx, id)`。
3. 从 `userActiveIndexKey` 读取所有未过期候选用户。
4. 对每个用户:
- 对 `userSlotKey(id)` 运行同一个“单 key 清理脚本”。
- 删除 `waitQueueKey(id)`。
- 调用 `refreshUserActiveIndex(ctx, id)`。
新增单 key Lua 脚本,替代当前 `startupCleanupScript` 的多 key 版本:
```lua
local key = KEYS[1]
local activePrefix = ARGV[1]
local slotTTL = tonumber(ARGV[2])
local removed = 0
local members = redis.call('ZRANGE', key, 0, -1)
for _, member in ipairs(members) do
if string.sub(member, 1, string.len(activePrefix)) ~= activePrefix then
removed = removed + redis.call('ZREM', key, member)
end
end
if redis.call('ZCARD', key) == 0 then
redis.call('DEL', key)
else
redis.call('EXPIRE', key, slotTTL)
end
return removed
```
该脚本只接受一个 slot key,避免 Redis Cluster CROSSSLOT。
如果索引不存在或为空:
- `CleanupStaleProcessSlots` 直接返回 nil。
- 不要 fallback 到 `SCAN`。
- 旧版本遗留 key 依赖 Redis TTL 自然过期。不要在 app 启动时做兼容性 keyspace backfill。
### 2.3 删除旧扫描函数
删除以下生产函数:
- `scanAccountIDsByPrefix`
- `parseAccountIDFromPrefixedKey`,如果没有其他生产调用
- `cleanupExpiredSlotKeysByPattern`
- `cleanupSlotsByPattern`
- `deleteKeysByPattern`
如果测试需要解析 key,测试内自建 helper,不要保留生产 helper。
## 第三部分:UMQ 锁索引
修改文件:
- `backend/internal/repository/user_msg_queue_cache.go`
- `backend/internal/service/user_msg_queue_service.go`
### 3.1 修改 service 接口
在 `backend/internal/service/user_msg_queue_service.go` 的 `UserMsgQueueCache` 接口中删除:
```go
ScanLockKeys(ctx context.Context, maxCount int) ([]int64, error)
ForceReleaseLock(ctx context.Context, accountID int64) error
```
替换为:
```go
ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (cleaned int, err error)
```
原因:cleanup worker 不应该知道 Redis lock key 的枚举方式,也不应该先枚举再逐个 `ForceReleaseLock`。候选读取、PTTL 校验、索引刷新应该封装在 cache 层。
### 3.2 修改 acquireLockScript 返回值
当前脚本只返回 0/1。改成返回数组:
```lua
redis.replicate_commands()
local cur = redis.call('GET', KEYS[1])
local ttl = tonumber(ARGV[2])
if cur == ARGV[1] then
redis.call('PEXPIRE', KEYS[1], ttl)
local t = redis.call('TIME')
local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000)
return {1, ms + ttl}
end
if cur ~= false then
return {0, 0}
end
redis.call('SET', KEYS[1], ARGV[1], 'PX', ttl)
local t = redis.call('TIME')
local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000)
return {1, ms + ttl}
```
Go 侧解析:
- 第一个元素是 acquired,1 表示拿到锁。
- 第二个元素是 Redis 时间计算出的 `expireAtUnixMs`。
- acquired 为 1 时,best-effort 写 `ZADD umqLockIndexKey expireAtMs accountID`。
- `ZADD` 失败不能让 `AcquireLock` 返回失败。
### 3.3 修改 ReleaseLock
`ReleaseLock` 主逻辑保持原子释放锁和写 last key。
释放成功时:
```go
if result == 1 {
_ = c.rdb.ZRem(ctx, umqLockIndexKey, strconv.FormatInt(accountID, 10)).Err()
}
```
释放失败时不要删除索引。失败可能是 requestID 不匹配或 lock 已过期;cleanup worker 会处理 stale index。
### 3.4 新增 reconcile 脚本
删除 `forceReleaseLockScript` 的外部使用。新增脚本:
```lua
local pttl = redis.call('PTTL', KEYS[1])
if pttl == -2 then
return {-2, 0}
end
if pttl == -1 then
redis.call('DEL', KEYS[1])
return {-1, 0}
end
return {1, pttl}
```
返回语义:
- `-2`:lock key 不存在。Go 侧 `ZREM umqLockIndexKey accountID`。
- `-1`:lock key 存在但无 TTL,脚本已删除。Go 侧 `ZREM umqLockIndexKey accountID`,cleaned++。
- `1`:lock key 仍有 TTL。Go 侧用 Redis 当前毫秒 + pttl 刷新 `umqLockIndexKey` score。
### 3.5 实现 ReconcileExpiredLockCandidates
实现步骤:
1. 用 `c.rdb.Time(ctx)` 获取 Redis 当前毫秒 `nowMs`。
2. 从 `umqLockIndexKey` 取到期候选:
```go
members, err := c.rdb.ZRangeByScore(ctx, umqLockIndexKey, &redis.ZRangeBy{
Min: "-inf",
Max: strconv.FormatInt(nowMs, 10),
Count: int64(maxCount),
}).Result()
```
3. 逐个解析 accountID。非法 member 直接 `ZREM`。
4. 对合法 accountID 运行 reconcile 脚本,key 为 `umqLockKey(accountID)`。
5. 根据返回值删除索引、刷新索引或累计 cleaned。
6. 函数返回 cleaned 数。
禁止:
- 禁止 fallback 到 `SCAN umq:{*}:lock`。
- 禁止用 `KEYS umq:*`。
- 禁止 cleanup worker 自己解析 lock key。
### 3.6 修改 StartCleanupWorker
当前 worker 先 `ScanLockKeys` 再逐个 `ForceReleaseLock`。改成:
```go
cleaned, err := s.cache.ReconcileExpiredLockCandidates(ctx, 1000)
if err != nil {
logger.LegacyPrintf("service.umq", "Cleanup reconcile failed: %v", err)
return
}
if cleaned > 0 {
logger.LegacyPrintf("service.umq", "Cleanup completed: released %d orphaned locks", cleaned)
}
```
worker 不再知道扫描、PTTL、索引等细节。
### 3.7 删除旧 UMQ 扫描函数
删除:
- `umqScanPattern`
- `ScanLockKeys`
- `ForceReleaseLock`,如果无生产调用
如果测试仍需要强造 PTTL == -1 的 key,只在测试里直接写 Redis。
## 测试要求
### 并发缓存测试
新增或修改 `backend/internal/repository/concurrency_cache_*_test.go`。
必须覆盖:
1. `AcquireAccountSlot` 成功后 `GetActiveAccountLoadMap` 能看到该账号。
2. `ReleaseAccountSlot` 后 `GetActiveAccountLoadMap` 不再返回该账号。
3. `IncrementAccountWaitCount` 成功后 `GetActiveAccountLoadMap` 能看到 waiting count。
4. `DecrementAccountWaitCount` 后如果无 slot,则索引被移除。
5. `CleanupExpiredAccountSlotKeys` 不依赖 keyspace scan:测试里只创建索引成员和对应 slot key,然后确认会清理;再创建未索引 slot key,确认不会被该函数主动发现。
6. `CleanupStaleProcessSlots` 只处理索引中的 account/user,删除旧 request prefix 成员,保留当前 prefix 成员,删除 account/user wait key。
7. 索引中存在非法 member 时,`GetActiveAccountLoadMap` 不报错,并移除非法 member。
### UMQ 测试
新增或修改 `backend/internal/repository/user_msg_queue_cache*_test.go` 和 `backend/internal/service/user_msg_queue_service*_test.go`。
必须覆盖:
1. `AcquireLock` 成功后写入 `umq:lock:index`,score 大于 Redis 当前毫秒。
2. `ReleaseLock` 成功后删除 `umq:lock:index` member。
3. lock 已自然过期时,`ReconcileExpiredLockCandidates` 删除 stale index member。
4. lock 仍有 TTL 但 index score 到期时,`ReconcileExpiredLockCandidates` 刷新 index score,不删除 lock。
5. lock 存在且 `PTTL == -1` 时,`ReconcileExpiredLockCandidates` 删除 lock,删除 index member,并返回 cleaned=1。
6. index 中非法 member 不导致错误,并被删除。
7. `StartCleanupWorker` 调用 `ReconcileExpiredLockCandidates`,不再调用 `ScanLockKeys` 或 `ForceReleaseLock`。
### 禁止项测试
实现完成后运行:
```powershell
rg -n "\.Scan\(" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go
rg -n "ScanLockKeys|umqScanPattern|scanAccountIDsByPrefix|cleanupExpiredSlotKeysByPattern|cleanupSlotsByPattern|deleteKeysByPattern" backend/internal/repository/concurrency_cache.go backend/internal/repository/user_msg_queue_cache.go backend/internal/service/user_msg_queue_service.go
```
上述命令必须无输出。
再运行相关测试。按项目约定,编译很慢时先把代码复制到 WSL 文件系统再跑:
```bash
cd backend
go test ./internal/repository ./internal/service
```
如果全量包太慢,至少先跑:
```bash
cd backend
go test ./internal/repository -run 'Concurrency|UserMsgQueue|Redis'
go test ./internal/service -run 'Concurrency|UserMessageQueue'
```
## 迁移和兼容
不要在应用启动时扫描旧 key 回填索引。
原因:
- 这会把问题从运行期 `SCAN` 搬到启动期 `SCAN`。
- 生产实例重启时 Redis 已经高 CPU,启动扫描会放大抖动。
- 并发 slot key 和 wait key 都有 TTL,新版本写路径会为新流量维护索引,旧 key 可自然过期。
兼容策略:
- 新版本上线后,新请求会逐步填充 `concurrency:*:active_index` 和 `umq:lock:index`。
- 旧并发 slot key 没有索引时,不影响并发限制本身;对应账号下一次 acquire/get-load 会清理自己的 slot。
- 旧 UMQ lock 如果有 TTL,会自然过期。
- 极少数历史 `PTTL == -1` UMQ lock 且没有 index 的情况,不由应用自动发现。需要人工离线维护时,单独写一次性脚本,维护窗口运行,不要放进服务启动或后台 worker。
## 代码审查检查表
提交前逐项确认:
- [ ] 没有新增 `KEYS`。
- [ ] 没有新增生产路径 `SCAN`。
- [ ] 没有在 Lua 脚本中同时操作全局索引 key 和账号/用户局部 key。
- [ ] 索引维护失败不会让 acquire/release/decrement 的主结果失败。
- [ ] `GetActiveAccountLoadMap` 对 stale index、非法 member、Redis nil 都能正常返回。
- [ ] `CleanupExpiredAccountSlotKeys` 不再遍历 keyspace。
- [ ] `CleanupStaleProcessSlots` 不再遍历 keyspace。
- [ ] UMQ cleanup worker 不再知道 lock key pattern。
- [ ] 所有旧扫描 helper 已删除或仅存在于测试文件。
- [ ] 新测试覆盖成功路径、stale index、非法 member、PTTL -1、自然过期。
+375 -114
View File
@@ -36,6 +36,15 @@ const (
// 默认槽位过期时间(分钟),可通过配置覆盖
defaultSlotTTLMinutes = 15
// 活跃索引用来替代后台任务全量 SCAN 槽位键。
// member 是账号/用户 ID,score 是“预计仍需关注到”的 Redis Unix 秒时间戳。
accountActiveIndexKey = "concurrency:account:active_index" // ZSET member=accountID, score=expireAtUnixSeconds
userActiveIndexKey = "concurrency:user:active_index" // ZSET member=userID, score=expireAtUnixSeconds
// 后台清理只按批处理索引候选,避免单次任务占用 Redis 太久。
activeIndexCleanupBatchSize = 1000
activeIndexPipelineChunkSize = 500
)
var (
@@ -198,50 +207,24 @@ var (
return 1
`)
// cleanupExpiredSlotKeysScript 批量清理实际存在的账号槽位键,避免后台任务从数据库加载全量账号。
// KEYS = 有序集合键列表,ARGV[1] = TTL(秒)。
cleanupExpiredSlotKeysScript = redis.NewScript(`
-- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
-- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
redis.replicate_commands()
local ttl = tonumber(ARGV[1])
local timeResult = redis.call('TIME')
local now = tonumber(timeResult[1])
local expireBefore = now - ttl
local removed = 0
for i = 1, #KEYS do
local key = KEYS[i]
removed = removed + redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
if redis.call('ZCARD', key) == 0 then
redis.call('DEL', key)
else
redis.call('EXPIRE', key, ttl)
end
end
return removed
`)
// startupCleanupScript 清理非当前进程前缀的槽位成员。
// KEYS 是有序集合键列表,ARGV[1] 是当前进程前缀,ARGV[2] 是槽位 TTL。
// 遍历每个 KEYS[i],移除前缀不匹配的成员,清空后删 key,否则刷新 EXPIRE。
startupCleanupScript = redis.NewScript(`
// startupCleanupSlotScript 清理单个槽位 key 中非当前进程前缀的成员,避免 Redis Cluster CROSSSLOT。
// KEYS[1] 是有序集合键,ARGV[1] 是当前进程前缀,ARGV[2] 是槽位 TTL。
startupCleanupSlotScript = redis.NewScript(`
local key = KEYS[1]
local activePrefix = ARGV[1]
local slotTTL = tonumber(ARGV[2])
local removed = 0
for i = 1, #KEYS do
local key = KEYS[i]
local members = redis.call('ZRANGE', key, 0, -1)
for _, member in ipairs(members) do
if string.sub(member, 1, string.len(activePrefix)) ~= activePrefix then
removed = removed + redis.call('ZREM', key, member)
end
end
if redis.call('ZCARD', key) == 0 then
redis.call('DEL', key)
else
redis.call('EXPIRE', key, slotTTL)
local members = redis.call('ZRANGE', key, 0, -1)
for _, member in ipairs(members) do
if string.sub(member, 1, string.len(activePrefix)) ~= activePrefix then
removed = removed + redis.call('ZREM', key, member)
end
end
if redis.call('ZCARD', key) == 0 then
redis.call('DEL', key)
else
redis.call('EXPIRE', key, slotTTL)
end
return removed
`)
)
@@ -290,6 +273,198 @@ func accountWaitKey(accountID int64) string {
return fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
}
// redisUnixSeconds 统一使用 Redis 服务器时间,避免多实例本地时钟漂移导致索引提前/延后过期。
func (c *concurrencyCache) redisUnixSeconds(ctx context.Context) (int64, error) {
now, err := c.rdb.Time(ctx).Result()
if err != nil {
return 0, fmt.Errorf("redis TIME: %w", err)
}
return now.Unix(), nil
}
func (c *concurrencyCache) touchAccountActiveIndex(ctx context.Context, accountID int64, ttlSeconds int) {
c.touchActiveIndex(ctx, accountActiveIndexKey, accountID, ttlSeconds)
}
func (c *concurrencyCache) touchUserActiveIndex(ctx context.Context, userID int64, ttlSeconds int) {
c.touchActiveIndex(ctx, userActiveIndexKey, userID, ttlSeconds)
}
// touchActiveIndex 是写路径上的轻量标记:主操作已成功时,尽力把 ID 放入活跃索引。
// 索引失败不影响并发槽位/等待队列本身,后续释放或清理会再次校正。
func (c *concurrencyCache) touchActiveIndex(ctx context.Context, indexKey string, id int64, ttlSeconds int) {
if c == nil || c.rdb == nil || id <= 0 || ttlSeconds <= 0 {
return
}
now, err := c.redisUnixSeconds(ctx)
if err != nil {
return
}
_ = c.rdb.ZAdd(ctx, indexKey, redis.Z{
Score: float64(now + int64(ttlSeconds)),
Member: strconv.FormatInt(id, 10),
}).Err()
}
func (c *concurrencyCache) refreshAccountActiveIndex(ctx context.Context, accountID int64) {
c.refreshActiveIndex(ctx, accountActiveIndexKey, accountID, accountSlotKey(accountID), accountWaitKey(accountID))
}
func (c *concurrencyCache) refreshUserActiveIndex(ctx context.Context, userID int64) {
c.refreshActiveIndex(ctx, userActiveIndexKey, userID, userSlotKey(userID), waitQueueKey(userID))
}
// refreshActiveIndex 以 Redis 中的真实槽位/等待数为准重建索引状态。
// 释放槽位、等待计数减少、清理过期成员后都会调用它,防止索引残留。
func (c *concurrencyCache) refreshActiveIndex(ctx context.Context, indexKey string, id int64, slotKey, waitKey string) {
if c == nil || c.rdb == nil || id <= 0 {
return
}
now, err := c.redisUnixSeconds(ctx)
if err != nil {
return
}
load, err := c.readActiveLoadForKey(ctx, id, slotKey, waitKey, now)
if err != nil {
return
}
member := strconv.FormatInt(id, 10)
if load.slotCount == 0 && load.waitCount <= 0 {
_ = c.rdb.ZRem(ctx, indexKey, member).Err()
return
}
ttlSeconds := c.activeIndexTTL(load.slotCount, load.waitCount)
if ttlSeconds <= 0 {
return
}
_ = c.rdb.ZAdd(ctx, indexKey, redis.Z{
Score: float64(now + int64(ttlSeconds)),
Member: member,
}).Err()
}
type activeIndexLoad struct {
id int64
member string
slotCount int
waitCount int
}
// activeIndexTTL 取槽位 TTL 与等待队列 TTL 中仍然需要关注的较大值。
// 只要并发槽位或等待计数还有负载,就保留索引;两者都为 0 时调用方会删除索引。
func (c *concurrencyCache) activeIndexTTL(slotCount int, waitCount int) int {
ttlSeconds := 0
if slotCount > 0 {
ttlSeconds = c.slotTTLSeconds
}
if waitCount > 0 && c.waitQueueTTLSeconds > ttlSeconds {
ttlSeconds = c.waitQueueTTLSeconds
}
return ttlSeconds
}
// readActiveLoadForKey 读取单个 ID 的当前负载,并顺手清理该槽位集合中的过期成员。
func (c *concurrencyCache) readActiveLoadForKey(ctx context.Context, id int64, slotKey, waitKey string, now int64) (activeIndexLoad, error) {
cutoffTime := now - int64(c.slotTTLSeconds)
pipe := c.rdb.Pipeline()
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
zcardCmd := pipe.ZCard(ctx, slotKey)
getCmd := pipe.Get(ctx, waitKey)
if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
return activeIndexLoad{}, fmt.Errorf("pipeline exec: %w", err)
}
waitCount := 0
if v, err := getCmd.Int(); err == nil && v > 0 {
waitCount = v
}
return activeIndexLoad{
id: id,
member: strconv.FormatInt(id, 10),
slotCount: int(zcardCmd.Val()),
waitCount: waitCount,
}, nil
}
// readAccountIndexLoads 批量读取账号索引候选的真实负载。
// 分块 Pipeline 可以减少 Redis 往返,同时避免一次 Pipeline 塞入过多命令。
func (c *concurrencyCache) readAccountIndexLoads(ctx context.Context, members []string, now int64) ([]activeIndexLoad, []string, error) {
loads := make([]activeIndexLoad, 0, len(members))
staleMembers := make([]string, 0)
candidates := make([]activeIndexLoad, 0, len(members))
for _, member := range members {
id, err := strconv.ParseInt(member, 10, 64)
if err != nil || id <= 0 {
staleMembers = append(staleMembers, member)
continue
}
candidates = append(candidates, activeIndexLoad{id: id, member: member})
}
cutoffTime := now - int64(c.slotTTLSeconds)
for start := 0; start < len(candidates); start += activeIndexPipelineChunkSize {
end := start + activeIndexPipelineChunkSize
if end > len(candidates) {
end = len(candidates)
}
chunk := candidates[start:end]
pipe := c.rdb.Pipeline()
type accountCmd struct {
activeIndexLoad
zcardCmd *redis.IntCmd
getCmd *redis.StringCmd
}
cmds := make([]accountCmd, 0, len(chunk))
for _, candidate := range chunk {
slotKey := accountSlotKey(candidate.id)
waitKey := accountWaitKey(candidate.id)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
cmds = append(cmds, accountCmd{
activeIndexLoad: candidate,
zcardCmd: pipe.ZCard(ctx, slotKey),
getCmd: pipe.Get(ctx, waitKey),
})
}
if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
return nil, nil, fmt.Errorf("pipeline exec: %w", err)
}
for _, cmd := range cmds {
waitCount := 0
if v, err := cmd.getCmd.Int(); err == nil && v > 0 {
waitCount = v
}
loads = append(loads, activeIndexLoad{
id: cmd.id,
member: cmd.member,
slotCount: int(cmd.zcardCmd.Val()),
waitCount: waitCount,
})
}
}
return loads, staleMembers, nil
}
// removeActiveIndexMembers 清理无效 member;这是辅助索引的维护动作,调用方无需因为失败中断主流程。
func (c *concurrencyCache) removeActiveIndexMembers(ctx context.Context, indexKey string, members []string) {
if len(members) == 0 {
return
}
args := make([]any, 0, len(members))
for _, member := range members {
args = append(args, member)
}
_ = c.rdb.ZRem(ctx, indexKey, args...).Err()
}
// touchActiveIndexForLoad 根据已读取的真实负载刷新索引过期时间。
func (c *concurrencyCache) touchActiveIndexForLoad(ctx context.Context, indexKey string, load activeIndexLoad) {
c.touchActiveIndex(ctx, indexKey, load.id, c.activeIndexTTL(load.slotCount, load.waitCount))
}
// Account slot operations
func (c *concurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -299,12 +474,21 @@ func (c *concurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int
if err != nil {
return false, err
}
if result == 1 {
// 成功占槽后标记活跃账号,后台清理即可从索引定位候选账号。
c.touchAccountActiveIndex(ctx, accountID, c.slotTTLSeconds)
}
return result == 1, nil
}
func (c *concurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
key := accountSlotKey(accountID)
return c.rdb.ZRem(ctx, key, requestID).Err()
if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil {
return err
}
// 释放后用真实负载刷新索引;若没有槽位和等待计数,会移除索引 member。
c.refreshAccountActiveIndex(ctx, accountID)
return nil
}
func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID int64) (int, error) {
@@ -363,12 +547,21 @@ func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, ma
if err != nil {
return false, err
}
if result == 1 {
// 成功占槽后标记活跃用户,避免启动清理依赖全量 SCAN。
c.touchUserActiveIndex(ctx, userID, c.slotTTLSeconds)
}
return result == 1, nil
}
func (c *concurrencyCache) ReleaseUserSlot(ctx context.Context, userID int64, requestID string) error {
key := userSlotKey(userID)
return c.rdb.ZRem(ctx, key, requestID).Err()
if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil {
return err
}
// 释放后按 Redis 中剩余负载修正索引状态。
c.refreshUserActiveIndex(ctx, userID)
return nil
}
func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) (int, error) {
@@ -437,12 +630,20 @@ func (c *concurrencyCache) IncrementWaitCount(ctx context.Context, userID int64,
if err != nil {
return false, err
}
if result == 1 {
// 等待队列也会让用户保持“活跃”,否则槽位为 0 时后台任务可能漏看等待计数。
c.touchUserActiveIndex(ctx, userID, c.waitQueueTTLSeconds)
}
return result == 1, nil
}
func (c *concurrencyCache) DecrementWaitCount(ctx context.Context, userID int64) error {
key := waitQueueKey(userID)
_, err := decrementWaitScript.Run(ctx, c.rdb, []string{key}).Result()
if err == nil {
// 等待数减少后重新判断是否还需要保留索引。
c.refreshUserActiveIndex(ctx, userID)
}
return err
}
@@ -454,12 +655,20 @@ func (c *concurrencyCache) IncrementAccountWaitCount(ctx context.Context, accoun
if err != nil {
return false, err
}
if result == 1 {
// 账号级等待队列同样写入账号活跃索引,供负载查询和清理任务使用。
c.touchAccountActiveIndex(ctx, accountID, c.waitQueueTTLSeconds)
}
return result == 1, nil
}
func (c *concurrencyCache) DecrementAccountWaitCount(ctx context.Context, accountID int64) error {
key := accountWaitKey(accountID)
_, err := decrementWaitScript.Run(ctx, c.rdb, []string{key}).Result()
if err == nil {
// 等待计数归零后索引需要同步删除,避免后台任务反复处理空账号。
c.refreshAccountActiveIndex(ctx, accountID)
}
return err
}
@@ -599,101 +808,153 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic
func (c *concurrencyCache) CleanupExpiredAccountSlots(ctx context.Context, accountID int64) error {
key := accountSlotKey(accountID)
_, err := cleanupExpiredSlotsScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Result()
if err == nil {
// 单账号清理后同步索引,保持后台批量清理的候选集准确。
c.refreshAccountActiveIndex(ctx, accountID)
}
return err
}
func (c *concurrencyCache) CleanupExpiredAccountSlotKeys(ctx context.Context) error {
return c.cleanupExpiredSlotKeysByPattern(ctx, accountSlotKeyPrefix+"*")
// GetActiveAccountLoadMap 只读取活跃账号索引中的账号负载。
// 这是给热路径使用的轻量视图,避免为获取全局账号负载而扫描所有槽位键。
func (c *concurrencyCache) GetActiveAccountLoadMap(ctx context.Context) (map[int64]*service.AccountLoadInfo, error) {
now, err := c.redisUnixSeconds(ctx)
if err != nil {
return nil, err
}
if err := c.rdb.ZRemRangeByScore(ctx, accountActiveIndexKey, "-inf", strconv.FormatInt(now, 10)).Err(); err != nil {
return nil, fmt.Errorf("cleanup account active index: %w", err)
}
members, err := c.rdb.ZRangeByScore(ctx, accountActiveIndexKey, &redis.ZRangeBy{
Min: strconv.FormatInt(now+1, 10),
Max: "+inf",
}).Result()
if err != nil {
return nil, fmt.Errorf("read account active index: %w", err)
}
loads, staleMembers, err := c.readAccountIndexLoads(ctx, members, now)
if err != nil {
return nil, err
}
loadMap := make(map[int64]*service.AccountLoadInfo, len(loads))
for _, load := range loads {
if load.slotCount == 0 && load.waitCount <= 0 {
// 索引候选已没有实际负载,删除 member 而不是返回空负载。
staleMembers = append(staleMembers, load.member)
continue
}
loadMap[load.id] = &service.AccountLoadInfo{
AccountID: load.id,
CurrentConcurrency: load.slotCount,
WaitingCount: load.waitCount,
}
c.touchActiveIndexForLoad(ctx, accountActiveIndexKey, load)
}
c.removeActiveIndexMembers(ctx, accountActiveIndexKey, staleMembers)
return loadMap, nil
}
// CleanupExpiredAccountSlotKeys 只处理索引中过期的账号候选。
// 若候选仍有真实负载,则刷新索引;若没有负载,则移除索引 member。
func (c *concurrencyCache) CleanupExpiredAccountSlotKeys(ctx context.Context) error {
now, err := c.redisUnixSeconds(ctx)
if err != nil {
return err
}
members, err := c.rdb.ZRangeByScore(ctx, accountActiveIndexKey, &redis.ZRangeBy{
Min: "-inf",
Max: strconv.FormatInt(now, 10),
Count: activeIndexCleanupBatchSize,
}).Result()
if err != nil {
return fmt.Errorf("read expired account active index: %w", err)
}
loads, staleMembers, err := c.readAccountIndexLoads(ctx, members, now)
if err != nil {
return err
}
for _, load := range loads {
if load.slotCount == 0 && load.waitCount <= 0 {
// 真实槽位和等待数都为空,说明这个索引 member 已经完成使命。
staleMembers = append(staleMembers, load.member)
continue
}
c.touchActiveIndexForLoad(ctx, accountActiveIndexKey, load)
}
c.removeActiveIndexMembers(ctx, accountActiveIndexKey, staleMembers)
return nil
}
// CleanupStaleProcessSlots 启动时清理非当前进程前缀的槽位。
// 清理范围来自活跃索引,避免在 Redis 上 SCAN 全部 concurrency:* 键。
// API Key 槽位(concurrency:api_key:*)是 stats-only 数据:每次 Track/读取都会按分数
// 裁剪过期成员,key 自带 TTL,可在一个 slot TTL 内自愈,因此不参与启动清理。
func (c *concurrencyCache) CleanupStaleProcessSlots(ctx context.Context, activeRequestPrefix string) error {
if activeRequestPrefix == "" {
return nil
}
// 1. 清理有序集合中非当前进程前缀的成员
slotPatterns := []string{accountSlotKeyPrefix + "*", userSlotKeyPrefix + "*", apiKeySlotKeyPrefix + "*"}
for _, pattern := range slotPatterns {
if err := c.cleanupSlotsByPattern(ctx, pattern, activeRequestPrefix); err != nil {
return err
}
now, err := c.redisUnixSeconds(ctx)
if err != nil {
return err
}
// 2. 删除所有等待队列计数器(重启后计数器失效)
waitPatterns := []string{accountWaitKeyPrefix + "*", waitQueueKeyPrefix + "*"}
for _, pattern := range waitPatterns {
if err := c.deleteKeysByPattern(ctx, pattern); err != nil {
return err
}
accountMembers, err := c.activeIndexMembers(ctx, accountActiveIndexKey, now)
if err != nil {
return err
}
if err := c.cleanupStaleProcessSlotsForIndex(ctx, accountActiveIndexKey, accountMembers, activeRequestPrefix, accountSlotKey, accountWaitKey, c.refreshAccountActiveIndex); err != nil {
return err
}
return nil
userMembers, err := c.activeIndexMembers(ctx, userActiveIndexKey, now)
if err != nil {
return err
}
return c.cleanupStaleProcessSlotsForIndex(ctx, userActiveIndexKey, userMembers, activeRequestPrefix, userSlotKey, waitQueueKey, c.refreshUserActiveIndex)
}
// cleanupExpiredSlotKeysByPattern 扫描实际存在的账号槽位键并批量清理过期成员。
func (c *concurrencyCache) cleanupExpiredSlotKeysByPattern(ctx context.Context, pattern string) error {
const scanCount = 200
var cursor uint64
for {
keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, scanCount).Result()
if err != nil {
return fmt.Errorf("scan %s: %w", pattern, err)
}
if len(keys) > 0 {
_, err := cleanupExpiredSlotKeysScript.Run(ctx, c.rdb, keys, c.slotTTLSeconds).Result()
if err != nil {
return fmt.Errorf("cleanup expired slots %s: %w", pattern, err)
}
}
cursor = nextCursor
if cursor == 0 {
break
}
// activeIndexMembers 只返回当前仍未过期的索引 member;过期 member 由对应清理任务处理。
func (c *concurrencyCache) activeIndexMembers(ctx context.Context, indexKey string, now int64) ([]string, error) {
members, err := c.rdb.ZRangeByScore(ctx, indexKey, &redis.ZRangeBy{
Min: strconv.FormatInt(now+1, 10),
Max: "+inf",
}).Result()
if err != nil {
return nil, fmt.Errorf("read active index %s: %w", indexKey, err)
}
return nil
return members, nil
}
// cleanupSlotsByPattern 扫描匹配 pattern 的有序集合键,批量调用 Lua 脚本清理非当前进程成员。
func (c *concurrencyCache) cleanupSlotsByPattern(ctx context.Context, pattern, activePrefix string) error {
const scanCount = 200
var cursor uint64
for {
keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, scanCount).Result()
if err != nil {
return fmt.Errorf("scan %s: %w", pattern, err)
// cleanupStaleProcessSlotsForIndex 逐个处理索引中的账号/用户。
// Lua 脚本一次只碰一个槽位 key,兼容 Redis Cluster,随后删除重启后已失效的等待计数。
func (c *concurrencyCache) cleanupStaleProcessSlotsForIndex(
ctx context.Context,
indexKey string,
members []string,
activeRequestPrefix string,
slotKeyForID func(int64) string,
waitKeyForID func(int64) string,
refreshIndex func(context.Context, int64),
) error {
staleMembers := make([]string, 0)
for _, member := range members {
id, err := strconv.ParseInt(member, 10, 64)
if err != nil || id <= 0 {
staleMembers = append(staleMembers, member)
continue
}
if len(keys) > 0 {
_, err := startupCleanupScript.Run(ctx, c.rdb, keys, activePrefix, c.slotTTLSeconds).Result()
if err != nil {
return fmt.Errorf("cleanup slots %s: %w", pattern, err)
}
}
cursor = nextCursor
if cursor == 0 {
break
}
}
return nil
}
// deleteKeysByPattern 扫描匹配 pattern 的键并删除。
func (c *concurrencyCache) deleteKeysByPattern(ctx context.Context, pattern string) error {
const scanCount = 200
var cursor uint64
for {
keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, scanCount).Result()
if err != nil {
return fmt.Errorf("scan %s: %w", pattern, err)
if _, err := startupCleanupSlotScript.Run(ctx, c.rdb, []string{slotKeyForID(id)}, activeRequestPrefix, c.slotTTLSeconds).Result(); err != nil {
return fmt.Errorf("cleanup stale process slots %s: %w", slotKeyForID(id), err)
}
if len(keys) > 0 {
if err := c.rdb.Del(ctx, keys...).Err(); err != nil {
return fmt.Errorf("del %s: %w", pattern, err)
}
}
cursor = nextCursor
if cursor == 0 {
break
if err := c.rdb.Del(ctx, waitKeyForID(id)).Err(); err != nil {
return fmt.Errorf("delete stale wait key %s: %w", waitKeyForID(id), err)
}
refreshIndex(ctx, id)
}
c.removeActiveIndexMembers(ctx, indexKey, staleMembers)
return nil
}
@@ -6,6 +6,7 @@ import (
"context"
"errors"
"fmt"
"strconv"
"testing"
"time"
@@ -23,7 +24,8 @@ var testSlotTTL = time.Duration(testSlotTTLMinutes) * time.Minute
type ConcurrencyCacheSuite struct {
IntegrationRedisSuite
cache service.ConcurrencyCache
cache service.ConcurrencyCache
rawCache *concurrencyCache
}
func TestConcurrencyCacheSuite(t *testing.T) {
@@ -32,7 +34,8 @@ func TestConcurrencyCacheSuite(t *testing.T) {
func (s *ConcurrencyCacheSuite) SetupTest() {
s.IntegrationRedisSuite.SetupTest()
s.cache = NewConcurrencyCache(s.rdb, testSlotTTLMinutes, int(testSlotTTL.Seconds()))
s.rawCache = NewConcurrencyCache(s.rdb, testSlotTTLMinutes, int(testSlotTTL.Seconds())).(*concurrencyCache)
s.cache = s.rawCache
}
type apiKeyConcurrencyCacheForTest interface {
@@ -74,6 +77,61 @@ func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() {
require.Equal(s.T(), 1, cur, "expected 1 after release")
}
func (s *ConcurrencyCacheSuite) TestActiveAccountLoadMap_AcquireAndRelease() {
accountID := int64(610)
reqID := "active-load-req"
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID)
require.NoError(s.T(), err)
require.True(s.T(), ok)
loadMap, err := s.rawCache.GetActiveAccountLoadMap(s.ctx)
require.NoError(s.T(), err)
require.Contains(s.T(), loadMap, accountID)
require.Equal(s.T(), 1, loadMap[accountID].CurrentConcurrency)
require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, reqID))
loadMap, err = s.rawCache.GetActiveAccountLoadMap(s.ctx)
require.NoError(s.T(), err)
require.NotContains(s.T(), loadMap, accountID)
}
func (s *ConcurrencyCacheSuite) TestActiveAccountLoadMap_AccountWaitIndexLifecycle() {
accountID := int64(611)
ok, err := s.cache.IncrementAccountWaitCount(s.ctx, accountID, 2)
require.NoError(s.T(), err)
require.True(s.T(), ok)
loadMap, err := s.rawCache.GetActiveAccountLoadMap(s.ctx)
require.NoError(s.T(), err)
require.Contains(s.T(), loadMap, accountID)
require.Equal(s.T(), 1, loadMap[accountID].WaitingCount)
require.NoError(s.T(), s.cache.DecrementAccountWaitCount(s.ctx, accountID))
loadMap, err = s.rawCache.GetActiveAccountLoadMap(s.ctx)
require.NoError(s.T(), err)
require.NotContains(s.T(), loadMap, accountID)
}
func (s *ConcurrencyCacheSuite) TestActiveAccountLoadMap_RemovesInvalidIndexMember() {
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
Score: float64(now + 60),
Member: "not-an-account-id",
}).Err())
loadMap, err := s.rawCache.GetActiveAccountLoadMap(s.ctx)
require.NoError(s.T(), err)
require.Empty(s.T(), loadMap)
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, "not-an-account-id").Result()
require.ErrorIs(s.T(), err, redis.Nil)
}
func (s *ConcurrencyCacheSuite) TestAccountSlot_TTL() {
accountID := int64(11)
reqID := "req_ttl_test"
@@ -296,13 +354,17 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() {
accountID := int64(901)
userID := int64(902)
apiKeyID := int64(903)
unindexedAccountID := int64(1901)
accountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID)
apiKeyKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
unindexedAccountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, unindexedAccountID)
userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
unindexedAccountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, unindexedAccountID)
now := time.Now().Unix()
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKey,
redis.Z{Score: float64(now), Member: "oldproc-1"},
redis.Z{Score: float64(now), Member: "keep-1"},
@@ -311,12 +373,24 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() {
redis.Z{Score: float64(now), Member: "oldproc-2"},
redis.Z{Score: float64(now), Member: "keep-2"},
).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, unindexedAccountKey,
redis.Z{Score: float64(now), Member: "oldproc-unindexed"},
).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, apiKeyKey,
redis.Z{Score: float64(now), Member: "oldproc-3"},
redis.Z{Score: float64(now), Member: "keep-3"},
).Err())
require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, time.Minute).Err())
require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, time.Minute).Err())
require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 2, time.Minute).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
Score: float64(now + 60),
Member: strconv.FormatInt(accountID, 10),
}).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{
Score: float64(now + 60),
Member: strconv.FormatInt(userID, 10),
}).Err())
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-"))
@@ -328,15 +402,22 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() {
require.NoError(s.T(), err)
require.Equal(s.T(), []string{"keep-2"}, userMembers)
// API Key 槽位(stats-only)不在启动清理范围内,靠分数裁剪与 key TTL 自愈。
apiKeyMembers, err := s.rdb.ZRange(s.ctx, apiKeyKey, 0, -1).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), []string{"keep-3"}, apiKeyMembers)
require.ElementsMatch(s.T(), []string{"keep-3", "oldproc-3"}, apiKeyMembers)
_, err = s.rdb.Get(s.ctx, userWaitKey).Result()
require.True(s.T(), errors.Is(err, redis.Nil))
_, err = s.rdb.Get(s.ctx, accountWaitKey).Result()
require.True(s.T(), errors.Is(err, redis.Nil))
unindexedMembers, err := s.rdb.ZRange(s.ctx, unindexedAccountKey, 0, -1).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), []string{"oldproc-unindexed"}, unindexedMembers)
_, err = s.rdb.Get(s.ctx, unindexedAccountWaitKey).Result()
require.NoError(s.T(), err)
}
func (s *ConcurrencyCacheSuite) TestGetAccountConcurrency_Missing() {
@@ -487,11 +568,13 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlots_NoExpired() {
}
func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() {
now := time.Now().Unix()
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
expiredTime := now - int64(testSlotTTL.Seconds()) - 10
accountKeyWithFresh := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 301)
accountKeyExpiredOnly := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 302)
userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, 303)
unindexedAccountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 304)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKeyWithFresh,
redis.Z{Score: float64(expiredTime), Member: "expired"},
@@ -503,6 +586,13 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() {
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKey,
redis.Z{Score: float64(expiredTime), Member: "user-expired"},
).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, unindexedAccountKey,
redis.Z{Score: float64(expiredTime), Member: "unindexed-expired"},
).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey,
redis.Z{Score: float64(now), Member: "301"},
redis.Z{Score: float64(now), Member: "302"},
).Err())
require.NoError(s.T(), s.cache.CleanupExpiredAccountSlotKeys(s.ctx))
@@ -517,6 +607,16 @@ func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() {
userMembers, err := s.rdb.ZRange(s.ctx, userKey, 0, -1).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), []string{"user-expired"}, userMembers)
unindexedMembers, err := s.rdb.ZRange(s.ctx, unindexedAccountKey, 0, -1).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), []string{"unindexed-expired"}, unindexedMembers)
score, err := s.rdb.ZScore(s.ctx, accountActiveIndexKey, "301").Result()
require.NoError(s.T(), err)
require.Greater(s.T(), int64(score), now)
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, "302").Result()
require.ErrorIs(s.T(), err, redis.Nil)
}
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesAndWaitCounters() {
@@ -527,19 +627,28 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesA
userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
now := float64(time.Now().Unix())
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey,
redis.Z{Score: now, Member: "oldproc-1"},
redis.Z{Score: now, Member: "activeproc-1"},
redis.Z{Score: float64(now), Member: "oldproc-1"},
redis.Z{Score: float64(now), Member: "activeproc-1"},
).Err())
require.NoError(s.T(), s.rdb.Expire(s.ctx, accountSlotKey, testSlotTTL).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userSlotKey,
redis.Z{Score: now, Member: "oldproc-2"},
redis.Z{Score: now, Member: "activeproc-2"},
redis.Z{Score: float64(now), Member: "oldproc-2"},
redis.Z{Score: float64(now), Member: "activeproc-2"},
).Err())
require.NoError(s.T(), s.rdb.Expire(s.ctx, userSlotKey, testSlotTTL).Err())
require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, testSlotTTL).Err())
require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, testSlotTTL).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
Score: float64(now + 60),
Member: strconv.FormatInt(accountID, 10),
}).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{
Score: float64(now + 60),
Member: strconv.FormatInt(userID, 10),
}).Err())
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "activeproc-"))
@@ -560,8 +669,14 @@ func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesA
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_DeletesEmptySlotKeys() {
accountID := int64(903)
accountSlotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey, redis.Z{Score: float64(time.Now().Unix()), Member: "oldproc-1"}).Err())
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey, redis.Z{Score: float64(now), Member: "oldproc-1"}).Err())
require.NoError(s.T(), s.rdb.Expire(s.ctx, accountSlotKey, testSlotTTL).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
Score: float64(now + 60),
Member: strconv.FormatInt(accountID, 10),
}).Err())
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "activeproc-"))
@@ -5,8 +5,6 @@ import (
"errors"
"fmt"
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/redis/go-redis/v9"
@@ -18,18 +16,30 @@ const (
umqKeyPrefix = "umq:"
umqLockSuffix = ":lock" // STRING (requestID), PX lockTtlMs
umqLastSuffix = ":last" // STRING (毫秒时间戳), EX 60s
// 锁索引用来替代后台清理对 umq:*:lock 的全量 SCAN。
// member 是 accountID,score 是锁预计过期的 Redis Unix 毫秒时间戳。
umqLockIndexKey = "umq:lock:index" // ZSET member=accountID, score=lockExpireAtUnixMs
umqLockIndexCleanupBatchSize = 1000
)
// Lua 脚本:原子获取串行锁(SET NX PX + 重入安全)
// 返回 {是否获取成功, 锁预计过期时间毫秒},让 Go 侧用同一 Redis 时间源更新索引。
var acquireLockScript = redis.NewScript(`
redis.replicate_commands()
local cur = redis.call('GET', KEYS[1])
local ttl = tonumber(ARGV[2])
if cur == ARGV[1] then
redis.call('PEXPIRE', KEYS[1], tonumber(ARGV[2]))
return 1
redis.call('PEXPIRE', KEYS[1], ttl)
local t = redis.call('TIME')
local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000)
return {1, ms + ttl}
end
if cur ~= false then return 0 end
redis.call('SET', KEYS[1], ARGV[1], 'PX', tonumber(ARGV[2]))
return 1
if cur ~= false then return {0, 0} end
redis.call('SET', KEYS[1], ARGV[1], 'PX', ttl)
local t = redis.call('TIME')
local ms = tonumber(t[1])*1000 + math.floor(tonumber(t[2])/1000)
return {1, ms + ttl}
`)
// Lua 脚本:原子释放锁 + 记录完成时间(使用 Redis TIME 避免时钟偏差)
@@ -48,14 +58,18 @@ end
return 0
`)
// Lua 脚本:原子清理孤儿锁(仅在 PTTL == -1 时删除,避免 TOCTOU 竞态误删合法锁)
var forceReleaseLockScript = redis.NewScript(`
// Lua 脚本:校验锁 TTL 状态,PTTL == -1 时原子删除异常锁。
// 返回状态: -2=锁不存在,-1=无 TTL 的异常锁已删除,1=锁仍存活并返回剩余 PTTL。
var reconcileLockScript = redis.NewScript(`
local pttl = redis.call('PTTL', KEYS[1])
if pttl == -2 then
return {-2, 0}
end
if pttl == -1 then
redis.call('DEL', KEYS[1])
return 1
return {-1, 0}
end
return 0
return {1, pttl}
`)
type userMsgQueueCache struct {
@@ -77,22 +91,33 @@ func umqLastKey(accountID int64) string {
return umqKeyPrefix + "{" + strconv.FormatInt(accountID, 10) + "}" + umqLastSuffix
}
// umqScanPattern 用于 SCAN 扫描锁 key
func umqScanPattern() string {
return umqKeyPrefix + "{*}" + umqLockSuffix
}
// AcquireLock 尝试获取账号级串行锁
// 成功后尽力写入锁索引,后台清理只需要看“到期候选”而不是扫描所有锁 key。
func (c *userMsgQueueCache) AcquireLock(ctx context.Context, accountID int64, requestID string, lockTtlMs int) (bool, error) {
key := umqLockKey(accountID)
result, err := acquireLockScript.Run(ctx, c.rdb, []string{key}, requestID, lockTtlMs).Int()
result, err := acquireLockScript.Run(ctx, c.rdb, []string{key}, requestID, lockTtlMs).Result()
if err != nil {
return false, fmt.Errorf("umq acquire lock: %w", err)
}
return result == 1, nil
acquired, err := redisScriptInt64At(result, 0)
if err != nil {
return false, fmt.Errorf("umq parse acquire lock result: %w", err)
}
expireAtMs, err := redisScriptInt64At(result, 1)
if err != nil {
return false, fmt.Errorf("umq parse acquire lock expire: %w", err)
}
if acquired == 1 {
_ = c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{
Score: float64(expireAtMs),
Member: strconv.FormatInt(accountID, 10),
}).Err()
}
return acquired == 1, nil
}
// ReleaseLock 释放锁并记录完成时间
// 只有 requestID 匹配时才删除锁索引,避免误删其他请求重入后写入的新锁。
func (c *userMsgQueueCache) ReleaseLock(ctx context.Context, accountID int64, requestID string) (bool, error) {
lockKey := umqLockKey(accountID)
lastKey := umqLastKey(accountID)
@@ -100,6 +125,9 @@ func (c *userMsgQueueCache) ReleaseLock(ctx context.Context, accountID int64, re
if err != nil {
return false, fmt.Errorf("umq release lock: %w", err)
}
if result == 1 {
_ = c.rdb.ZRem(ctx, umqLockIndexKey, strconv.FormatInt(accountID, 10)).Err()
}
return result == 1, nil
}
@@ -120,65 +148,6 @@ func (c *userMsgQueueCache) GetLastCompletedMs(ctx context.Context, accountID in
return ms, nil
}
// ForceReleaseLock 原子清理孤儿锁(仅在 PTTL == -1 时删除,防止 TOCTOU 竞态误删合法锁)
func (c *userMsgQueueCache) ForceReleaseLock(ctx context.Context, accountID int64) error {
key := umqLockKey(accountID)
_, err := forceReleaseLockScript.Run(ctx, c.rdb, []string{key}).Result()
if err != nil && !errors.Is(err, redis.Nil) {
return fmt.Errorf("umq force release lock: %w", err)
}
return nil
}
// ScanLockKeys 扫描所有锁 key,仅返回 PTTL == -1(无过期时间)的孤儿锁 accountID 列表
// 正常的锁都有 PX 过期时间,PTTL == -1 表示异常状态(如 Redis 故障恢复后丢失 TTL)
func (c *userMsgQueueCache) ScanLockKeys(ctx context.Context, maxCount int) ([]int64, error) {
var accountIDs []int64
var cursor uint64
pattern := umqScanPattern()
for {
keys, nextCursor, err := c.rdb.Scan(ctx, cursor, pattern, 100).Result()
if err != nil {
return nil, fmt.Errorf("umq scan lock keys: %w", err)
}
for _, key := range keys {
// 检查 PTTL:只清理 PTTL == -1(无过期时间)的异常锁
pttl, err := c.rdb.PTTL(ctx, key).Result()
if err != nil {
continue
}
// PTTL 返回值:-2 = key 不存在,-1 = 无过期时间,>0 = 剩余毫秒
// go-redis 对哨兵值 -1/-2 不乘精度系数,直接返回 time.Duration(-1)/-2
// 只删除 -1(无过期时间的异常锁),跳过正常持有的锁
if pttl != time.Duration(-1) {
continue
}
// 从 key 中提取 accountID: umq:{123}:lock → 提取 {} 内的数字
openBrace := strings.IndexByte(key, '{')
closeBrace := strings.IndexByte(key, '}')
if openBrace < 0 || closeBrace <= openBrace+1 {
continue
}
idStr := key[openBrace+1 : closeBrace]
id, err := strconv.ParseInt(idStr, 10, 64)
if err != nil {
continue
}
accountIDs = append(accountIDs, id)
if len(accountIDs) >= maxCount {
return accountIDs, nil
}
}
cursor = nextCursor
if cursor == 0 {
break
}
}
return accountIDs, nil
}
// GetCurrentTimeMs 通过 Redis TIME 命令获取当前服务器时间(毫秒),确保与锁记录的时间源一致
func (c *userMsgQueueCache) GetCurrentTimeMs(ctx context.Context) (int64, error) {
t, err := c.rdb.Time(ctx).Result()
@@ -187,3 +156,85 @@ func (c *userMsgQueueCache) GetCurrentTimeMs(ctx context.Context) (int64, error)
}
return t.UnixMilli(), nil
}
// ReconcileExpiredLockCandidates 只处理索引里已经到期的候选锁。
// 候选到期不等于锁一定失效:可能是续租后索引滞后,所以必须再用 PTTL 二次确认。
func (c *userMsgQueueCache) ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (int, error) {
if maxCount <= 0 {
maxCount = umqLockIndexCleanupBatchSize
}
nowMs, err := c.GetCurrentTimeMs(ctx)
if err != nil {
return 0, err
}
members, err := c.rdb.ZRangeByScore(ctx, umqLockIndexKey, &redis.ZRangeBy{
Min: "-inf",
Max: strconv.FormatInt(nowMs, 10),
Count: int64(maxCount),
}).Result()
if err != nil {
return 0, fmt.Errorf("umq read lock index: %w", err)
}
cleaned := 0
for _, member := range members {
accountID, err := strconv.ParseInt(member, 10, 64)
if err != nil || accountID <= 0 {
_ = c.rdb.ZRem(ctx, umqLockIndexKey, member).Err()
continue
}
result, err := reconcileLockScript.Run(ctx, c.rdb, []string{umqLockKey(accountID)}).Result()
if err != nil && !errors.Is(err, redis.Nil) {
return cleaned, fmt.Errorf("umq reconcile lock: %w", err)
}
status, err := redisScriptInt64At(result, 0)
if err != nil {
return cleaned, fmt.Errorf("umq parse reconcile status: %w", err)
}
pttl, err := redisScriptInt64At(result, 1)
if err != nil {
return cleaned, fmt.Errorf("umq parse reconcile pttl: %w", err)
}
switch status {
case -2:
// 锁自然过期或已释放,只需移除索引残留。
_ = c.rdb.ZRem(ctx, umqLockIndexKey, member).Err()
case -1:
// 无 TTL 的锁会永久阻塞队列,Lua 已原子删除它,这里统计一次清理。
_ = c.rdb.ZRem(ctx, umqLockIndexKey, member).Err()
cleaned++
case 1:
// 锁仍存活,说明索引过期时间滞后;按剩余 PTTL 重新排期。
_ = c.rdb.ZAdd(ctx, umqLockIndexKey, redis.Z{
Score: float64(nowMs + pttl),
Member: member,
}).Err()
}
}
return cleaned, nil
}
// redisScriptInt64At 兼容 go-redis 对 Lua 数组元素的不同返回类型。
func redisScriptInt64At(result any, index int) (int64, error) {
values, ok := result.([]any)
if !ok {
return 0, fmt.Errorf("expected redis script array, got %T", result)
}
if index < 0 || index >= len(values) {
return 0, fmt.Errorf("redis script array missing index %d", index)
}
switch v := values[index].(type) {
case int64:
return v, nil
case int:
return int64(v), nil
case string:
return strconv.ParseInt(v, 10, 64)
case []byte:
return strconv.ParseInt(string(v), 10, 64)
default:
return 0, fmt.Errorf("unexpected redis script value %T", v)
}
}
@@ -0,0 +1,128 @@
//go:build integration
package repository
import (
"errors"
"testing"
"time"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
type UserMsgQueueCacheSuite struct {
IntegrationRedisSuite
cache *userMsgQueueCache
}
func TestUserMsgQueueCacheSuite(t *testing.T) {
suite.Run(t, new(UserMsgQueueCacheSuite))
}
func (s *UserMsgQueueCacheSuite) SetupTest() {
s.IntegrationRedisSuite.SetupTest()
s.cache = NewUserMsgQueueCache(s.rdb).(*userMsgQueueCache)
}
func (s *UserMsgQueueCacheSuite) TestAcquireLockWritesIndexAndReleaseRemovesIt() {
accountID := int64(701)
nowMs, err := s.cache.GetCurrentTimeMs(s.ctx)
require.NoError(s.T(), err)
acquired, err := s.cache.AcquireLock(s.ctx, accountID, "req-701", 10_000)
require.NoError(s.T(), err)
require.True(s.T(), acquired)
score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "701").Result()
require.NoError(s.T(), err)
require.Greater(s.T(), int64(score), nowMs)
released, err := s.cache.ReleaseLock(s.ctx, accountID, "req-701")
require.NoError(s.T(), err)
require.True(s.T(), released)
_, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "701").Result()
require.ErrorIs(s.T(), err, redis.Nil)
}
func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRemovesNaturallyExpiredLockIndex() {
accountID := int64(702)
acquired, err := s.cache.AcquireLock(s.ctx, accountID, "req-702", 20)
require.NoError(s.T(), err)
require.True(s.T(), acquired)
score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "702").Result()
require.NoError(s.T(), err)
require.Eventually(s.T(), func() bool {
nowMs, err := s.cache.GetCurrentTimeMs(s.ctx)
return err == nil && nowMs >= int64(score)
}, time.Second, 10*time.Millisecond)
cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000)
require.NoError(s.T(), err)
require.Equal(s.T(), 0, cleaned)
_, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "702").Result()
require.ErrorIs(s.T(), err, redis.Nil)
}
func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRefreshesLiveLockIndex() {
accountID := int64(703)
nowMs, err := s.cache.GetCurrentTimeMs(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.Set(s.ctx, umqLockKey(accountID), "req-703", time.Minute).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, umqLockIndexKey, redis.Z{
Score: float64(nowMs - 1),
Member: "703",
}).Err())
cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000)
require.NoError(s.T(), err)
require.Equal(s.T(), 0, cleaned)
score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "703").Result()
require.NoError(s.T(), err)
require.Greater(s.T(), int64(score), nowMs)
exists, err := s.rdb.Exists(s.ctx, umqLockKey(accountID)).Result()
require.NoError(s.T(), err)
require.EqualValues(s.T(), 1, exists)
}
func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesDeletesNoTTLLock() {
accountID := int64(704)
nowMs, err := s.cache.GetCurrentTimeMs(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.Set(s.ctx, umqLockKey(accountID), "req-704", 0).Err())
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, umqLockIndexKey, redis.Z{
Score: float64(nowMs),
Member: "704",
}).Err())
cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000)
require.NoError(s.T(), err)
require.Equal(s.T(), 1, cleaned)
exists, err := s.rdb.Exists(s.ctx, umqLockKey(accountID)).Result()
require.NoError(s.T(), err)
require.EqualValues(s.T(), 0, exists)
_, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "704").Result()
require.ErrorIs(s.T(), err, redis.Nil)
}
func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRemovesInvalidMember() {
nowMs, err := s.cache.GetCurrentTimeMs(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, umqLockIndexKey, redis.Z{
Score: float64(nowMs),
Member: "not-an-account-id",
}).Err())
cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000)
require.NoError(s.T(), err)
require.Equal(s.T(), 0, cleaned)
_, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "not-an-account-id").Result()
require.True(s.T(), errors.Is(err, redis.Nil))
}
@@ -37,7 +37,7 @@ type ConcurrencyCache interface {
ReleaseUserSlot(ctx context.Context, userID int64, requestID string) error
GetUserConcurrency(ctx context.Context, userID int64) (int, error)
// 等待队列计数(只在首次创建时设置 TTL)
// 等待队列计数(每次入队都会刷新 TTL,避免长时间排队时计数提前过期)
IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error)
DecrementWaitCount(ctx context.Context, userID int64) error
@@ -25,10 +25,8 @@ type UserMsgQueueCache interface {
GetLastCompletedMs(ctx context.Context, accountID int64) (int64, error)
// GetCurrentTimeMs 获取 Redis 服务器当前时间(毫秒),与 ReleaseLock 记录的时间源一致
GetCurrentTimeMs(ctx context.Context) (int64, error)
// ForceReleaseLock 强制释放锁(孤儿锁清理)
ForceReleaseLock(ctx context.Context, accountID int64) error
// ScanLockKeys 扫描 PTTL == -1 的孤儿锁 key,返回 accountID 列表
ScanLockKeys(ctx context.Context, maxCount int) ([]int64, error)
// ReconcileExpiredLockCandidates 处理锁索引中的到期候选,按真实 PTTL 清理或刷新索引
ReconcileExpiredLockCandidates(ctx context.Context, maxCount int) (cleaned int, err error)
}
// QueueLockResult 锁获取结果
@@ -246,8 +244,8 @@ func (s *UserMessageQueueService) CalculateRPMAwareDelay(ctx context.Context, ac
return applyJitter(baseDelay, 0.15)
}
// StartCleanupWorker 启动孤儿锁清理 worker
// 定期 SCAN umq:*:lock 并清理 PTTL == -1 的异常锁(PTTL 检查在 cache.ScanLockKeys 内完成)
// StartCleanupWorker 启动孤儿锁清理 worker。
// worker 只处理锁索引中的到期候选,真正删除前由 cache 层再次校验锁 PTTL。
func (s *UserMessageQueueService) StartCleanupWorker(interval time.Duration) {
if s == nil || s.cache == nil || interval <= 0 {
return
@@ -257,23 +255,13 @@ func (s *UserMessageQueueService) StartCleanupWorker(interval time.Duration) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
accountIDs, err := s.cache.ScanLockKeys(ctx, 1000)
// 每轮限制处理数量,避免清理任务在大量过期候选时长时间占用 Redis。
cleaned, err := s.cache.ReconcileExpiredLockCandidates(ctx, 1000)
if err != nil {
logger.LegacyPrintf("service.umq", "Cleanup scan failed: %v", err)
logger.LegacyPrintf("service.umq", "Cleanup reconcile failed: %v", err)
return
}
cleaned := 0
for _, accountID := range accountIDs {
cleanCtx, cleanCancel := context.WithTimeout(context.Background(), 2*time.Second)
if err := s.cache.ForceReleaseLock(cleanCtx, accountID); err != nil {
logger.LegacyPrintf("service.umq", "Cleanup force release failed for account %d: %v", accountID, err)
} else {
cleaned++
}
cleanCancel()
}
if cleaned > 0 {
logger.LegacyPrintf("service.umq", "Cleanup completed: released %d orphaned locks", cleaned)
}
@@ -0,0 +1,54 @@
//go:build unit
package service
import (
"context"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type cleanupWorkerUserMsgQueueCache struct {
reconcileCalls atomic.Int64
maxCount atomic.Int64
}
var _ UserMsgQueueCache = (*cleanupWorkerUserMsgQueueCache)(nil)
func (c *cleanupWorkerUserMsgQueueCache) AcquireLock(context.Context, int64, string, int) (bool, error) {
return true, nil
}
func (c *cleanupWorkerUserMsgQueueCache) ReleaseLock(context.Context, int64, string) (bool, error) {
return true, nil
}
func (c *cleanupWorkerUserMsgQueueCache) GetLastCompletedMs(context.Context, int64) (int64, error) {
return 0, nil
}
func (c *cleanupWorkerUserMsgQueueCache) GetCurrentTimeMs(context.Context) (int64, error) {
return time.Now().UnixMilli(), nil
}
func (c *cleanupWorkerUserMsgQueueCache) ReconcileExpiredLockCandidates(_ context.Context, maxCount int) (int, error) {
c.reconcileCalls.Add(1)
c.maxCount.Store(int64(maxCount))
return 1, nil
}
func TestStartCleanupWorker_ReconcilesExpiredLockCandidates(t *testing.T) {
cache := &cleanupWorkerUserMsgQueueCache{}
svc := NewUserMessageQueueService(cache, nil, nil)
defer svc.Stop()
svc.StartCleanupWorker(time.Millisecond)
require.Eventually(t, func() bool {
return cache.reconcileCalls.Load() > 0
}, time.Second, 10*time.Millisecond)
require.EqualValues(t, 1000, cache.maxCount.Load())
}