concurrency_cache.go 44 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262
  1. package repository
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "strconv"
  7. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/redis/go-redis/v9"
  10. )
  11. // 并发控制缓存常量定义
  12. //
  13. // 性能优化说明:
  14. // 原实现使用 SCAN 命令遍历独立的槽位键(concurrency:account:{id}:{requestID}),
  15. // 在高并发场景下 SCAN 需要多次往返,且遍历大量键时性能下降明显。
  16. //
  17. // 新实现改用 Redis 有序集合(Sorted Set):
  18. // 1. 每个账号/用户只有一个键,成员为 requestID,分数为时间戳
  19. // 2. 使用 ZCARD 原子获取并发数,时间复杂度 O(1)
  20. // 3. 使用 ZREMRANGEBYSCORE 清理过期槽位,避免手动管理 TTL
  21. // 4. 单次 Redis 调用完成计数,减少网络往返
  22. const (
  23. // 并发槽位键前缀(有序集合)
  24. // 格式: concurrency:account:{accountID}
  25. accountSlotKeyPrefix = "concurrency:account:"
  26. // 格式: concurrency:user:{userID}
  27. userSlotKeyPrefix = "concurrency:user:"
  28. // 格式: concurrency:api_key:{apiKeyID}
  29. apiKeySlotKeyPrefix = "concurrency:api_key:"
  30. liveAccountSlotKeyPrefix = "concurrency:live:account:"
  31. liveUserSlotKeyPrefix = "concurrency:live:user:"
  32. liveAPIKeySlotKeyPrefix = "concurrency:live:api_key:"
  33. // API-key-scoped client WebSocket ingress leases use a shorter TTL than
  34. // ordinary request slots, because idle ingress sessions do not hold a turn slot.
  35. openAIWSIngressLeaseKeyPrefix = "concurrency:openai_ws_ingress:api_key:"
  36. openAIWSIngressLeaseTTLSeconds = 60
  37. liveLeaseTTLSeconds = 60
  38. // 等待队列计数器格式: concurrency:wait:{userID}
  39. waitQueueKeyPrefix = "concurrency:wait:"
  40. // 账号级等待队列计数器格式: wait:account:{accountID}
  41. accountWaitKeyPrefix = "wait:account:"
  42. // 默认槽位过期时间(分钟),可通过配置覆盖
  43. defaultSlotTTLMinutes = 15
  44. // 活跃索引用来替代后台任务全量 SCAN 槽位键。
  45. // member 是账号/用户 ID,score 是“预计仍需关注到”的 Redis Unix 秒时间戳。
  46. accountActiveIndexKey = "concurrency:account:active_index" // ZSET member=accountID, score=expireAtUnixSeconds
  47. userActiveIndexKey = "concurrency:user:active_index" // ZSET member=userID, score=expireAtUnixSeconds
  48. // 后台清理只按批处理索引候选,避免单次任务占用 Redis 太久。
  49. activeIndexCleanupBatchSize = 1000
  50. activeIndexPipelineChunkSize = 500
  51. // 一次性迁移 marker:活跃索引机制上线前遗留的等待计数键无法被索引发现,
  52. // 且有流量时 TTL 会被不断刷新,必须清扫一次。marker 存在即代表已完成。
  53. legacyWaitSweepMarkerKey = "concurrency:startup:legacy_wait_sweep:v1"
  54. )
  55. var (
  56. // acquireScript 使用有序集合计数并在未达上限时添加槽位
  57. // 使用 Redis TIME 命令获取服务器时间,避免多实例时钟不同步问题
  58. // KEYS[1] = 普通槽位键,KEYS[2] = 对应 Live 槽位键
  59. // ARGV[1] = maxConcurrency
  60. // ARGV[2] = TTL(秒)
  61. // ARGV[3] = requestID
  62. // 返回 {是否成功, Redis 当前秒},Go 侧复用同一时间源写活跃索引,省去额外 TIME 往返。
  63. acquireScript = redis.NewScript(`
  64. -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
  65. -- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
  66. redis.replicate_commands()
  67. local key = KEYS[1]
  68. local liveKey = KEYS[2]
  69. local maxConcurrency = tonumber(ARGV[1])
  70. local ttl = tonumber(ARGV[2])
  71. local requestID = ARGV[3]
  72. -- 使用 Redis 服务器时间,确保多实例时钟一致
  73. local timeResult = redis.call('TIME')
  74. local now = tonumber(timeResult[1])
  75. local expireBefore = now - ttl
  76. -- 清理过期槽位
  77. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  78. redis.call('ZREMRANGEBYSCORE', liveKey, '-inf', now - 60)
  79. -- 检查是否已存在(支持重试场景刷新时间戳)
  80. local exists = redis.call('ZSCORE', key, requestID)
  81. if exists ~= false then
  82. redis.call('ZADD', key, now, requestID)
  83. redis.call('EXPIRE', key, ttl)
  84. return {1, now}
  85. end
  86. -- 检查是否达到并发上限
  87. local count = redis.call('ZCARD', key) + redis.call('ZCARD', liveKey)
  88. if count < maxConcurrency then
  89. redis.call('ZADD', key, now, requestID)
  90. redis.call('EXPIRE', key, ttl)
  91. return {1, now}
  92. end
  93. return {0, now}
  94. `)
  95. // getCountScript 统计有序集合中的槽位数量并清理过期条目
  96. // 使用 Redis TIME 命令获取服务器时间
  97. // KEYS[1] = 普通槽位键,KEYS[2] = 对应 Live 槽位键
  98. // ARGV[1] = TTL(秒)
  99. getCountScript = redis.NewScript(`
  100. -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
  101. -- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
  102. redis.replicate_commands()
  103. local key = KEYS[1]
  104. local liveKey = KEYS[2]
  105. local ttl = tonumber(ARGV[1])
  106. -- 使用 Redis 服务器时间
  107. local timeResult = redis.call('TIME')
  108. local now = tonumber(timeResult[1])
  109. local expireBefore = now - ttl
  110. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  111. redis.call('ZREMRANGEBYSCORE', liveKey, '-inf', now - 60)
  112. return redis.call('ZCARD', key) + redis.call('ZCARD', liveKey)
  113. `)
  114. acquireLiveLeaseScript = redis.NewScript(`
  115. redis.replicate_commands()
  116. local accountRegular = KEYS[1]
  117. local accountLive = KEYS[2]
  118. local userRegular = KEYS[3]
  119. local userLive = KEYS[4]
  120. local apiLive = KEYS[5]
  121. local accountMax = tonumber(ARGV[1])
  122. local userMax = tonumber(ARGV[2])
  123. local ttl = tonumber(ARGV[3])
  124. local leaseID = ARGV[4]
  125. local replacing = tonumber(ARGV[5])
  126. local now = tonumber(redis.call('TIME')[1])
  127. local liveExpireBefore = now - ttl
  128. redis.call('ZREMRANGEBYSCORE', accountLive, '-inf', liveExpireBefore)
  129. redis.call('ZREMRANGEBYSCORE', userLive, '-inf', liveExpireBefore)
  130. redis.call('ZREMRANGEBYSCORE', apiLive, '-inf', liveExpireBefore)
  131. if redis.call('ZSCORE', accountLive, leaseID) ~= false then
  132. return 1
  133. end
  134. local accountCount = redis.call('ZCARD', accountRegular) + redis.call('ZCARD', accountLive)
  135. local userCount = redis.call('ZCARD', userRegular) + redis.call('ZCARD', userLive)
  136. local allowance = 0
  137. if replacing == 1 then allowance = 1 end
  138. if accountMax > 0 and accountCount >= accountMax + allowance then return 0 end
  139. if userMax > 0 and userCount >= userMax + allowance then return 0 end
  140. redis.call('ZADD', accountLive, now, leaseID)
  141. redis.call('ZADD', userLive, now, leaseID)
  142. redis.call('ZADD', apiLive, now, leaseID)
  143. redis.call('EXPIRE', accountLive, ttl)
  144. redis.call('EXPIRE', userLive, ttl)
  145. redis.call('EXPIRE', apiLive, ttl)
  146. return 1
  147. `)
  148. refreshLiveLeaseScript = redis.NewScript(`
  149. redis.replicate_commands()
  150. local ttl = tonumber(ARGV[1])
  151. local leaseID = ARGV[2]
  152. local now = tonumber(redis.call('TIME')[1])
  153. local expireBefore = now - ttl
  154. for _, key in ipairs(KEYS) do
  155. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  156. if redis.call('ZSCORE', key, leaseID) == false then return 0 end
  157. end
  158. for _, key in ipairs(KEYS) do
  159. redis.call('ZADD', key, now, leaseID)
  160. redis.call('EXPIRE', key, ttl)
  161. end
  162. return 1
  163. `)
  164. // trackSlotScript 记录 stats-only 槽位,不做并发上限判断。
  165. // KEYS[1] = 有序集合键
  166. // ARGV[1] = TTL(秒)
  167. // ARGV[2] = requestID
  168. trackSlotScript = redis.NewScript(`
  169. -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
  170. -- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
  171. redis.replicate_commands()
  172. local key = KEYS[1]
  173. local ttl = tonumber(ARGV[1])
  174. local requestID = ARGV[2]
  175. local timeResult = redis.call('TIME')
  176. local now = tonumber(timeResult[1])
  177. local expireBefore = now - ttl
  178. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  179. redis.call('ZADD', key, now, requestID)
  180. redis.call('EXPIRE', key, ttl)
  181. return 1
  182. `)
  183. // acquireOpenAIWSIngressLeaseScript atomically reaps crashed members and
  184. // acquires or refreshes one API-key-scoped ingress lease using Redis TIME.
  185. acquireOpenAIWSIngressLeaseScript = redis.NewScript(`
  186. redis.replicate_commands()
  187. local key = KEYS[1]
  188. local maxConnections = tonumber(ARGV[1])
  189. local ttl = tonumber(ARGV[2])
  190. local leaseID = ARGV[3]
  191. local now = tonumber(redis.call('TIME')[1])
  192. local expireBefore = now - ttl
  193. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  194. if redis.call('ZSCORE', key, leaseID) ~= false then
  195. redis.call('ZADD', key, now, leaseID)
  196. redis.call('EXPIRE', key, ttl)
  197. return 1
  198. end
  199. if redis.call('ZCARD', key) < maxConnections then
  200. redis.call('ZADD', key, now, leaseID)
  201. redis.call('EXPIRE', key, ttl)
  202. return 1
  203. end
  204. return 0
  205. `)
  206. // refreshOpenAIWSIngressLeaseScript does not recreate a missing member: a
  207. // process that lost its lease must terminate its local WebSocket instead of
  208. // silently continuing beyond the distributed cap.
  209. refreshOpenAIWSIngressLeaseScript = redis.NewScript(`
  210. redis.replicate_commands()
  211. local key = KEYS[1]
  212. local ttl = tonumber(ARGV[1])
  213. local leaseID = ARGV[2]
  214. local now = tonumber(redis.call('TIME')[1])
  215. local expireBefore = now - ttl
  216. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  217. if redis.call('ZSCORE', key, leaseID) == false then
  218. return 0
  219. end
  220. redis.call('ZADD', key, now, leaseID)
  221. redis.call('EXPIRE', key, ttl)
  222. return 1
  223. `)
  224. // incrementWaitScript - refreshes TTL on each increment to keep queue depth accurate
  225. // KEYS[1] = wait queue key
  226. // ARGV[1] = maxWait
  227. // ARGV[2] = TTL in seconds
  228. // 返回 {是否成功, Redis 当前秒},供 Go 侧免额外 TIME 往返写活跃索引。
  229. incrementWaitScript = redis.NewScript(`
  230. -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
  231. -- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
  232. redis.replicate_commands()
  233. local current = redis.call('GET', KEYS[1])
  234. if current == false then
  235. current = 0
  236. else
  237. current = tonumber(current)
  238. end
  239. local now = tonumber(redis.call('TIME')[1])
  240. if current >= tonumber(ARGV[1]) then
  241. return {0, now}
  242. end
  243. redis.call('INCR', KEYS[1])
  244. -- Refresh TTL so long-running traffic doesn't expire active queue counters.
  245. redis.call('EXPIRE', KEYS[1], ARGV[2])
  246. return {1, now}
  247. `)
  248. // incrementAccountWaitScript - account-level wait queue count (refresh TTL on each increment)
  249. // 返回值同 incrementWaitScript:{是否成功, Redis 当前秒}。
  250. incrementAccountWaitScript = redis.NewScript(`
  251. -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
  252. -- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
  253. redis.replicate_commands()
  254. local current = redis.call('GET', KEYS[1])
  255. if current == false then
  256. current = 0
  257. else
  258. current = tonumber(current)
  259. end
  260. local now = tonumber(redis.call('TIME')[1])
  261. if current >= tonumber(ARGV[1]) then
  262. return {0, now}
  263. end
  264. redis.call('INCR', KEYS[1])
  265. -- Refresh TTL so long-running traffic doesn't expire active queue counters.
  266. redis.call('EXPIRE', KEYS[1], ARGV[2])
  267. return {1, now}
  268. `)
  269. // decrementWaitScript - same as before
  270. decrementWaitScript = redis.NewScript(`
  271. local current = redis.call('GET', KEYS[1])
  272. if current ~= false and tonumber(current) > 0 then
  273. redis.call('DECR', KEYS[1])
  274. end
  275. return 1
  276. `)
  277. // cleanupExpiredSlotsScript 清理单个账号/用户有序集合中过期槽位
  278. // KEYS[1] = 有序集合键
  279. // ARGV[1] = TTL(秒)
  280. cleanupExpiredSlotsScript = redis.NewScript(`
  281. -- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
  282. -- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
  283. redis.replicate_commands()
  284. local key = KEYS[1]
  285. local ttl = tonumber(ARGV[1])
  286. local timeResult = redis.call('TIME')
  287. local now = tonumber(timeResult[1])
  288. local expireBefore = now - ttl
  289. redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
  290. if redis.call('ZCARD', key) == 0 then
  291. redis.call('DEL', key)
  292. else
  293. redis.call('EXPIRE', key, ttl)
  294. end
  295. return 1
  296. `)
  297. // startupCleanupSlotScript 清理单个槽位 key 中非当前进程前缀的成员,避免 Redis Cluster CROSSSLOT。
  298. // KEYS[1] 是有序集合键,ARGV[1] 是当前进程前缀,ARGV[2] 是槽位 TTL。
  299. // 返回 {清除数量, 剩余成员数},Go 侧据剩余数决定索引 member 去留,无需再回读槽位。
  300. startupCleanupSlotScript = redis.NewScript(`
  301. local key = KEYS[1]
  302. local activePrefix = ARGV[1]
  303. local slotTTL = tonumber(ARGV[2])
  304. local removed = 0
  305. local members = redis.call('ZRANGE', key, 0, -1)
  306. for _, member in ipairs(members) do
  307. if string.sub(member, 1, string.len(activePrefix)) ~= activePrefix then
  308. removed = removed + redis.call('ZREM', key, member)
  309. end
  310. end
  311. local remaining = redis.call('ZCARD', key)
  312. if remaining == 0 then
  313. redis.call('DEL', key)
  314. else
  315. redis.call('EXPIRE', key, slotTTL)
  316. end
  317. return {removed, remaining}
  318. `)
  319. )
  320. type concurrencyCache struct {
  321. rdb *redis.Client
  322. slotTTLSeconds int // 槽位过期时间(秒)
  323. waitQueueTTLSeconds int // 等待队列过期时间(秒)
  324. }
  325. // NewConcurrencyCache 创建并发控制缓存
  326. // slotTTLMinutes: 槽位过期时间(分钟),0 或负数使用默认值 15 分钟
  327. // waitQueueTTLSeconds: 等待队列过期时间(秒),0 或负数使用 slot TTL
  328. func NewConcurrencyCache(rdb *redis.Client, slotTTLMinutes int, waitQueueTTLSeconds int) service.ConcurrencyCache {
  329. if slotTTLMinutes <= 0 {
  330. slotTTLMinutes = defaultSlotTTLMinutes
  331. }
  332. if waitQueueTTLSeconds <= 0 {
  333. waitQueueTTLSeconds = slotTTLMinutes * 60
  334. }
  335. return &concurrencyCache{
  336. rdb: rdb,
  337. slotTTLSeconds: slotTTLMinutes * 60,
  338. waitQueueTTLSeconds: waitQueueTTLSeconds,
  339. }
  340. }
  341. // Helper functions for key generation
  342. func accountSlotKey(accountID int64) string {
  343. return fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
  344. }
  345. func userSlotKey(userID int64) string {
  346. return fmt.Sprintf("%s%d", userSlotKeyPrefix, userID)
  347. }
  348. func apiKeySlotKey(apiKeyID int64) string {
  349. return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
  350. }
  351. func liveAccountSlotKey(accountID int64) string {
  352. return fmt.Sprintf("%s%d", liveAccountSlotKeyPrefix, accountID)
  353. }
  354. func liveUserSlotKey(userID int64) string {
  355. return fmt.Sprintf("%s%d", liveUserSlotKeyPrefix, userID)
  356. }
  357. func liveAPIKeySlotKey(apiKeyID int64) string {
  358. return fmt.Sprintf("%s%d", liveAPIKeySlotKeyPrefix, apiKeyID)
  359. }
  360. func openAIWSIngressLeaseKey(apiKeyID int64) string {
  361. return fmt.Sprintf("%s%d", openAIWSIngressLeaseKeyPrefix, apiKeyID)
  362. }
  363. func waitQueueKey(userID int64) string {
  364. return fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
  365. }
  366. func accountWaitKey(accountID int64) string {
  367. return fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
  368. }
  369. // redisUnixSeconds 统一使用 Redis 服务器时间,避免多实例本地时钟漂移导致索引提前/延后过期。
  370. func (c *concurrencyCache) redisUnixSeconds(ctx context.Context) (int64, error) {
  371. now, err := c.rdb.Time(ctx).Result()
  372. if err != nil {
  373. return 0, fmt.Errorf("redis TIME: %w", err)
  374. }
  375. return now.Unix(), nil
  376. }
  377. // slotIndexSpec 描述一个活跃索引及其对应的槽位/等待键构造方式。
  378. // 用具名字段避免把 slotKey/waitKey 两个同签名函数按位置传参时写反。
  379. type slotIndexSpec struct {
  380. indexKey string
  381. slotKey func(int64) string
  382. waitKey func(int64) string
  383. }
  384. var (
  385. accountSlotIndex = slotIndexSpec{indexKey: accountActiveIndexKey, slotKey: accountSlotKey, waitKey: accountWaitKey}
  386. userSlotIndex = slotIndexSpec{indexKey: userActiveIndexKey, slotKey: userSlotKey, waitKey: waitQueueKey}
  387. )
  388. // touchActiveIndexAt 是写路径上的轻量标记:主操作已成功时,尽力把 ID 放入活跃索引,
  389. // score 为给定的绝对过期时间(Redis Unix 秒)。索引失败不影响并发槽位/等待队列本身,
  390. // 后续释放或清理会再次校正,因此只记日志不上抛。
  391. func (c *concurrencyCache) touchActiveIndexAt(ctx context.Context, indexKey string, id int64, expireAt int64) {
  392. if c == nil || c.rdb == nil || id <= 0 || expireAt <= 0 {
  393. return
  394. }
  395. if err := c.rdb.ZAdd(ctx, indexKey, redis.Z{
  396. Score: float64(expireAt),
  397. Member: strconv.FormatInt(id, 10),
  398. }).Err(); err != nil {
  399. logger.LegacyPrintf("repository.concurrency", "Warning: touch active index %s for %d failed: %v", indexKey, id, err)
  400. }
  401. }
  402. func (c *concurrencyCache) refreshAccountActiveIndex(ctx context.Context, accountID int64) {
  403. c.refreshActiveIndex(ctx, accountActiveIndexKey, accountID, accountSlotKey(accountID), accountWaitKey(accountID))
  404. }
  405. func (c *concurrencyCache) refreshUserActiveIndex(ctx context.Context, userID int64) {
  406. c.refreshActiveIndex(ctx, userActiveIndexKey, userID, userSlotKey(userID), waitQueueKey(userID))
  407. }
  408. // refreshActiveIndex 以 Redis 中的真实槽位/等待数为准重建索引状态。
  409. // 释放槽位、等待计数减少、清理过期成员后都会调用它,防止索引残留。
  410. // 索引维护是 best-effort:失败只记日志,不影响主流程。
  411. func (c *concurrencyCache) refreshActiveIndex(ctx context.Context, indexKey string, id int64, slotKey, waitKey string) {
  412. if c == nil || c.rdb == nil || id <= 0 {
  413. return
  414. }
  415. now, err := c.redisUnixSeconds(ctx)
  416. if err != nil {
  417. logger.LegacyPrintf("repository.concurrency", "Warning: refresh active index %s for %d failed: %v", indexKey, id, err)
  418. return
  419. }
  420. load, err := c.readActiveLoadForKey(ctx, id, slotKey, waitKey, now)
  421. if err != nil {
  422. logger.LegacyPrintf("repository.concurrency", "Warning: refresh active index %s for %d failed: %v", indexKey, id, err)
  423. return
  424. }
  425. member := strconv.FormatInt(id, 10)
  426. if load.slotCount == 0 && load.waitCount <= 0 {
  427. if err := c.rdb.ZRem(ctx, indexKey, member).Err(); err != nil {
  428. logger.LegacyPrintf("repository.concurrency", "Warning: remove active index member %s from %s failed: %v", member, indexKey, err)
  429. }
  430. return
  431. }
  432. ttlSeconds := c.activeIndexTTL(load.slotCount, load.waitCount)
  433. if ttlSeconds <= 0 {
  434. return
  435. }
  436. c.touchActiveIndexAt(ctx, indexKey, id, now+int64(ttlSeconds))
  437. }
  438. type activeIndexLoad struct {
  439. id int64
  440. member string
  441. slotCount int
  442. waitCount int
  443. }
  444. // activeIndexTTL 取槽位 TTL 与等待队列 TTL 中仍然需要关注的较大值。
  445. // 只要并发槽位或等待计数还有负载,就保留索引;两者都为 0 时调用方会删除索引。
  446. func (c *concurrencyCache) activeIndexTTL(slotCount int, waitCount int) int {
  447. ttlSeconds := 0
  448. if slotCount > 0 {
  449. ttlSeconds = c.slotTTLSeconds
  450. }
  451. if waitCount > 0 && c.waitQueueTTLSeconds > ttlSeconds {
  452. ttlSeconds = c.waitQueueTTLSeconds
  453. }
  454. return ttlSeconds
  455. }
  456. // readActiveLoadForKey 读取单个 ID 的当前负载,并顺手清理该槽位集合中的过期成员。
  457. func (c *concurrencyCache) readActiveLoadForKey(ctx context.Context, id int64, slotKey, waitKey string, now int64) (activeIndexLoad, error) {
  458. cutoffTime := now - int64(c.slotTTLSeconds)
  459. pipe := c.rdb.Pipeline()
  460. pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
  461. zcardCmd := pipe.ZCard(ctx, slotKey)
  462. getCmd := pipe.Get(ctx, waitKey)
  463. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  464. return activeIndexLoad{}, fmt.Errorf("pipeline exec: %w", err)
  465. }
  466. waitCount := 0
  467. if v, err := getCmd.Int(); err == nil && v > 0 {
  468. waitCount = v
  469. }
  470. return activeIndexLoad{
  471. id: id,
  472. member: strconv.FormatInt(id, 10),
  473. slotCount: int(zcardCmd.Val()),
  474. waitCount: waitCount,
  475. }, nil
  476. }
  477. // readIndexLoads 批量读取索引候选的真实负载(账号/用户通用)。
  478. // 分块 Pipeline 可以减少 Redis 往返,同时避免一次 Pipeline 塞入过多命令。
  479. func (c *concurrencyCache) readIndexLoads(ctx context.Context, spec slotIndexSpec, members []string, now int64) ([]activeIndexLoad, []string, error) {
  480. loads := make([]activeIndexLoad, 0, len(members))
  481. staleMembers := make([]string, 0)
  482. candidates := make([]activeIndexLoad, 0, len(members))
  483. for _, member := range members {
  484. id, err := strconv.ParseInt(member, 10, 64)
  485. if err != nil || id <= 0 {
  486. staleMembers = append(staleMembers, member)
  487. continue
  488. }
  489. candidates = append(candidates, activeIndexLoad{id: id, member: member})
  490. }
  491. cutoffTime := now - int64(c.slotTTLSeconds)
  492. for start := 0; start < len(candidates); start += activeIndexPipelineChunkSize {
  493. end := start + activeIndexPipelineChunkSize
  494. if end > len(candidates) {
  495. end = len(candidates)
  496. }
  497. chunk := candidates[start:end]
  498. pipe := c.rdb.Pipeline()
  499. type loadCmd struct {
  500. activeIndexLoad
  501. zcardCmd *redis.IntCmd
  502. getCmd *redis.StringCmd
  503. }
  504. cmds := make([]loadCmd, 0, len(chunk))
  505. for _, candidate := range chunk {
  506. slotKey := spec.slotKey(candidate.id)
  507. waitKey := spec.waitKey(candidate.id)
  508. pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
  509. cmds = append(cmds, loadCmd{
  510. activeIndexLoad: candidate,
  511. zcardCmd: pipe.ZCard(ctx, slotKey),
  512. getCmd: pipe.Get(ctx, waitKey),
  513. })
  514. }
  515. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  516. return nil, nil, fmt.Errorf("pipeline exec: %w", err)
  517. }
  518. for _, cmd := range cmds {
  519. waitCount := 0
  520. if v, err := cmd.getCmd.Int(); err == nil && v > 0 {
  521. waitCount = v
  522. }
  523. loads = append(loads, activeIndexLoad{
  524. id: cmd.id,
  525. member: cmd.member,
  526. slotCount: int(cmd.zcardCmd.Val()),
  527. waitCount: waitCount,
  528. })
  529. }
  530. }
  531. return loads, staleMembers, nil
  532. }
  533. // removeActiveIndexMembers 清理无效 member;这是辅助索引的维护动作,调用方无需因为失败中断主流程。
  534. func (c *concurrencyCache) removeActiveIndexMembers(ctx context.Context, indexKey string, members []string) {
  535. if len(members) == 0 {
  536. return
  537. }
  538. args := make([]any, 0, len(members))
  539. for _, member := range members {
  540. args = append(args, member)
  541. }
  542. if err := c.rdb.ZRem(ctx, indexKey, args...).Err(); err != nil {
  543. logger.LegacyPrintf("repository.concurrency", "Warning: remove %d active index members from %s failed: %v", len(members), indexKey, err)
  544. }
  545. }
  546. // runScriptInt64Pair 执行返回两元素整数数组的 Lua 脚本并解析(如 {result, now}、{removed, remaining})。
  547. func runScriptInt64Pair(ctx context.Context, rdb *redis.Client, script *redis.Script, keys []string, args ...any) (int64, int64, error) {
  548. raw, err := script.Run(ctx, rdb, keys, args...).Result()
  549. if err != nil {
  550. return 0, 0, err
  551. }
  552. first, err := redisScriptInt64At(raw, 0)
  553. if err != nil {
  554. return 0, 0, fmt.Errorf("parse script value 0: %w", err)
  555. }
  556. second, err := redisScriptInt64At(raw, 1)
  557. if err != nil {
  558. return 0, 0, fmt.Errorf("parse script value 1: %w", err)
  559. }
  560. return first, second, nil
  561. }
  562. // Account slot operations
  563. func (c *concurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
  564. key := accountSlotKey(accountID)
  565. // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致
  566. result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key, liveAccountSlotKey(accountID)}, maxConcurrency, c.slotTTLSeconds, requestID)
  567. if err != nil {
  568. return false, err
  569. }
  570. if result == 1 {
  571. // 成功占槽后标记活跃账号,后台清理即可从索引定位候选账号。
  572. c.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.slotTTLSeconds))
  573. }
  574. return result == 1, nil
  575. }
  576. func (c *concurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
  577. key := accountSlotKey(accountID)
  578. if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil {
  579. return err
  580. }
  581. // 释放后用真实负载刷新索引;若没有槽位和等待计数,会移除索引 member。
  582. c.refreshAccountActiveIndex(ctx, accountID)
  583. return nil
  584. }
  585. func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID int64) (int, error) {
  586. key := accountSlotKey(accountID)
  587. // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取
  588. result, err := getCountScript.Run(ctx, c.rdb, []string{key, liveAccountSlotKey(accountID)}, c.slotTTLSeconds).Int()
  589. if err != nil {
  590. return 0, err
  591. }
  592. return result, nil
  593. }
  594. func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error) {
  595. if len(accountIDs) == 0 {
  596. return map[int64]int{}, nil
  597. }
  598. now, err := c.rdb.Time(ctx).Result()
  599. if err != nil {
  600. return nil, fmt.Errorf("redis TIME: %w", err)
  601. }
  602. cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
  603. pipe := c.rdb.Pipeline()
  604. type accountCmd struct {
  605. accountID int64
  606. zcardCmd *redis.IntCmd
  607. liveCmd *redis.IntCmd
  608. }
  609. cmds := make([]accountCmd, 0, len(accountIDs))
  610. for _, accountID := range accountIDs {
  611. slotKey := accountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
  612. liveKey := liveAccountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
  613. pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
  614. pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
  615. cmds = append(cmds, accountCmd{
  616. accountID: accountID,
  617. zcardCmd: pipe.ZCard(ctx, slotKey),
  618. liveCmd: pipe.ZCard(ctx, liveKey),
  619. })
  620. }
  621. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  622. return nil, fmt.Errorf("pipeline exec: %w", err)
  623. }
  624. result := make(map[int64]int, len(accountIDs))
  625. for _, cmd := range cmds {
  626. result[cmd.accountID] = int(cmd.zcardCmd.Val() + cmd.liveCmd.Val())
  627. }
  628. return result, nil
  629. }
  630. // User slot operations
  631. func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
  632. key := userSlotKey(userID)
  633. // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致
  634. result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key, liveUserSlotKey(userID)}, maxConcurrency, c.slotTTLSeconds, requestID)
  635. if err != nil {
  636. return false, err
  637. }
  638. if result == 1 {
  639. // 成功占槽后标记活跃用户,避免启动清理依赖全量 SCAN。
  640. c.touchActiveIndexAt(ctx, userActiveIndexKey, userID, now+int64(c.slotTTLSeconds))
  641. }
  642. return result == 1, nil
  643. }
  644. func (c *concurrencyCache) ReleaseUserSlot(ctx context.Context, userID int64, requestID string) error {
  645. key := userSlotKey(userID)
  646. if err := c.rdb.ZRem(ctx, key, requestID).Err(); err != nil {
  647. return err
  648. }
  649. // 释放后按 Redis 中剩余负载修正索引状态。
  650. c.refreshUserActiveIndex(ctx, userID)
  651. return nil
  652. }
  653. func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) (int, error) {
  654. key := userSlotKey(userID)
  655. // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取
  656. result, err := getCountScript.Run(ctx, c.rdb, []string{key, liveUserSlotKey(userID)}, c.slotTTLSeconds).Int()
  657. if err != nil {
  658. return 0, err
  659. }
  660. return result, nil
  661. }
  662. func (c *concurrencyCache) TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error {
  663. key := apiKeySlotKey(apiKeyID)
  664. _, err := trackSlotScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds, requestID).Result()
  665. return err
  666. }
  667. func (c *concurrencyCache) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error {
  668. key := apiKeySlotKey(apiKeyID)
  669. return c.rdb.ZRem(ctx, key, requestID).Err()
  670. }
  671. func (c *concurrencyCache) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
  672. if c == nil || c.rdb == nil || apiKeyID <= 0 || maxConnections <= 0 || leaseID == "" {
  673. return false, nil
  674. }
  675. result, err := acquireOpenAIWSIngressLeaseScript.Run(
  676. ctx,
  677. c.rdb,
  678. []string{openAIWSIngressLeaseKey(apiKeyID)},
  679. maxConnections,
  680. openAIWSIngressLeaseTTLSeconds,
  681. leaseID,
  682. ).Int()
  683. if err != nil {
  684. return false, err
  685. }
  686. return result == 1, nil
  687. }
  688. func (c *concurrencyCache) RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error) {
  689. if c == nil || c.rdb == nil || apiKeyID <= 0 || leaseID == "" {
  690. return false, nil
  691. }
  692. result, err := refreshOpenAIWSIngressLeaseScript.Run(
  693. ctx,
  694. c.rdb,
  695. []string{openAIWSIngressLeaseKey(apiKeyID)},
  696. openAIWSIngressLeaseTTLSeconds,
  697. leaseID,
  698. ).Int()
  699. if err != nil {
  700. return false, err
  701. }
  702. return result == 1, nil
  703. }
  704. func (c *concurrencyCache) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
  705. if c == nil || c.rdb == nil || apiKeyID <= 0 || leaseID == "" {
  706. return nil
  707. }
  708. return c.rdb.ZRem(ctx, openAIWSIngressLeaseKey(apiKeyID), leaseID).Err()
  709. }
  710. func (c *concurrencyCache) AcquireLiveLease(
  711. ctx context.Context,
  712. accountID int64,
  713. accountMax int,
  714. userID int64,
  715. userMax int,
  716. apiKeyID int64,
  717. leaseID string,
  718. replacingRegularSlots bool,
  719. ) (bool, error) {
  720. if c == nil || c.rdb == nil || accountID <= 0 || userID <= 0 || apiKeyID <= 0 || leaseID == "" {
  721. return false, nil
  722. }
  723. replacing := 0
  724. if replacingRegularSlots {
  725. replacing = 1
  726. }
  727. result, err := acquireLiveLeaseScript.Run(ctx, c.rdb, []string{
  728. accountSlotKey(accountID),
  729. liveAccountSlotKey(accountID),
  730. userSlotKey(userID),
  731. liveUserSlotKey(userID),
  732. liveAPIKeySlotKey(apiKeyID),
  733. }, accountMax, userMax, liveLeaseTTLSeconds, leaseID, replacing).Int()
  734. return result == 1, err
  735. }
  736. func (c *concurrencyCache) RefreshLiveLease(ctx context.Context, accountID, userID, apiKeyID int64, leaseID string) (bool, error) {
  737. if c == nil || c.rdb == nil || leaseID == "" {
  738. return false, nil
  739. }
  740. result, err := refreshLiveLeaseScript.Run(ctx, c.rdb, []string{
  741. liveAccountSlotKey(accountID),
  742. liveUserSlotKey(userID),
  743. liveAPIKeySlotKey(apiKeyID),
  744. }, liveLeaseTTLSeconds, leaseID).Int()
  745. return result == 1, err
  746. }
  747. func (c *concurrencyCache) ReleaseLiveLease(ctx context.Context, accountID, userID, apiKeyID int64, leaseID string) error {
  748. if c == nil || c.rdb == nil || leaseID == "" {
  749. return nil
  750. }
  751. pipe := c.rdb.TxPipeline()
  752. pipe.ZRem(ctx, liveAccountSlotKey(accountID), leaseID)
  753. pipe.ZRem(ctx, liveUserSlotKey(userID), leaseID)
  754. pipe.ZRem(ctx, liveAPIKeySlotKey(apiKeyID), leaseID)
  755. _, err := pipe.Exec(ctx)
  756. return err
  757. }
  758. func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) {
  759. if len(apiKeyIDs) == 0 {
  760. return map[int64]int{}, nil
  761. }
  762. now, err := c.rdb.Time(ctx).Result()
  763. if err != nil {
  764. return nil, fmt.Errorf("redis TIME: %w", err)
  765. }
  766. cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
  767. pipe := c.rdb.Pipeline()
  768. type apiKeyCmd struct {
  769. apiKeyID int64
  770. zcardCmd *redis.IntCmd
  771. liveCmd *redis.IntCmd
  772. }
  773. cmds := make([]apiKeyCmd, 0, len(apiKeyIDs))
  774. for _, apiKeyID := range apiKeyIDs {
  775. slotKey := apiKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10)
  776. liveKey := liveAPIKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10)
  777. pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
  778. pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
  779. cmds = append(cmds, apiKeyCmd{
  780. apiKeyID: apiKeyID,
  781. zcardCmd: pipe.ZCard(ctx, slotKey),
  782. liveCmd: pipe.ZCard(ctx, liveKey),
  783. })
  784. }
  785. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  786. return nil, fmt.Errorf("pipeline exec: %w", err)
  787. }
  788. result := make(map[int64]int, len(apiKeyIDs))
  789. for _, cmd := range cmds {
  790. result[cmd.apiKeyID] = int(cmd.zcardCmd.Val() + cmd.liveCmd.Val())
  791. }
  792. return result, nil
  793. }
  794. // Wait queue operations
  795. func (c *concurrencyCache) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) {
  796. key := waitQueueKey(userID)
  797. result, now, err := runScriptInt64Pair(ctx, c.rdb, incrementWaitScript, []string{key}, maxWait, c.waitQueueTTLSeconds)
  798. if err != nil {
  799. return false, err
  800. }
  801. if result == 1 {
  802. // 等待队列也会让用户保持“活跃”,否则槽位为 0 时后台任务可能漏看等待计数。
  803. c.touchActiveIndexAt(ctx, userActiveIndexKey, userID, now+int64(c.waitQueueTTLSeconds))
  804. }
  805. return result == 1, nil
  806. }
  807. func (c *concurrencyCache) DecrementWaitCount(ctx context.Context, userID int64) error {
  808. key := waitQueueKey(userID)
  809. _, err := decrementWaitScript.Run(ctx, c.rdb, []string{key}).Result()
  810. if err == nil {
  811. // 等待数减少后重新判断是否还需要保留索引。
  812. c.refreshUserActiveIndex(ctx, userID)
  813. }
  814. return err
  815. }
  816. // Account wait queue operations
  817. func (c *concurrencyCache) IncrementAccountWaitCount(ctx context.Context, accountID int64, maxWait int) (bool, error) {
  818. key := accountWaitKey(accountID)
  819. result, now, err := runScriptInt64Pair(ctx, c.rdb, incrementAccountWaitScript, []string{key}, maxWait, c.waitQueueTTLSeconds)
  820. if err != nil {
  821. return false, err
  822. }
  823. if result == 1 {
  824. // 账号级等待队列同样写入账号活跃索引,供负载查询和清理任务使用。
  825. c.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.waitQueueTTLSeconds))
  826. }
  827. return result == 1, nil
  828. }
  829. func (c *concurrencyCache) DecrementAccountWaitCount(ctx context.Context, accountID int64) error {
  830. key := accountWaitKey(accountID)
  831. _, err := decrementWaitScript.Run(ctx, c.rdb, []string{key}).Result()
  832. if err == nil {
  833. // 等待计数归零后索引需要同步删除,避免后台任务反复处理空账号。
  834. c.refreshAccountActiveIndex(ctx, accountID)
  835. }
  836. return err
  837. }
  838. func (c *concurrencyCache) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
  839. key := accountWaitKey(accountID)
  840. val, err := c.rdb.Get(ctx, key).Int()
  841. if err != nil && !errors.Is(err, redis.Nil) {
  842. return 0, err
  843. }
  844. if errors.Is(err, redis.Nil) {
  845. return 0, nil
  846. }
  847. return val, nil
  848. }
  849. func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []service.AccountWithConcurrency) (map[int64]*service.AccountLoadInfo, error) {
  850. if len(accounts) == 0 {
  851. return map[int64]*service.AccountLoadInfo{}, nil
  852. }
  853. // 使用 Pipeline 替代 Lua 脚本,兼容 Redis Cluster(Lua 内动态拼 key 会 CROSSSLOT)。
  854. // 每个账号执行 3 个命令:ZREMRANGEBYSCORE(清理过期)、ZCARD(并发数)、GET(等待数)。
  855. now, err := c.rdb.Time(ctx).Result()
  856. if err != nil {
  857. return nil, fmt.Errorf("redis TIME: %w", err)
  858. }
  859. cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
  860. pipe := c.rdb.Pipeline()
  861. type accountCmds struct {
  862. id int64
  863. maxConcurrency int
  864. zcardCmd *redis.IntCmd
  865. liveCmd *redis.IntCmd
  866. getCmd *redis.StringCmd
  867. }
  868. cmds := make([]accountCmds, 0, len(accounts))
  869. for _, acc := range accounts {
  870. slotKey := accountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
  871. liveKey := liveAccountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
  872. waitKey := accountWaitKeyPrefix + strconv.FormatInt(acc.ID, 10)
  873. pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
  874. pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
  875. ac := accountCmds{
  876. id: acc.ID,
  877. maxConcurrency: acc.MaxConcurrency,
  878. zcardCmd: pipe.ZCard(ctx, slotKey),
  879. liveCmd: pipe.ZCard(ctx, liveKey),
  880. getCmd: pipe.Get(ctx, waitKey),
  881. }
  882. cmds = append(cmds, ac)
  883. }
  884. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  885. return nil, fmt.Errorf("pipeline exec: %w", err)
  886. }
  887. loadMap := make(map[int64]*service.AccountLoadInfo, len(accounts))
  888. for _, ac := range cmds {
  889. currentConcurrency := int(ac.zcardCmd.Val() + ac.liveCmd.Val())
  890. waitingCount := 0
  891. if v, err := ac.getCmd.Int(); err == nil {
  892. waitingCount = v
  893. }
  894. loadRate := 0
  895. if ac.maxConcurrency > 0 {
  896. loadRate = (currentConcurrency + waitingCount) * 100 / ac.maxConcurrency
  897. }
  898. loadMap[ac.id] = &service.AccountLoadInfo{
  899. AccountID: ac.id,
  900. CurrentConcurrency: currentConcurrency,
  901. WaitingCount: waitingCount,
  902. LoadRate: loadRate,
  903. }
  904. }
  905. return loadMap, nil
  906. }
  907. func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []service.UserWithConcurrency) (map[int64]*service.UserLoadInfo, error) {
  908. if len(users) == 0 {
  909. return map[int64]*service.UserLoadInfo{}, nil
  910. }
  911. // 使用 Pipeline 替代 Lua 脚本,兼容 Redis Cluster。
  912. now, err := c.rdb.Time(ctx).Result()
  913. if err != nil {
  914. return nil, fmt.Errorf("redis TIME: %w", err)
  915. }
  916. cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
  917. pipe := c.rdb.Pipeline()
  918. type userCmds struct {
  919. id int64
  920. maxConcurrency int
  921. zcardCmd *redis.IntCmd
  922. liveCmd *redis.IntCmd
  923. getCmd *redis.StringCmd
  924. }
  925. cmds := make([]userCmds, 0, len(users))
  926. for _, u := range users {
  927. slotKey := userSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
  928. liveKey := liveUserSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
  929. waitKey := waitQueueKeyPrefix + strconv.FormatInt(u.ID, 10)
  930. pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
  931. pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
  932. uc := userCmds{
  933. id: u.ID,
  934. maxConcurrency: u.MaxConcurrency,
  935. zcardCmd: pipe.ZCard(ctx, slotKey),
  936. liveCmd: pipe.ZCard(ctx, liveKey),
  937. getCmd: pipe.Get(ctx, waitKey),
  938. }
  939. cmds = append(cmds, uc)
  940. }
  941. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  942. return nil, fmt.Errorf("pipeline exec: %w", err)
  943. }
  944. loadMap := make(map[int64]*service.UserLoadInfo, len(users))
  945. for _, uc := range cmds {
  946. currentConcurrency := int(uc.zcardCmd.Val() + uc.liveCmd.Val())
  947. waitingCount := 0
  948. if v, err := uc.getCmd.Int(); err == nil {
  949. waitingCount = v
  950. }
  951. loadRate := 0
  952. if uc.maxConcurrency > 0 {
  953. loadRate = (currentConcurrency + waitingCount) * 100 / uc.maxConcurrency
  954. }
  955. loadMap[uc.id] = &service.UserLoadInfo{
  956. UserID: uc.id,
  957. CurrentConcurrency: currentConcurrency,
  958. WaitingCount: waitingCount,
  959. LoadRate: loadRate,
  960. }
  961. }
  962. return loadMap, nil
  963. }
  964. func (c *concurrencyCache) CleanupExpiredAccountSlots(ctx context.Context, accountID int64) error {
  965. key := accountSlotKey(accountID)
  966. _, err := cleanupExpiredSlotsScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Result()
  967. if err == nil {
  968. // 单账号清理后同步索引,保持后台批量清理的候选集准确。
  969. c.refreshAccountActiveIndex(ctx, accountID)
  970. }
  971. return err
  972. }
  973. // CleanupExpiredAccountSlotKeys 处理账号与用户两个活跃索引中已到期的候选。
  974. // (方法名中的 Account 是历史遗留,保留以避免接口变更;实际同时回收两个索引,
  975. // 否则 user 索引的过期成员没有任何清理路径,会无界累积。)
  976. func (c *concurrencyCache) CleanupExpiredAccountSlotKeys(ctx context.Context) error {
  977. if err := c.reconcileExpiredIndexCandidates(ctx, accountSlotIndex); err != nil {
  978. return err
  979. }
  980. return c.reconcileExpiredIndexCandidates(ctx, userSlotIndex)
  981. }
  982. // reconcileExpiredIndexCandidates 处理单个活跃索引中 score 已到期的候选:
  983. // 无真实负载则移除 member;仍有负载则按真实负载批量刷新 score。
  984. func (c *concurrencyCache) reconcileExpiredIndexCandidates(ctx context.Context, spec slotIndexSpec) error {
  985. now, err := c.redisUnixSeconds(ctx)
  986. if err != nil {
  987. return err
  988. }
  989. members, err := c.rdb.ZRangeArgs(ctx, redis.ZRangeArgs{
  990. Key: spec.indexKey,
  991. Start: "-inf",
  992. Stop: strconv.FormatInt(now, 10),
  993. ByScore: true,
  994. Count: activeIndexCleanupBatchSize,
  995. }).Result()
  996. if err != nil {
  997. return fmt.Errorf("read expired index %s: %w", spec.indexKey, err)
  998. }
  999. loads, staleMembers, err := c.readIndexLoads(ctx, spec, members, now)
  1000. if err != nil {
  1001. return err
  1002. }
  1003. refreshed := make([]redis.Z, 0, len(loads))
  1004. for _, load := range loads {
  1005. if load.slotCount == 0 && load.waitCount <= 0 {
  1006. // 真实槽位和等待数都为空,说明这个索引 member 已经完成使命。
  1007. staleMembers = append(staleMembers, load.member)
  1008. continue
  1009. }
  1010. refreshed = append(refreshed, redis.Z{
  1011. Score: float64(now + int64(c.activeIndexTTL(load.slotCount, load.waitCount))),
  1012. Member: load.member,
  1013. })
  1014. }
  1015. if len(refreshed) > 0 {
  1016. if err := c.rdb.ZAdd(ctx, spec.indexKey, refreshed...).Err(); err != nil {
  1017. logger.LegacyPrintf("repository.concurrency", "Warning: refresh %d active index members in %s failed: %v", len(refreshed), spec.indexKey, err)
  1018. }
  1019. }
  1020. c.removeActiveIndexMembers(ctx, spec.indexKey, staleMembers)
  1021. return nil
  1022. }
  1023. // CleanupStaleProcessSlots 启动时清理非当前进程前缀的槽位。
  1024. // 清理范围来自活跃索引(含 score 已过期的成员——它们往往正是崩溃进程留下的残留),
  1025. // 避免在 Redis 上 SCAN 全部 concurrency:* 键;另有一次性迁移清扫兜底索引机制上线前的遗留等待计数。
  1026. // API Key 槽位(concurrency:api_key:*)是 stats-only 数据:每次 Track/读取都会按分数
  1027. // 裁剪过期成员,key 自带 TTL,可在一个 slot TTL 内自愈,因此不参与启动清理。
  1028. func (c *concurrencyCache) CleanupStaleProcessSlots(ctx context.Context, activeRequestPrefix string) error {
  1029. if activeRequestPrefix == "" {
  1030. return nil
  1031. }
  1032. if err := c.sweepLegacyWaitKeysOnce(ctx); err != nil {
  1033. return err
  1034. }
  1035. now, err := c.redisUnixSeconds(ctx)
  1036. if err != nil {
  1037. return err
  1038. }
  1039. accountMembers, err := c.allIndexMembers(ctx, accountActiveIndexKey)
  1040. if err != nil {
  1041. return err
  1042. }
  1043. if err := c.cleanupStaleProcessSlotsForIndex(ctx, accountSlotIndex, accountMembers, activeRequestPrefix, now); err != nil {
  1044. return err
  1045. }
  1046. userMembers, err := c.allIndexMembers(ctx, userActiveIndexKey)
  1047. if err != nil {
  1048. return err
  1049. }
  1050. return c.cleanupStaleProcessSlotsForIndex(ctx, userSlotIndex, userMembers, activeRequestPrefix, now)
  1051. }
  1052. // sweepLegacyWaitKeysOnce 一次性清扫活跃索引机制上线前遗留的等待计数键。
  1053. // 等待计数在有流量时会不断刷新 TTL、无法自然过期,而索引不认识旧键,
  1054. // 因此这里例外地做一次 SCAN,用 marker 键保证整个 Redis 数据生命周期内只执行一次。
  1055. // 先清扫后写 marker:清扫失败时下次启动会重试;并发实例重复清扫是幂等的。
  1056. func (c *concurrencyCache) sweepLegacyWaitKeysOnce(ctx context.Context) error {
  1057. exists, err := c.rdb.Exists(ctx, legacyWaitSweepMarkerKey).Result()
  1058. if err != nil {
  1059. return fmt.Errorf("check legacy wait sweep marker: %w", err)
  1060. }
  1061. if exists > 0 {
  1062. return nil
  1063. }
  1064. for _, pattern := range []string{accountWaitKeyPrefix + "*", waitQueueKeyPrefix + "*"} {
  1065. var cursor uint64
  1066. for {
  1067. keys, next, err := c.rdb.Scan(ctx, cursor, pattern, 200).Result()
  1068. if err != nil {
  1069. return fmt.Errorf("scan legacy wait keys %s: %w", pattern, err)
  1070. }
  1071. if len(keys) > 0 {
  1072. if err := c.rdb.Del(ctx, keys...).Err(); err != nil {
  1073. return fmt.Errorf("delete legacy wait keys: %w", err)
  1074. }
  1075. }
  1076. cursor = next
  1077. if cursor == 0 {
  1078. break
  1079. }
  1080. }
  1081. }
  1082. if err := c.rdb.Set(ctx, legacyWaitSweepMarkerKey, "1", 0).Err(); err != nil {
  1083. return fmt.Errorf("set legacy wait sweep marker: %w", err)
  1084. }
  1085. return nil
  1086. }
  1087. // allIndexMembers 返回索引中全部 member(含 score 已过期的)。
  1088. // 启动清理必须覆盖过期成员:长时间停机后 score 过期的候选恰恰最可能持有死进程残留。
  1089. func (c *concurrencyCache) allIndexMembers(ctx context.Context, indexKey string) ([]string, error) {
  1090. members, err := c.rdb.ZRange(ctx, indexKey, 0, -1).Result()
  1091. if err != nil {
  1092. return nil, fmt.Errorf("read active index %s: %w", indexKey, err)
  1093. }
  1094. return members, nil
  1095. }
  1096. // cleanupStaleProcessSlotsForIndex 逐个处理索引中的账号/用户。
  1097. // Lua 脚本一次只碰一个槽位 key,兼容 Redis Cluster,随后删除重启后已失效的等待计数;
  1098. // 索引 member 的去留由脚本返回的剩余槽位数决定,最后批量写回。
  1099. func (c *concurrencyCache) cleanupStaleProcessSlotsForIndex(
  1100. ctx context.Context,
  1101. spec slotIndexSpec,
  1102. members []string,
  1103. activeRequestPrefix string,
  1104. now int64,
  1105. ) error {
  1106. staleMembers := make([]string, 0)
  1107. refreshed := make([]redis.Z, 0)
  1108. for _, member := range members {
  1109. id, err := strconv.ParseInt(member, 10, 64)
  1110. if err != nil || id <= 0 {
  1111. staleMembers = append(staleMembers, member)
  1112. continue
  1113. }
  1114. _, remaining, err := runScriptInt64Pair(ctx, c.rdb, startupCleanupSlotScript, []string{spec.slotKey(id)}, activeRequestPrefix, c.slotTTLSeconds)
  1115. if err != nil {
  1116. return fmt.Errorf("cleanup stale process slots %s: %w", spec.slotKey(id), err)
  1117. }
  1118. // 等待计数属于已死进程,直接删除;剩余槽位(当前进程前缀)决定索引 member 去留。
  1119. if err := c.rdb.Del(ctx, spec.waitKey(id)).Err(); err != nil {
  1120. return fmt.Errorf("delete stale wait key %s: %w", spec.waitKey(id), err)
  1121. }
  1122. if remaining > 0 {
  1123. refreshed = append(refreshed, redis.Z{
  1124. Score: float64(now + int64(c.slotTTLSeconds)),
  1125. Member: member,
  1126. })
  1127. } else {
  1128. staleMembers = append(staleMembers, member)
  1129. }
  1130. }
  1131. if len(refreshed) > 0 {
  1132. if err := c.rdb.ZAdd(ctx, spec.indexKey, refreshed...).Err(); err != nil {
  1133. logger.LegacyPrintf("repository.concurrency", "Warning: refresh %d active index members in %s failed: %v", len(refreshed), spec.indexKey, err)
  1134. }
  1135. }
  1136. c.removeActiveIndexMembers(ctx, spec.indexKey, staleMembers)
  1137. return nil
  1138. }