scheduler_cache.go 34 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094
  1. package repository
  2. import (
  3. "context"
  4. "crypto/rand"
  5. "encoding/hex"
  6. "encoding/json"
  7. "fmt"
  8. "log/slog"
  9. "strconv"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. "github.com/redis/go-redis/v9"
  13. )
  14. const (
  15. schedulerBucketSetKey = "sched:buckets"
  16. schedulerOutboxWatermarkKey = "sched:outbox:watermark"
  17. schedulerAccountPrefix = "sched:acc:"
  18. schedulerAccountMetaPrefix = "sched:meta:"
  19. schedulerAccountLastUsedPrefix = "sched:acc:last_used:"
  20. schedulerActivePrefix = "sched:active:"
  21. schedulerReadyPrefix = "sched:ready:"
  22. schedulerVersionPrefix = "sched:ver:"
  23. schedulerEpochPrefix = "sched:epoch:"
  24. schedulerRetiredPrefix = "sched:retired:"
  25. schedulerSnapshotPrefix = "sched:"
  26. schedulerLockPrefix = "sched:lock:"
  27. defaultSchedulerSnapshotMGetChunkSize = 128
  28. defaultSchedulerSnapshotWriteChunkSize = 256
  29. schedulerLastUsedUpdateChunkSize = 256
  30. // snapshotGraceTTLSeconds 旧快照过期的宽限期(秒)。
  31. // 替代立即 DEL,让正在读取旧版本的 reader 有足够时间完成 ZRANGE。
  32. snapshotGraceTTLSeconds = 60
  33. )
  34. const (
  35. schedulerGroupLifecycleLockPrefix = "sched:group:lifecycle-lock:"
  36. schedulerGroupLifecycleOwnerTokenBytes = 16
  37. )
  38. var updateSchedulerLastUsedScript = redis.NewScript(`
  39. local updated = 0
  40. for index = 1, #ARGV do
  41. local key_index = (index - 1) * 2 + 1
  42. local candidate = tonumber(ARGV[index])
  43. if candidate == nil then
  44. return redis.error_reply('invalid last_used value')
  45. end
  46. if redis.call('EXISTS', KEYS[key_index]) == 1 then
  47. local current = tonumber(redis.call('GET', KEYS[key_index + 1]))
  48. if current == nil or candidate > current then
  49. redis.call('SET', KEYS[key_index + 1], ARGV[index])
  50. updated = updated + 1
  51. end
  52. end
  53. end
  54. return updated
  55. `)
  56. var (
  57. // epoch 标识 bucket writer 的代际,retired key 是持久退休标记。
  58. // Capture、allocate、activate 都在 Lua 内同时校验两者:-1 表示已退休,-2 表示 epoch 无效或与 token 代际不匹配;
  59. // allocate 与 activate 的双重校验可拦截快照写入期间发生的 Retire。
  60. // Retire 仅在首次退休时推进 epoch,Reopen 只清除标记并沿用该代际,因此重复调用保持幂等。
  61. captureBucketWriteTokenScript = redis.NewScript(`
  62. if redis.call('EXISTS', KEYS[2]) == 1 then
  63. return -1
  64. end
  65. local currentEpoch = redis.call('GET', KEYS[1])
  66. if currentEpoch == false then
  67. redis.call('SET', KEYS[1], '1')
  68. return 1
  69. end
  70. local parsedEpoch = tonumber(currentEpoch)
  71. if parsedEpoch == nil or parsedEpoch < 1 then
  72. return -2
  73. end
  74. return parsedEpoch
  75. `)
  76. allocateSnapshotVersionScript = redis.NewScript(`
  77. if redis.call('EXISTS', KEYS[2]) == 1 then
  78. return -1
  79. end
  80. local currentEpoch = tonumber(redis.call('GET', KEYS[1]))
  81. local expectedEpoch = tonumber(ARGV[1])
  82. if currentEpoch == nil or expectedEpoch == nil or currentEpoch ~= expectedEpoch then
  83. return -2
  84. end
  85. return redis.call('INCR', KEYS[3])
  86. `)
  87. retireBucketScript = redis.NewScript(`
  88. local retired = redis.call('GET', KEYS[2])
  89. local currentEpoch = tonumber(redis.call('GET', KEYS[1])) or 0
  90. if retired == false then
  91. currentEpoch = currentEpoch + 1
  92. if currentEpoch < 1 then
  93. currentEpoch = 1
  94. end
  95. redis.call('SET', KEYS[1], tostring(currentEpoch))
  96. redis.call('SET', KEYS[2], tostring(currentEpoch))
  97. elseif currentEpoch < 1 then
  98. currentEpoch = tonumber(retired) or 1
  99. redis.call('SET', KEYS[1], tostring(currentEpoch))
  100. end
  101. redis.call('SREM', KEYS[3], ARGV[1])
  102. local currentActive = redis.call('GET', KEYS[5])
  103. if currentActive ~= false then
  104. redis.call('EXPIRE', ARGV[2] .. currentActive, tonumber(ARGV[3]))
  105. end
  106. redis.call('DEL', KEYS[4], KEYS[5])
  107. return currentEpoch
  108. `)
  109. reopenBucketScript = redis.NewScript(`
  110. local currentEpochRaw = redis.call('GET', KEYS[1])
  111. local currentEpoch = tonumber(currentEpochRaw)
  112. local retiredEpochRaw = redis.call('GET', KEYS[2])
  113. if retiredEpochRaw == false then
  114. if currentEpochRaw == false then
  115. redis.call('SET', KEYS[1], '1')
  116. return 1
  117. end
  118. if currentEpoch == nil or currentEpoch < 1 then
  119. return -2
  120. end
  121. return currentEpoch
  122. end
  123. local retiredEpoch = tonumber(retiredEpochRaw)
  124. if retiredEpoch == nil or retiredEpoch < 1 then
  125. return -2
  126. end
  127. if currentEpoch == nil or currentEpoch < retiredEpoch then
  128. currentEpoch = retiredEpoch
  129. end
  130. redis.call('SET', KEYS[1], tostring(currentEpoch))
  131. redis.call('DEL', KEYS[2])
  132. redis.call('SREM', KEYS[3], ARGV[1])
  133. local currentActive = redis.call('GET', KEYS[5])
  134. if currentActive ~= false then
  135. redis.call('EXPIRE', ARGV[2] .. currentActive, tonumber(ARGV[3]))
  136. end
  137. redis.call('DEL', KEYS[4], KEYS[5])
  138. return currentEpoch
  139. `)
  140. // 释放租约必须先比较所有者令牌再删除,过期持有者的延迟释放不能误删继任租约。
  141. releaseGroupLifecycleLeaseScript = redis.NewScript(`
  142. if redis.call('GET', KEYS[1]) == ARGV[1] then
  143. return redis.call('DEL', KEYS[1])
  144. end
  145. return 0
  146. `)
  147. // activateSnapshotScript 原子 CAS 切换快照版本。
  148. // 仅当新版本号 >= 当前激活版本时才切换,防止并发写入导致版本回滚。
  149. // 旧快照使用 EXPIRE 设置宽限期而非立即 DEL,避免与 reader 竞态。
  150. //
  151. // KEYS[1] = activeKey (sched:active:{bucket})
  152. // KEYS[2] = readyKey (sched:ready:{bucket})
  153. // KEYS[3] = bucketSetKey (sched:buckets)
  154. // KEYS[4] = snapshotKey (新写入的快照 key)
  155. // KEYS[5] = epochKey
  156. // KEYS[6] = retiredKey
  157. // ARGV[1] = 新版本号字符串
  158. // ARGV[2] = bucket 字符串 (用于 SADD)
  159. // ARGV[3] = 快照 key 前缀 (用于构造旧快照 key)
  160. // ARGV[4] = 宽限期 TTL 秒数
  161. // ARGV[5] = writer epoch
  162. //
  163. // 返回 1 = 已激活, 0 = 版本过旧未激活
  164. activateSnapshotScript = redis.NewScript(`
  165. if redis.call('EXISTS', KEYS[6]) == 1 then
  166. redis.call('DEL', KEYS[4])
  167. return -1
  168. end
  169. local currentEpoch = tonumber(redis.call('GET', KEYS[5]))
  170. local expectedEpoch = tonumber(ARGV[5])
  171. if currentEpoch == nil or expectedEpoch == nil or currentEpoch ~= expectedEpoch then
  172. redis.call('DEL', KEYS[4])
  173. return -2
  174. end
  175. local currentActive = redis.call('GET', KEYS[1])
  176. local newVersion = tonumber(ARGV[1])
  177. if currentActive ~= false then
  178. local curVersion = tonumber(currentActive)
  179. if curVersion and newVersion < curVersion then
  180. redis.call('DEL', KEYS[4])
  181. return 0
  182. end
  183. end
  184. redis.call('SET', KEYS[1], ARGV[1])
  185. redis.call('SET', KEYS[2], '1')
  186. redis.call('SADD', KEYS[3], ARGV[2])
  187. if currentActive ~= false and currentActive ~= ARGV[1] then
  188. redis.call('EXPIRE', ARGV[3] .. currentActive, tonumber(ARGV[4]))
  189. end
  190. return 1
  191. `)
  192. )
  193. type schedulerCache struct {
  194. rdb *redis.Client
  195. mgetChunkSize int
  196. writeChunkSize int
  197. }
  198. func NewSchedulerCache(rdb *redis.Client) service.SchedulerCache {
  199. return newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize)
  200. }
  201. func newSchedulerCacheWithChunkSizes(rdb *redis.Client, mgetChunkSize, writeChunkSize int) service.SchedulerCache {
  202. if mgetChunkSize <= 0 {
  203. mgetChunkSize = defaultSchedulerSnapshotMGetChunkSize
  204. }
  205. if writeChunkSize <= 0 {
  206. writeChunkSize = defaultSchedulerSnapshotWriteChunkSize
  207. }
  208. return &schedulerCache{
  209. rdb: rdb,
  210. mgetChunkSize: mgetChunkSize,
  211. writeChunkSize: writeChunkSize,
  212. }
  213. }
  214. func (c *schedulerCache) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
  215. readyKey := schedulerBucketKey(schedulerReadyPrefix, bucket)
  216. readyVal, err := c.rdb.Get(ctx, readyKey).Result()
  217. if err == redis.Nil {
  218. return nil, false, nil
  219. }
  220. if err != nil {
  221. return nil, false, err
  222. }
  223. if readyVal != "1" {
  224. return nil, false, nil
  225. }
  226. activeKey := schedulerBucketKey(schedulerActivePrefix, bucket)
  227. activeVal, err := c.rdb.Get(ctx, activeKey).Result()
  228. if err == redis.Nil {
  229. return nil, false, nil
  230. }
  231. if err != nil {
  232. return nil, false, err
  233. }
  234. snapshotKey := schedulerSnapshotKey(bucket, activeVal)
  235. ids, err := c.rdb.ZRange(ctx, snapshotKey, 0, -1).Result()
  236. if err != nil {
  237. return nil, false, err
  238. }
  239. if len(ids) == 0 {
  240. // 空快照视为缓存未命中,触发数据库回退查询
  241. // 这解决了新分组创建后立即绑定账号时的竞态条件问题
  242. return nil, false, nil
  243. }
  244. keys := make([]string, 0, len(ids))
  245. lastUsedKeys := make([]string, 0, len(ids))
  246. for _, id := range ids {
  247. keys = append(keys, schedulerAccountMetaKey(id))
  248. lastUsedKeys = append(lastUsedKeys, schedulerLastUsedKey(id))
  249. }
  250. values, err := c.mgetChunked(ctx, keys)
  251. if err != nil {
  252. return nil, false, err
  253. }
  254. lastUsedValues, err := c.mgetChunked(ctx, lastUsedKeys)
  255. if err != nil {
  256. return nil, false, err
  257. }
  258. accounts := make([]*service.Account, 0, len(values))
  259. for i, val := range values {
  260. if val == nil {
  261. return nil, false, nil
  262. }
  263. account, err := decodeCachedAccount(val)
  264. if err != nil {
  265. return nil, false, err
  266. }
  267. if err := applySchedulerLastUsed(account, lastUsedValues[i]); err != nil {
  268. return nil, false, err
  269. }
  270. accounts = append(accounts, account)
  271. }
  272. return accounts, true, nil
  273. }
  274. func (c *schedulerCache) CaptureBucketWriteToken(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
  275. result, err := captureBucketWriteTokenScript.Run(ctx, c.rdb, []string{
  276. schedulerBucketKey(schedulerEpochPrefix, bucket),
  277. schedulerBucketKey(schedulerRetiredPrefix, bucket),
  278. }).Int64()
  279. if err != nil {
  280. return service.SchedulerBucketWriteToken{}, err
  281. }
  282. if err := schedulerBucketWriteResultError(result, bucket); err != nil {
  283. return service.SchedulerBucketWriteToken{}, err
  284. }
  285. return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: result}, nil
  286. }
  287. func (c *schedulerCache) RetireBucket(ctx context.Context, bucket service.SchedulerBucket) error {
  288. snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode)
  289. result, err := retireBucketScript.Run(ctx, c.rdb, []string{
  290. schedulerBucketKey(schedulerEpochPrefix, bucket),
  291. schedulerBucketKey(schedulerRetiredPrefix, bucket),
  292. schedulerBucketSetKey,
  293. schedulerBucketKey(schedulerReadyPrefix, bucket),
  294. schedulerBucketKey(schedulerActivePrefix, bucket),
  295. }, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds).Int64()
  296. if err != nil {
  297. return err
  298. }
  299. if result < 1 {
  300. return fmt.Errorf("retire scheduler bucket %s returned invalid epoch %d", bucket.String(), result)
  301. }
  302. return nil
  303. }
  304. func (c *schedulerCache) ReopenBucket(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
  305. snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode)
  306. result, err := reopenBucketScript.Run(ctx, c.rdb, []string{
  307. schedulerBucketKey(schedulerEpochPrefix, bucket),
  308. schedulerBucketKey(schedulerRetiredPrefix, bucket),
  309. schedulerBucketSetKey,
  310. schedulerBucketKey(schedulerReadyPrefix, bucket),
  311. schedulerBucketKey(schedulerActivePrefix, bucket),
  312. }, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds).Int64()
  313. if err != nil {
  314. return service.SchedulerBucketWriteToken{}, err
  315. }
  316. if err := schedulerBucketWriteResultError(result, bucket); err != nil {
  317. return service.SchedulerBucketWriteToken{}, err
  318. }
  319. return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: result}, nil
  320. }
  321. func (c *schedulerCache) TryAcquireGroupLifecycleLease(ctx context.Context, groupID int64, ttl time.Duration) (service.SchedulerGroupLifecycleLease, bool, error) {
  322. if groupID <= 0 {
  323. return service.SchedulerGroupLifecycleLease{}, false, fmt.Errorf("%w: group id must be positive", service.ErrSchedulerGroupLifecycleLeaseInvalid)
  324. }
  325. if ttl <= 0 {
  326. return service.SchedulerGroupLifecycleLease{}, false, fmt.Errorf("%w: ttl must be positive", service.ErrSchedulerGroupLifecycleLeaseInvalid)
  327. }
  328. ownerToken, err := newSchedulerGroupLifecycleOwnerToken()
  329. if err != nil {
  330. return service.SchedulerGroupLifecycleLease{}, false, err
  331. }
  332. acquired, err := c.rdb.SetNX(ctx, schedulerGroupLifecycleLockKey(groupID), ownerToken, ttl).Result()
  333. if err != nil {
  334. return service.SchedulerGroupLifecycleLease{}, false, err
  335. }
  336. if !acquired {
  337. return service.SchedulerGroupLifecycleLease{}, false, nil
  338. }
  339. return service.SchedulerGroupLifecycleLease{GroupID: groupID, OwnerToken: ownerToken}, true, nil
  340. }
  341. func (c *schedulerCache) ReleaseGroupLifecycleLease(ctx context.Context, lease service.SchedulerGroupLifecycleLease) error {
  342. if !lease.ValidFor(lease.GroupID) {
  343. return service.ErrSchedulerGroupLifecycleLeaseInvalid
  344. }
  345. result, err := releaseGroupLifecycleLeaseScript.Run(
  346. ctx,
  347. c.rdb,
  348. []string{schedulerGroupLifecycleLockKey(lease.GroupID)},
  349. lease.OwnerToken,
  350. ).Int64()
  351. if err != nil {
  352. return err
  353. }
  354. if result == 0 {
  355. return fmt.Errorf("%w: group=%d", service.ErrSchedulerGroupLifecycleLeaseLost, lease.GroupID)
  356. }
  357. if result != 1 {
  358. return fmt.Errorf("release scheduler group lifecycle lease returned %d", result)
  359. }
  360. return nil
  361. }
  362. func newSchedulerGroupLifecycleOwnerToken() (string, error) {
  363. raw := make([]byte, schedulerGroupLifecycleOwnerTokenBytes)
  364. if _, err := rand.Read(raw); err != nil {
  365. return "", fmt.Errorf("generate scheduler group lifecycle owner token: %w", err)
  366. }
  367. return hex.EncodeToString(raw), nil
  368. }
  369. func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) error {
  370. if !token.ValidFor(bucket) {
  371. return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
  372. }
  373. // 分配版本与激活指针是两个 fencing 边界;中间写入的数据只有通过第二次校验才能发布。
  374. version, err := c.allocateSnapshotVersion(ctx, bucket, token)
  375. if err != nil {
  376. return err
  377. }
  378. // 快照成员最终只依赖可编码账号的有序 ID;直接复用 ID 路径,避免为
  379. // 随后立即丢弃的完整 Account 再分配一份临时切片。
  380. if _, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts); err != nil {
  381. return err
  382. }
  383. return c.activateSnapshotVersion(ctx, bucket, token, version)
  384. }
  385. // SetSnapshotAndReturnAccountIDs 完整发布快照,并返回实际成功编码并写入的有序账号 ID。
  386. // 该可选能力只供同一重建批次复用,返回前仍会完成版本激活与 fencing 校验。
  387. func (c *schedulerCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) ([]int64, error) {
  388. if !token.ValidFor(bucket) {
  389. return nil, fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
  390. }
  391. // 分配版本与激活指针是两个 fencing 边界;中间写入的数据只有通过第二次校验才能发布。
  392. version, err := c.allocateSnapshotVersion(ctx, bucket, token)
  393. if err != nil {
  394. return nil, err
  395. }
  396. accountIDs, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts)
  397. if err != nil {
  398. return nil, err
  399. }
  400. if err := c.activateSnapshotVersion(ctx, bucket, token, version); err != nil {
  401. return nil, err
  402. }
  403. return accountIDs, nil
  404. }
  405. // SetSnapshotByAccountIDs 复用同批次首次完整写入后得到的账号成员。
  406. // 每个桶仍独立分配版本、写入有序集合并执行激活 fencing,只省略重复的账号 JSON 与全局键写入。
  407. func (c *schedulerCache) SetSnapshotByAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accountIDs []int64) error {
  408. if !token.ValidFor(bucket) {
  409. return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
  410. }
  411. version, err := c.allocateSnapshotVersion(ctx, bucket, token)
  412. if err != nil {
  413. return err
  414. }
  415. if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil {
  416. return err
  417. }
  418. return c.activateSnapshotVersion(ctx, bucket, token, version)
  419. }
  420. func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken) (string, error) {
  421. result, err := allocateSnapshotVersionScript.Run(ctx, c.rdb, []string{
  422. schedulerBucketKey(schedulerEpochPrefix, bucket),
  423. schedulerBucketKey(schedulerRetiredPrefix, bucket),
  424. schedulerBucketKey(schedulerVersionPrefix, bucket),
  425. }, token.Epoch).Int64()
  426. if err != nil {
  427. return "", err
  428. }
  429. if err := schedulerBucketWriteResultError(result, bucket); err != nil {
  430. return "", err
  431. }
  432. return strconv.FormatInt(result, 10), nil
  433. }
  434. func (c *schedulerCache) writeSnapshotVersionAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) ([]int64, error) {
  435. accountIDs, err := c.writeAccountIDs(ctx, accounts)
  436. if err != nil {
  437. return nil, err
  438. }
  439. if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil {
  440. return nil, err
  441. }
  442. return accountIDs, nil
  443. }
  444. func (c *schedulerCache) writeSnapshotAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accountIDs []int64) error {
  445. members := schedulerSnapshotMembers(accountIDs)
  446. return c.writeSnapshotMembers(ctx, bucket, version, members)
  447. }
  448. func schedulerSnapshotMembers(accountIDs []int64) []redis.Z {
  449. if len(accountIDs) == 0 {
  450. return nil
  451. }
  452. // 使用序号作为 score,保持数据库返回的排序语义;重复 ID 继续交由 Redis ZADD
  453. // 按最后一个 score 覆盖,与直接从账号切片构造成员时的行为一致。
  454. members := make([]redis.Z, 0, len(accountIDs))
  455. for idx, accountID := range accountIDs {
  456. members = append(members, redis.Z{
  457. Score: float64(idx),
  458. Member: strconv.FormatInt(accountID, 10),
  459. })
  460. }
  461. return members
  462. }
  463. func (c *schedulerCache) writeSnapshotMembers(ctx context.Context, bucket service.SchedulerBucket, version string, members []redis.Z) error {
  464. if len(members) == 0 {
  465. return nil
  466. }
  467. snapshotKey := schedulerSnapshotKey(bucket, version)
  468. pipe := c.rdb.Pipeline()
  469. for start := 0; start < len(members); start += c.writeChunkSize {
  470. end := start + c.writeChunkSize
  471. if end > len(members) {
  472. end = len(members)
  473. }
  474. pipe.ZAdd(ctx, snapshotKey, members[start:end]...)
  475. }
  476. _, err := pipe.Exec(ctx)
  477. return err
  478. }
  479. func (c *schedulerCache) activateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, version string) error {
  480. snapshotKey := schedulerSnapshotKey(bucket, version)
  481. // Phase 2: 原子 CAS 切换版本,同时再次校验退休状态与 writer epoch。
  482. // Lua 脚本保证:仅当新版本 >= 当前激活版本时才切换 active 指针,
  483. // 防止并发写入导致版本回滚。
  484. // 旧快照使用 EXPIRE 宽限期而非立即 DEL,避免 reader 竞态。
  485. activeKey := schedulerBucketKey(schedulerActivePrefix, bucket)
  486. readyKey := schedulerBucketKey(schedulerReadyPrefix, bucket)
  487. snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode)
  488. keys := []string{
  489. activeKey,
  490. readyKey,
  491. schedulerBucketSetKey,
  492. snapshotKey,
  493. schedulerBucketKey(schedulerEpochPrefix, bucket),
  494. schedulerBucketKey(schedulerRetiredPrefix, bucket),
  495. }
  496. args := []any{version, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds, token.Epoch}
  497. result, err := activateSnapshotScript.Run(ctx, c.rdb, keys, args...).Int64()
  498. if err != nil {
  499. return err
  500. }
  501. return schedulerBucketWriteResultError(result, bucket)
  502. }
  503. func schedulerBucketWriteResultError(result int64, bucket service.SchedulerBucket) error {
  504. switch result {
  505. case -1:
  506. return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketRetired, bucket.String())
  507. case -2:
  508. return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
  509. default:
  510. return nil
  511. }
  512. }
  513. func (c *schedulerCache) GetAccount(ctx context.Context, accountID int64) (*service.Account, error) {
  514. id := strconv.FormatInt(accountID, 10)
  515. values, err := c.rdb.MGet(ctx, schedulerAccountKey(id), schedulerLastUsedKey(id)).Result()
  516. if err != nil {
  517. return nil, err
  518. }
  519. if len(values) != 2 || values[0] == nil {
  520. return nil, nil
  521. }
  522. account, err := decodeCachedAccount(values[0])
  523. if err != nil {
  524. return nil, err
  525. }
  526. if err := applySchedulerLastUsed(account, values[1]); err != nil {
  527. return nil, err
  528. }
  529. return account, nil
  530. }
  531. func (c *schedulerCache) SetAccount(ctx context.Context, account *service.Account) error {
  532. if account == nil || account.ID <= 0 {
  533. return nil
  534. }
  535. accountIDs, err := c.writeAccountIDs(ctx, []service.Account{*account})
  536. if err != nil {
  537. return err
  538. }
  539. if len(accountIDs) == 0 {
  540. return c.DeleteAccount(ctx, account.ID)
  541. }
  542. return nil
  543. }
  544. func (c *schedulerCache) DeleteAccount(ctx context.Context, accountID int64) error {
  545. if accountID <= 0 {
  546. return nil
  547. }
  548. id := strconv.FormatInt(accountID, 10)
  549. return c.rdb.Del(ctx, schedulerAccountKey(id), schedulerAccountMetaKey(id), schedulerLastUsedKey(id)).Err()
  550. }
  551. func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
  552. if len(updates) == 0 {
  553. return nil
  554. }
  555. pipe := c.rdb.Pipeline()
  556. queued := 0
  557. keys := make([]string, 0, schedulerLastUsedUpdateChunkSize*2)
  558. args := make([]any, 0, schedulerLastUsedUpdateChunkSize)
  559. queueBatch := func() {
  560. if len(args) == 0 {
  561. return
  562. }
  563. updateSchedulerLastUsedScript.Eval(ctx, pipe, keys, args...)
  564. queued++
  565. keys = make([]string, 0, schedulerLastUsedUpdateChunkSize*2)
  566. args = make([]any, 0, schedulerLastUsedUpdateChunkSize)
  567. }
  568. for id, usedAt := range updates {
  569. if id <= 0 {
  570. continue
  571. }
  572. millis, err := schedulerLastUsedMillis(usedAt)
  573. if err != nil {
  574. slog.Warn("scheduler cache removes account with unencodable payload",
  575. "account_id", id,
  576. "error", err,
  577. )
  578. idText := strconv.FormatInt(id, 10)
  579. pipe.Del(ctx, schedulerAccountKey(idText), schedulerAccountMetaKey(idText), schedulerLastUsedKey(idText))
  580. queued++
  581. continue
  582. }
  583. idText := strconv.FormatInt(id, 10)
  584. keys = append(keys, schedulerAccountKey(idText), schedulerLastUsedKey(idText))
  585. args = append(args, millis)
  586. if len(args) >= schedulerLastUsedUpdateChunkSize {
  587. queueBatch()
  588. }
  589. }
  590. queueBatch()
  591. if queued == 0 {
  592. return nil
  593. }
  594. _, err := pipe.Exec(ctx)
  595. return err
  596. }
  597. func (c *schedulerCache) TryLockBucket(ctx context.Context, bucket service.SchedulerBucket, ttl time.Duration) (bool, error) {
  598. key := schedulerBucketKey(schedulerLockPrefix, bucket)
  599. return c.rdb.SetNX(ctx, key, time.Now().UnixNano(), ttl).Result()
  600. }
  601. func (c *schedulerCache) UnlockBucket(ctx context.Context, bucket service.SchedulerBucket) error {
  602. key := schedulerBucketKey(schedulerLockPrefix, bucket)
  603. return c.rdb.Del(ctx, key).Err()
  604. }
  605. func (c *schedulerCache) ListBuckets(ctx context.Context) ([]service.SchedulerBucket, error) {
  606. raw, err := c.rdb.SMembers(ctx, schedulerBucketSetKey).Result()
  607. if err != nil {
  608. return nil, err
  609. }
  610. out := make([]service.SchedulerBucket, 0, len(raw))
  611. for _, entry := range raw {
  612. bucket, ok := service.ParseSchedulerBucket(entry)
  613. if !ok {
  614. continue
  615. }
  616. out = append(out, bucket)
  617. }
  618. return out, nil
  619. }
  620. func (c *schedulerCache) GetOutboxWatermark(ctx context.Context) (int64, error) {
  621. val, err := c.rdb.Get(ctx, schedulerOutboxWatermarkKey).Result()
  622. if err == redis.Nil {
  623. return 0, nil
  624. }
  625. if err != nil {
  626. return 0, err
  627. }
  628. id, err := strconv.ParseInt(val, 10, 64)
  629. if err != nil {
  630. return 0, err
  631. }
  632. return id, nil
  633. }
  634. func (c *schedulerCache) SetOutboxWatermark(ctx context.Context, id int64) error {
  635. return c.rdb.Set(ctx, schedulerOutboxWatermarkKey, strconv.FormatInt(id, 10), 0).Err()
  636. }
  637. func schedulerBucketKey(prefix string, bucket service.SchedulerBucket) string {
  638. return fmt.Sprintf("%s%d:%s:%s", prefix, bucket.GroupID, bucket.Platform, bucket.Mode)
  639. }
  640. func schedulerGroupLifecycleLockKey(groupID int64) string {
  641. return schedulerGroupLifecycleLockPrefix + strconv.FormatInt(groupID, 10)
  642. }
  643. func schedulerSnapshotKey(bucket service.SchedulerBucket, version string) string {
  644. return fmt.Sprintf("%s%d:%s:%s:v%s", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode, version)
  645. }
  646. func schedulerAccountKey(id string) string {
  647. return schedulerAccountPrefix + id
  648. }
  649. func schedulerAccountMetaKey(id string) string {
  650. return schedulerAccountMetaPrefix + id
  651. }
  652. func schedulerLastUsedKey(id string) string {
  653. return schedulerAccountLastUsedPrefix + id
  654. }
  655. func ptrTime(t time.Time) *time.Time {
  656. return &t
  657. }
  658. func schedulerLastUsedMillis(value time.Time) (int64, error) {
  659. if _, err := value.MarshalJSON(); err != nil {
  660. return 0, err
  661. }
  662. return value.UTC().UnixMilli(), nil
  663. }
  664. func applySchedulerLastUsed(account *service.Account, value any) error {
  665. if account == nil || value == nil {
  666. return nil
  667. }
  668. var raw string
  669. switch typed := value.(type) {
  670. case string:
  671. raw = typed
  672. case []byte:
  673. raw = string(typed)
  674. default:
  675. return fmt.Errorf("unexpected last_used cache type: %T", value)
  676. }
  677. millis, err := strconv.ParseInt(raw, 10, 64)
  678. if err != nil {
  679. return fmt.Errorf("invalid last_used cache value %q: %w", raw, err)
  680. }
  681. lastUsedAt := time.UnixMilli(millis).UTC()
  682. if account.LastUsedAt == nil || lastUsedAt.After(*account.LastUsedAt) {
  683. account.LastUsedAt = ptrTime(lastUsedAt)
  684. }
  685. return nil
  686. }
  687. func decodeCachedAccount(val any) (*service.Account, error) {
  688. var payload []byte
  689. switch raw := val.(type) {
  690. case string:
  691. payload = []byte(raw)
  692. case []byte:
  693. payload = raw
  694. default:
  695. return nil, fmt.Errorf("unexpected account cache type: %T", val)
  696. }
  697. var account service.Account
  698. if err := json.Unmarshal(payload, &account); err != nil {
  699. return nil, err
  700. }
  701. return &account, nil
  702. }
  703. func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service.Account) ([]int64, error) {
  704. if len(accounts) == 0 {
  705. return nil, nil
  706. }
  707. pipe := c.rdb.Pipeline()
  708. accountIDs := make([]int64, 0, len(accounts))
  709. pending := 0
  710. flush := func() error {
  711. if pending == 0 {
  712. return nil
  713. }
  714. if _, err := pipe.Exec(ctx); err != nil {
  715. return err
  716. }
  717. pipe = c.rdb.Pipeline()
  718. pending = 0
  719. return nil
  720. }
  721. for _, account := range accounts {
  722. fullPayload, metaPayload, err := marshalSchedulerCacheAccount(account)
  723. if err != nil {
  724. slog.Warn("scheduler cache skips account with unencodable payload",
  725. "account_id", account.ID,
  726. "error", err,
  727. )
  728. continue
  729. }
  730. id := strconv.FormatInt(account.ID, 10)
  731. pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0)
  732. pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0)
  733. // Keep the hot LastUsedAt side key untouched: a lagging snapshot rebuild
  734. // must not overwrite a newer scheduler update.
  735. accountIDs = append(accountIDs, account.ID)
  736. pending++
  737. if pending >= c.writeChunkSize {
  738. if err := flush(); err != nil {
  739. return nil, err
  740. }
  741. }
  742. }
  743. if err := flush(); err != nil {
  744. return nil, err
  745. }
  746. return accountIDs, nil
  747. }
  748. func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) {
  749. fullPayload, err := json.Marshal(account)
  750. if err != nil {
  751. return nil, nil, fmt.Errorf("marshal account: %w", err)
  752. }
  753. metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account))
  754. if err != nil {
  755. return nil, nil, fmt.Errorf("marshal account metadata: %w", err)
  756. }
  757. return fullPayload, metaPayload, nil
  758. }
  759. func (c *schedulerCache) mgetChunked(ctx context.Context, keys []string) ([]any, error) {
  760. if len(keys) == 0 {
  761. return []any{}, nil
  762. }
  763. out := make([]any, 0, len(keys))
  764. chunkSize := c.mgetChunkSize
  765. if chunkSize <= 0 {
  766. chunkSize = defaultSchedulerSnapshotMGetChunkSize
  767. }
  768. for start := 0; start < len(keys); start += chunkSize {
  769. end := start + chunkSize
  770. if end > len(keys) {
  771. end = len(keys)
  772. }
  773. part, err := c.rdb.MGet(ctx, keys[start:end]...).Result()
  774. if err != nil {
  775. return nil, err
  776. }
  777. out = append(out, part...)
  778. }
  779. return out, nil
  780. }
  781. func buildSchedulerMetadataAccount(account service.Account) service.Account {
  782. return service.Account{
  783. ID: account.ID,
  784. Name: account.Name,
  785. Platform: account.Platform,
  786. Type: account.Type,
  787. Concurrency: account.Concurrency,
  788. LoadFactor: account.LoadFactor,
  789. Priority: account.Priority,
  790. RateMultiplier: account.RateMultiplier,
  791. Status: account.Status,
  792. LastUsedAt: account.LastUsedAt,
  793. ExpiresAt: account.ExpiresAt,
  794. AutoPauseOnExpired: account.AutoPauseOnExpired,
  795. Schedulable: account.Schedulable,
  796. RateLimitedAt: account.RateLimitedAt,
  797. RateLimitResetAt: account.RateLimitResetAt,
  798. OverloadUntil: account.OverloadUntil,
  799. TempUnschedulableUntil: account.TempUnschedulableUntil,
  800. TempUnschedulableReason: account.TempUnschedulableReason,
  801. SessionWindowStart: account.SessionWindowStart,
  802. SessionWindowEnd: account.SessionWindowEnd,
  803. SessionWindowStatus: account.SessionWindowStatus,
  804. ParentAccountID: account.ParentAccountID,
  805. QuotaDimension: account.QuotaDimension,
  806. AccountGroups: filterSchedulerAccountGroups(account.AccountGroups),
  807. GroupIDs: filterSchedulerGroupIDs(account.GroupIDs, account.AccountGroups),
  808. Credentials: filterSchedulerCredentials(account.Credentials),
  809. Extra: filterSchedulerExtra(account.Extra),
  810. }
  811. }
  812. func filterSchedulerAccountGroups(accountGroups []service.AccountGroup) []service.AccountGroup {
  813. if len(accountGroups) == 0 {
  814. return nil
  815. }
  816. filtered := make([]service.AccountGroup, 0, len(accountGroups))
  817. for _, ag := range accountGroups {
  818. if ag.GroupID <= 0 {
  819. continue
  820. }
  821. filtered = append(filtered, service.AccountGroup{
  822. AccountID: ag.AccountID,
  823. GroupID: ag.GroupID,
  824. Priority: ag.Priority,
  825. CreatedAt: ag.CreatedAt,
  826. })
  827. }
  828. if len(filtered) == 0 {
  829. return nil
  830. }
  831. return filtered
  832. }
  833. func filterSchedulerGroupIDs(groupIDs []int64, accountGroups []service.AccountGroup) []int64 {
  834. if len(groupIDs) == 0 && len(accountGroups) == 0 {
  835. return nil
  836. }
  837. seen := make(map[int64]struct{}, len(groupIDs)+len(accountGroups))
  838. filtered := make([]int64, 0, len(groupIDs)+len(accountGroups))
  839. for _, id := range groupIDs {
  840. if id <= 0 {
  841. continue
  842. }
  843. if _, ok := seen[id]; ok {
  844. continue
  845. }
  846. seen[id] = struct{}{}
  847. filtered = append(filtered, id)
  848. }
  849. for _, ag := range accountGroups {
  850. if ag.GroupID <= 0 {
  851. continue
  852. }
  853. if _, ok := seen[ag.GroupID]; ok {
  854. continue
  855. }
  856. seen[ag.GroupID] = struct{}{}
  857. filtered = append(filtered, ag.GroupID)
  858. }
  859. if len(filtered) == 0 {
  860. return nil
  861. }
  862. return filtered
  863. }
  864. func filterSchedulerCredentials(credentials map[string]any) map[string]any {
  865. if len(credentials) == 0 {
  866. return nil
  867. }
  868. keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type"}
  869. filtered := make(map[string]any)
  870. for _, key := range keys {
  871. if value, ok := credentials[key]; ok && value != nil {
  872. filtered[key] = value
  873. }
  874. }
  875. if len(filtered) == 0 {
  876. return nil
  877. }
  878. return filtered
  879. }
  880. func filterSchedulerExtra(extra map[string]any) map[string]any {
  881. if len(extra) == 0 {
  882. return nil
  883. }
  884. keys := []string{
  885. "quota_limit",
  886. "quota_used",
  887. "quota_daily_limit",
  888. "quota_daily_used",
  889. "quota_daily_start",
  890. "quota_daily_reset_mode",
  891. "quota_daily_reset_hour",
  892. "quota_weekly_limit",
  893. "quota_weekly_used",
  894. "quota_weekly_start",
  895. "quota_weekly_reset_mode",
  896. "quota_weekly_reset_day",
  897. "quota_weekly_reset_hour",
  898. "quota_reset_timezone",
  899. "mixed_scheduling",
  900. "window_cost_limit",
  901. "window_cost_sticky_reserve",
  902. "max_sessions",
  903. "session_idle_timeout_minutes",
  904. "openai_oauth_responses_websockets_v2_enabled",
  905. "openai_oauth_responses_websockets_v2_mode",
  906. "openai_apikey_responses_websockets_v2_enabled",
  907. "openai_apikey_responses_websockets_v2_mode",
  908. "responses_websockets_v2_enabled",
  909. "openai_ws_enabled",
  910. "openai_ws_force_http",
  911. "openai_responses_mode",
  912. "openai_responses_supported",
  913. // 透传开关必须进投影:候选过滤(ListSchedulableAccounts)读的是本投影,
  914. // 而 Account.IsModelSupported 靠 extra 上的这两个键短路 model_mapping 白名单。
  915. // 裁掉它们,透传账号在选号阶段会退回按(常为过期的)白名单判定并被误判为
  916. // model_not_supported —— 转发阶段却仍按透传工作,表现为"单独测账号能通、
  917. // 走网关报 no available accounts"。
  918. "openai_passthrough",
  919. "openai_oauth_passthrough",
  920. "codex_fingerprint_mode",
  921. "request_integrity_mode",
  922. "account_traffic_control",
  923. "codex_fingerprint_seed",
  924. "anti_degrade",
  925. "anti_degradation",
  926. "protection_scope",
  927. "proxy_mode",
  928. "enable_tls_fingerprint",
  929. "tls_fingerprint_builtin",
  930. "tls_fingerprint_profile_id",
  931. "codex_5h_used_percent",
  932. "codex_7d_used_percent",
  933. "codex_5h_reset_at",
  934. "codex_7d_reset_at",
  935. "codex_5h_reset_after_seconds",
  936. "codex_7d_reset_after_seconds",
  937. "codex_usage_updated_at",
  938. "auto_pause_5h_threshold",
  939. "auto_pause_7d_threshold",
  940. "auto_pause_5h_disabled",
  941. "auto_pause_7d_disabled",
  942. "model_rate_limits",
  943. service.UpstreamBillingProbeExtraKey,
  944. service.GrokMediaEligibleExtraKey,
  945. "grok_billing_snapshot",
  946. }
  947. filtered := make(map[string]any)
  948. for _, key := range keys {
  949. if value, ok := extra[key]; ok && value != nil {
  950. if key == service.UpstreamBillingProbeExtraKey {
  951. filteredProbe := filterSchedulerUpstreamBillingProbe(value)
  952. if filteredProbe == nil {
  953. continue
  954. }
  955. value = filteredProbe
  956. }
  957. filtered[key] = value
  958. }
  959. }
  960. if len(filtered) == 0 {
  961. return nil
  962. }
  963. return filtered
  964. }
  965. func filterSchedulerUpstreamBillingProbe(value any) map[string]any {
  966. source, ok := value.(map[string]any)
  967. if !ok {
  968. return nil
  969. }
  970. status, ok := source["status"].(string)
  971. if !ok || status == "" {
  972. return nil
  973. }
  974. filtered := map[string]any{"status": status}
  975. for _, key := range []string{"received_at", "fresh_until", "next_probe_at"} {
  976. if field, exists := source[key]; exists && field != nil {
  977. filtered[key] = field
  978. }
  979. }
  980. data, ok := source["data"].(map[string]any)
  981. if !ok {
  982. return filtered
  983. }
  984. filteredData := make(map[string]any)
  985. for _, key := range []string{
  986. "billing_scope",
  987. "resolved_rate_multiplier",
  988. "peak_rate_enabled",
  989. "peak_start",
  990. "peak_end",
  991. "peak_rate_multiplier",
  992. "timezone",
  993. } {
  994. if field, exists := data[key]; exists && field != nil {
  995. filteredData[key] = field
  996. }
  997. }
  998. if len(filteredData) > 0 {
  999. filtered["data"] = filteredData
  1000. }
  1001. return filtered
  1002. }