concurrency_service.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771
  1. package service
  2. import (
  3. "context"
  4. "crypto/rand"
  5. "crypto/sha256"
  6. "encoding/binary"
  7. "encoding/hex"
  8. "errors"
  9. "os"
  10. "strconv"
  11. "sync"
  12. "sync/atomic"
  13. "time"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  15. "go.uber.org/zap"
  16. "golang.org/x/sync/singleflight"
  17. )
  18. // ConcurrencyCache 定义并发控制的缓存接口
  19. // 使用有序集合存储槽位,按时间戳清理过期条目
  20. type ConcurrencyCache interface {
  21. // 账号槽位管理
  22. // 键格式: concurrency:account:{accountID}(有序集合,成员为 requestID)
  23. AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
  24. ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error
  25. GetAccountConcurrency(ctx context.Context, accountID int64) (int, error)
  26. GetAccountConcurrencyBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error)
  27. // 账号等待队列(账号级)
  28. IncrementAccountWaitCount(ctx context.Context, accountID int64, maxWait int) (bool, error)
  29. DecrementAccountWaitCount(ctx context.Context, accountID int64) error
  30. GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error)
  31. // 用户槽位管理
  32. // 键格式: concurrency:user:{userID}(有序集合,成员为 requestID)
  33. AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
  34. ReleaseUserSlot(ctx context.Context, userID int64, requestID string) error
  35. GetUserConcurrency(ctx context.Context, userID int64) (int, error)
  36. // 等待队列计数(每次入队都会刷新 TTL,避免长时间排队时计数提前过期)
  37. IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error)
  38. DecrementWaitCount(ctx context.Context, userID int64) error
  39. // 批量负载查询(只读)
  40. GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error)
  41. GetUsersLoadBatch(ctx context.Context, users []UserWithConcurrency) (map[int64]*UserLoadInfo, error)
  42. // 清理过期槽位(后台任务)
  43. CleanupExpiredAccountSlots(ctx context.Context, accountID int64) error
  44. CleanupExpiredAccountSlotKeys(ctx context.Context) error
  45. // 启动时清理旧进程遗留槽位与等待计数
  46. CleanupStaleProcessSlots(ctx context.Context, activeRequestPrefix string) error
  47. }
  48. type APIKeyConcurrencyCache interface {
  49. TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error
  50. ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error
  51. GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error)
  52. }
  53. // OpenAIWSIngressLeaseCache owns the short-lived distributed lease used to
  54. // bound live client WebSocket sessions. It is deliberately independent of the
  55. // request-slot namespace: idle ingress connections do not occupy turn slots.
  56. type OpenAIWSIngressLeaseCache interface {
  57. AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error)
  58. RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error)
  59. ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error
  60. }
  61. const (
  62. openAIWSIngressLeaseTTL = 60 * time.Second
  63. openAIWSIngressLeaseRefreshInterval = 20 * time.Second
  64. openAIWSIngressLeaseOperationTO = 2 * time.Second
  65. )
  66. var ErrOpenAIWSIngressLeaseLost = errors.New("openai websocket ingress lease lost")
  67. // OpenAIWSIngressLease keeps a Redis-backed ingress lease alive and cancels
  68. // its context if Redis cannot confirm ownership for a full lease lifetime.
  69. // Call Release on every handler exit to reclaim capacity immediately.
  70. type OpenAIWSIngressLease struct {
  71. ctx context.Context
  72. cancel context.CancelCauseFunc
  73. cache OpenAIWSIngressLeaseCache
  74. apiKeyID int64
  75. leaseID string
  76. stopOnce sync.Once
  77. stopCh chan struct{}
  78. refreshDone chan struct{}
  79. }
  80. func (l *OpenAIWSIngressLease) Context() context.Context {
  81. if l == nil || l.ctx == nil {
  82. return context.Background()
  83. }
  84. return l.ctx
  85. }
  86. func (l *OpenAIWSIngressLease) Release() {
  87. if l == nil {
  88. return
  89. }
  90. l.stopOnce.Do(func() {
  91. if l.stopCh != nil {
  92. close(l.stopCh)
  93. }
  94. if l.cancel != nil {
  95. l.cancel(nil)
  96. }
  97. if l.refreshDone != nil {
  98. <-l.refreshDone
  99. }
  100. if l.cache == nil || l.apiKeyID <= 0 || l.leaseID == "" {
  101. return
  102. }
  103. releaseCtx, releaseCancel := context.WithTimeout(context.Background(), openAIWSIngressLeaseOperationTO)
  104. defer releaseCancel()
  105. if err := l.cache.ReleaseOpenAIWSIngressLease(releaseCtx, l.apiKeyID, l.leaseID); err != nil {
  106. logger.L().Warn("openai_ws_ingress_lease_release_failed",
  107. zap.Int64("api_key_id", l.apiKeyID),
  108. zap.Error(err),
  109. )
  110. }
  111. })
  112. }
  113. func (l *OpenAIWSIngressLease) refreshLoop() {
  114. defer func() {
  115. if l != nil && l.refreshDone != nil {
  116. close(l.refreshDone)
  117. }
  118. }()
  119. if l == nil || l.cache == nil {
  120. return
  121. }
  122. ticker := time.NewTicker(openAIWSIngressLeaseRefreshInterval)
  123. defer ticker.Stop()
  124. lastConfirmedAt := time.Now()
  125. for {
  126. select {
  127. case <-l.ctx.Done():
  128. return
  129. case <-l.stopCh:
  130. return
  131. case <-ticker.C:
  132. var lost bool
  133. lastConfirmedAt, lost = l.refresh(lastConfirmedAt)
  134. if lost {
  135. l.cancel(ErrOpenAIWSIngressLeaseLost)
  136. return
  137. }
  138. }
  139. }
  140. }
  141. // refresh confirms the lease is still owned. A missing member is an immediate
  142. // lease loss; transient Redis errors are tolerated only for one full lease TTL.
  143. func (l *OpenAIWSIngressLease) refresh(lastConfirmedAt time.Time) (time.Time, bool) {
  144. refreshCtx, refreshCancel := context.WithTimeout(context.Background(), openAIWSIngressLeaseOperationTO)
  145. owned, err := l.cache.RefreshOpenAIWSIngressLease(refreshCtx, l.apiKeyID, l.leaseID)
  146. refreshCancel()
  147. if err == nil && owned {
  148. return time.Now(), false
  149. }
  150. if err == nil {
  151. err = ErrOpenAIWSIngressLeaseLost
  152. }
  153. elapsed := time.Since(lastConfirmedAt)
  154. logger.L().Warn("openai_ws_ingress_lease_refresh_failed",
  155. zap.Int64("api_key_id", l.apiKeyID),
  156. zap.Duration("unconfirmed_for", elapsed),
  157. zap.Error(err),
  158. )
  159. if errors.Is(err, ErrOpenAIWSIngressLeaseLost) || elapsed >= openAIWSIngressLeaseTTL {
  160. logger.L().Error("openai_ws_ingress_lease_lost",
  161. zap.Int64("api_key_id", l.apiKeyID),
  162. zap.Duration("unconfirmed_for", elapsed),
  163. zap.Error(err),
  164. )
  165. return lastConfirmedAt, true
  166. }
  167. return lastConfirmedAt, false
  168. }
  169. var (
  170. requestIDPrefix = initRequestIDPrefix()
  171. requestIDCounter atomic.Uint64
  172. )
  173. func initRequestIDPrefix() string {
  174. b := make([]byte, 8)
  175. if _, err := rand.Read(b); err == nil {
  176. return "r" + strconv.FormatUint(binary.BigEndian.Uint64(b), 36)
  177. }
  178. fallback := uint64(time.Now().UnixNano()) ^ (uint64(os.Getpid()) << 16)
  179. return "r" + strconv.FormatUint(fallback, 36)
  180. }
  181. func RequestIDPrefix() string {
  182. return requestIDPrefix
  183. }
  184. func generateRequestID() string {
  185. seq := requestIDCounter.Add(1)
  186. return requestIDPrefix + "-" + strconv.FormatUint(seq, 36)
  187. }
  188. func (s *ConcurrencyService) CleanupStaleProcessSlots(ctx context.Context) error {
  189. if s == nil || s.cache == nil {
  190. return nil
  191. }
  192. return s.cache.CleanupStaleProcessSlots(ctx, RequestIDPrefix())
  193. }
  194. const (
  195. // 默认等待队列额外槽位
  196. defaultExtraWaitSlots = 20
  197. defaultAccountLoadBatchCacheTTL = 200 * time.Millisecond
  198. accountLoadBatchFetchTimeout = 3 * time.Second
  199. maxAccountLoadBatchCacheEntries = 256
  200. apiKeyConcurrencyFetchTimeout = 3 * time.Second
  201. apiKeySlotTrackTimeout = 2 * time.Second
  202. )
  203. // ConcurrencyService 管理账号和用户的并发限制。
  204. type ConcurrencyService struct {
  205. cache ConcurrencyCache
  206. accountLoadCacheTTL atomic.Int64
  207. accountLoadCacheMu sync.RWMutex
  208. accountLoadCache map[string]cachedAccountLoadBatch
  209. accountLoadGroup singleflight.Group
  210. }
  211. type cachedAccountLoadBatch struct {
  212. loadMap map[int64]*AccountLoadInfo
  213. expiresAt time.Time
  214. }
  215. // NewConcurrencyService 创建并发控制服务。
  216. func NewConcurrencyService(cache ConcurrencyCache) *ConcurrencyService {
  217. svc := &ConcurrencyService{
  218. cache: cache,
  219. accountLoadCache: make(map[string]cachedAccountLoadBatch),
  220. }
  221. svc.SetAccountLoadBatchCacheTTL(defaultAccountLoadBatchCacheTTL)
  222. return svc
  223. }
  224. // AcquireOpenAIWSIngressLease atomically reserves one live ingress connection
  225. // for an API key. A non-positive limit explicitly disables this protection.
  226. func (s *ConcurrencyService) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int) (*OpenAIWSIngressLease, bool, error) {
  227. if maxConnections <= 0 {
  228. return nil, true, nil
  229. }
  230. if s == nil || s.cache == nil || apiKeyID <= 0 {
  231. return nil, false, errors.New("openai websocket ingress lease cache is unavailable")
  232. }
  233. cache, ok := s.cache.(OpenAIWSIngressLeaseCache)
  234. if !ok {
  235. return nil, false, errors.New("openai websocket ingress lease cache is unsupported")
  236. }
  237. leaseID := generateRequestID()
  238. baseCtx := context.Background()
  239. if ctx != nil {
  240. baseCtx = context.WithoutCancel(ctx)
  241. }
  242. acquireCtx, acquireCancel := context.WithTimeout(baseCtx, openAIWSIngressLeaseOperationTO)
  243. acquired, err := cache.AcquireOpenAIWSIngressLease(acquireCtx, apiKeyID, maxConnections, leaseID)
  244. acquireCancel()
  245. if err != nil || !acquired {
  246. return nil, acquired, err
  247. }
  248. if ctx == nil {
  249. ctx = context.Background()
  250. }
  251. leaseCtx, leaseCancel := context.WithCancelCause(ctx)
  252. lease := &OpenAIWSIngressLease{
  253. ctx: leaseCtx,
  254. cancel: leaseCancel,
  255. cache: cache,
  256. apiKeyID: apiKeyID,
  257. leaseID: leaseID,
  258. stopCh: make(chan struct{}),
  259. refreshDone: make(chan struct{}),
  260. }
  261. go lease.refreshLoop()
  262. return lease, true, nil
  263. }
  264. // SetAccountLoadBatchCacheTTL 设置账号负载批量读取的极短 TTL 缓存;非正数表示禁用缓存。
  265. func (s *ConcurrencyService) SetAccountLoadBatchCacheTTL(ttl time.Duration) {
  266. if s == nil {
  267. return
  268. }
  269. s.accountLoadCacheTTL.Store(int64(ttl))
  270. if ttl <= 0 {
  271. s.accountLoadCacheMu.Lock()
  272. s.accountLoadCache = make(map[string]cachedAccountLoadBatch)
  273. s.accountLoadCacheMu.Unlock()
  274. }
  275. }
  276. // AcquireResult represents the result of acquiring a concurrency slot
  277. type AcquireResult struct {
  278. Acquired bool
  279. ReleaseFunc func() // Must be called when done (typically via defer)
  280. }
  281. type AccountWithConcurrency struct {
  282. ID int64
  283. MaxConcurrency int
  284. }
  285. type UserWithConcurrency struct {
  286. ID int64
  287. MaxConcurrency int
  288. }
  289. type AccountLoadInfo struct {
  290. AccountID int64
  291. CurrentConcurrency int
  292. WaitingCount int
  293. LoadRate int // 0-100+ (percent)
  294. }
  295. type UserLoadInfo struct {
  296. UserID int64
  297. CurrentConcurrency int
  298. WaitingCount int
  299. LoadRate int // 0-100+ (percent)
  300. }
  301. // AcquireAccountSlot attempts to acquire a concurrency slot for an account.
  302. // If the account is at max concurrency, it waits until a slot is available or timeout.
  303. // Returns a release function that MUST be called when the request completes.
  304. func (s *ConcurrencyService) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (*AcquireResult, error) {
  305. // If maxConcurrency is 0 or negative, no limit
  306. if maxConcurrency <= 0 {
  307. return &AcquireResult{
  308. Acquired: true,
  309. ReleaseFunc: func() {}, // no-op
  310. }, nil
  311. }
  312. // Generate unique request ID for this slot
  313. requestID := generateRequestID()
  314. acquired, err := s.cache.AcquireAccountSlot(ctx, accountID, maxConcurrency, requestID)
  315. if err != nil {
  316. return nil, err
  317. }
  318. if acquired {
  319. return &AcquireResult{
  320. Acquired: true,
  321. ReleaseFunc: func() {
  322. bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  323. defer cancel()
  324. if err := s.cache.ReleaseAccountSlot(bgCtx, accountID, requestID); err != nil {
  325. logger.LegacyPrintf("service.concurrency", "Warning: failed to release account slot for %d (req=%s): %v", accountID, requestID, err)
  326. }
  327. },
  328. }, nil
  329. }
  330. return &AcquireResult{
  331. Acquired: false,
  332. ReleaseFunc: nil,
  333. }, nil
  334. }
  335. // AcquireUserSlot attempts to acquire a concurrency slot for a user.
  336. // If the user is at max concurrency, it waits until a slot is available or timeout.
  337. // Returns a release function that MUST be called when the request completes.
  338. func (s *ConcurrencyService) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int) (*AcquireResult, error) {
  339. // If maxConcurrency is 0 or negative, no limit
  340. if maxConcurrency <= 0 {
  341. return &AcquireResult{
  342. Acquired: true,
  343. ReleaseFunc: func() {}, // no-op
  344. }, nil
  345. }
  346. // Generate unique request ID for this slot
  347. requestID := generateRequestID()
  348. acquired, err := s.cache.AcquireUserSlot(ctx, userID, maxConcurrency, requestID)
  349. if err != nil {
  350. return nil, err
  351. }
  352. if acquired {
  353. return &AcquireResult{
  354. Acquired: true,
  355. ReleaseFunc: func() {
  356. bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  357. defer cancel()
  358. if err := s.cache.ReleaseUserSlot(bgCtx, userID, requestID); err != nil {
  359. logger.LegacyPrintf("service.concurrency", "Warning: failed to release user slot for %d (req=%s): %v", userID, requestID, err)
  360. }
  361. },
  362. }, nil
  363. }
  364. return &AcquireResult{
  365. Acquired: false,
  366. ReleaseFunc: nil,
  367. }, nil
  368. }
  369. // TrackAPIKeySlot records one active request slot for an API key without
  370. // applying key-level concurrency limits. It is fail-open: Redis errors are
  371. // logged and return a no-op release function.
  372. func (s *ConcurrencyService) TrackAPIKeySlot(ctx context.Context, apiKeyID int64) func() {
  373. if s == nil || s.cache == nil || apiKeyID <= 0 {
  374. return func() {}
  375. }
  376. cache, ok := s.cache.(APIKeyConcurrencyCache)
  377. if !ok {
  378. return func() {}
  379. }
  380. requestID := generateRequestID()
  381. baseCtx := context.Background()
  382. if ctx != nil {
  383. baseCtx = context.WithoutCancel(ctx)
  384. }
  385. trackCtx, cancel := context.WithTimeout(baseCtx, apiKeySlotTrackTimeout)
  386. err := cache.TrackAPIKeySlot(trackCtx, apiKeyID, requestID)
  387. cancel()
  388. if err != nil {
  389. logger.LegacyPrintf("service.concurrency", "Warning: failed to track api key slot for %d (req=%s): %v", apiKeyID, requestID, err)
  390. return func() {}
  391. }
  392. return func() {
  393. bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  394. defer cancel()
  395. if err := cache.ReleaseAPIKeySlot(bgCtx, apiKeyID, requestID); err != nil {
  396. logger.LegacyPrintf("service.concurrency", "Warning: failed to release api key slot for %d (req=%s): %v", apiKeyID, requestID, err)
  397. }
  398. }
  399. }
  400. // GetAPIKeyConcurrencyBatch gets real-time active request counts for API keys.
  401. // Stats are best-effort: missing Redis support or Redis errors return zeroes.
  402. func (s *ConcurrencyService) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) {
  403. result := zeroAPIKeyConcurrencyMap(apiKeyIDs)
  404. if len(apiKeyIDs) == 0 {
  405. return result, nil
  406. }
  407. if s == nil || s.cache == nil {
  408. return result, nil
  409. }
  410. cache, ok := s.cache.(APIKeyConcurrencyCache)
  411. if !ok {
  412. return result, nil
  413. }
  414. redisCtx, cancel := context.WithTimeout(context.Background(), apiKeyConcurrencyFetchTimeout)
  415. defer cancel()
  416. counts, err := cache.GetAPIKeyConcurrencyBatch(redisCtx, apiKeyIDs)
  417. if err != nil {
  418. logger.LegacyPrintf("service.concurrency", "Warning: get api key concurrency batch failed: %v", err)
  419. return result, nil
  420. }
  421. for _, apiKeyID := range apiKeyIDs {
  422. result[apiKeyID] = counts[apiKeyID]
  423. }
  424. return result, nil
  425. }
  426. func zeroAPIKeyConcurrencyMap(apiKeyIDs []int64) map[int64]int {
  427. result := make(map[int64]int, len(apiKeyIDs))
  428. for _, apiKeyID := range apiKeyIDs {
  429. result[apiKeyID] = 0
  430. }
  431. return result
  432. }
  433. // ============================================
  434. // Wait Queue Count Methods
  435. // ============================================
  436. // IncrementWaitCount attempts to increment the wait queue counter for a user.
  437. // Returns true if successful, false if the wait queue is full.
  438. // maxWait should be user.Concurrency + defaultExtraWaitSlots
  439. func (s *ConcurrencyService) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) {
  440. if s.cache == nil {
  441. // Redis not available, allow request
  442. return true, nil
  443. }
  444. result, err := s.cache.IncrementWaitCount(ctx, userID, maxWait)
  445. if err != nil {
  446. // On error, allow the request to proceed (fail open)
  447. logger.LegacyPrintf("service.concurrency", "Warning: increment wait count failed for user %d: %v", userID, err)
  448. return true, nil
  449. }
  450. return result, nil
  451. }
  452. // DecrementWaitCount decrements the wait queue counter for a user.
  453. // Should be called when a request completes or exits the wait queue.
  454. func (s *ConcurrencyService) DecrementWaitCount(ctx context.Context, userID int64) {
  455. if s.cache == nil {
  456. return
  457. }
  458. // Use background context to ensure decrement even if original context is cancelled
  459. bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  460. defer cancel()
  461. if err := s.cache.DecrementWaitCount(bgCtx, userID); err != nil {
  462. logger.LegacyPrintf("service.concurrency", "Warning: decrement wait count failed for user %d: %v", userID, err)
  463. }
  464. }
  465. // IncrementAccountWaitCount increments the wait queue counter for an account.
  466. func (s *ConcurrencyService) IncrementAccountWaitCount(ctx context.Context, accountID int64, maxWait int) (bool, error) {
  467. if s.cache == nil {
  468. return true, nil
  469. }
  470. result, err := s.cache.IncrementAccountWaitCount(ctx, accountID, maxWait)
  471. if err != nil {
  472. logger.LegacyPrintf("service.concurrency", "Warning: increment wait count failed for account %d: %v", accountID, err)
  473. return true, nil
  474. }
  475. return result, nil
  476. }
  477. // DecrementAccountWaitCount decrements the wait queue counter for an account.
  478. func (s *ConcurrencyService) DecrementAccountWaitCount(ctx context.Context, accountID int64) {
  479. if s.cache == nil {
  480. return
  481. }
  482. bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  483. defer cancel()
  484. if err := s.cache.DecrementAccountWaitCount(bgCtx, accountID); err != nil {
  485. logger.LegacyPrintf("service.concurrency", "Warning: decrement wait count failed for account %d: %v", accountID, err)
  486. }
  487. }
  488. // GetAccountWaitingCount gets current wait queue count for an account.
  489. func (s *ConcurrencyService) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
  490. if s.cache == nil {
  491. return 0, nil
  492. }
  493. return s.cache.GetAccountWaitingCount(ctx, accountID)
  494. }
  495. // CalculateMaxWait calculates the maximum wait queue size for a user
  496. // maxWait = userConcurrency + defaultExtraWaitSlots
  497. func CalculateMaxWait(userConcurrency int) int {
  498. if userConcurrency <= 0 {
  499. userConcurrency = 1
  500. }
  501. return userConcurrency + defaultExtraWaitSlots
  502. }
  503. // GetAccountsLoadBatch 批量获取账号负载信息。
  504. func (s *ConcurrencyService) GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
  505. return s.getAccountsLoadBatch(ctx, accounts, true)
  506. }
  507. // GetAccountsLoadBatchFresh 绕过极短 TTL 缓存,用于抢槽失败后的实时刷新兜底。
  508. func (s *ConcurrencyService) GetAccountsLoadBatchFresh(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
  509. return s.getAccountsLoadBatch(ctx, accounts, false)
  510. }
  511. func (s *ConcurrencyService) getAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency, allowCache bool) (map[int64]*AccountLoadInfo, error) {
  512. if len(accounts) == 0 {
  513. return map[int64]*AccountLoadInfo{}, nil
  514. }
  515. if s.cache == nil {
  516. return map[int64]*AccountLoadInfo{}, nil
  517. }
  518. ttl := time.Duration(s.accountLoadCacheTTL.Load())
  519. if !allowCache || ttl <= 0 {
  520. return s.fetchAccountsLoadBatch(ctx, accounts)
  521. }
  522. key := accountLoadBatchCacheKey(accounts)
  523. if cached, ok := s.getCachedAccountLoadBatch(key, time.Now()); ok {
  524. return cached, nil
  525. }
  526. value, err, _ := s.accountLoadGroup.Do(key, func() (any, error) {
  527. now := time.Now()
  528. if cached, ok := s.getCachedAccountLoadBatch(key, now); ok {
  529. return cached, nil
  530. }
  531. loadMap, fetchErr := s.fetchAccountsLoadBatch(ctx, accounts)
  532. if fetchErr != nil {
  533. return nil, fetchErr
  534. }
  535. cached := cloneAccountLoadMap(loadMap)
  536. s.storeCachedAccountLoadBatch(key, cached, now.Add(ttl))
  537. return cached, nil
  538. })
  539. if err != nil {
  540. return nil, err
  541. }
  542. loadMap, _ := value.(map[int64]*AccountLoadInfo)
  543. if loadMap == nil {
  544. return map[int64]*AccountLoadInfo{}, nil
  545. }
  546. return loadMap, nil
  547. }
  548. func (s *ConcurrencyService) fetchAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
  549. if s.cache == nil {
  550. return map[int64]*AccountLoadInfo{}, nil
  551. }
  552. baseCtx := context.Background()
  553. if ctx != nil {
  554. baseCtx = context.WithoutCancel(ctx)
  555. }
  556. redisCtx, cancel := context.WithTimeout(baseCtx, accountLoadBatchFetchTimeout)
  557. defer cancel()
  558. return s.cache.GetAccountsLoadBatch(redisCtx, accounts)
  559. }
  560. func (s *ConcurrencyService) getCachedAccountLoadBatch(key string, now time.Time) (map[int64]*AccountLoadInfo, bool) {
  561. s.accountLoadCacheMu.RLock()
  562. cached, ok := s.accountLoadCache[key]
  563. s.accountLoadCacheMu.RUnlock()
  564. if !ok {
  565. return nil, false
  566. }
  567. if !now.Before(cached.expiresAt) {
  568. s.accountLoadCacheMu.Lock()
  569. if current, exists := s.accountLoadCache[key]; exists && !now.Before(current.expiresAt) {
  570. delete(s.accountLoadCache, key)
  571. }
  572. s.accountLoadCacheMu.Unlock()
  573. return nil, false
  574. }
  575. return cached.loadMap, true
  576. }
  577. func (s *ConcurrencyService) storeCachedAccountLoadBatch(key string, loadMap map[int64]*AccountLoadInfo, expiresAt time.Time) {
  578. s.accountLoadCacheMu.Lock()
  579. if s.accountLoadCache == nil {
  580. s.accountLoadCache = make(map[string]cachedAccountLoadBatch)
  581. }
  582. if len(s.accountLoadCache) >= maxAccountLoadBatchCacheEntries {
  583. now := time.Now()
  584. for cacheKey, cached := range s.accountLoadCache {
  585. if !now.Before(cached.expiresAt) {
  586. delete(s.accountLoadCache, cacheKey)
  587. }
  588. }
  589. for len(s.accountLoadCache) >= maxAccountLoadBatchCacheEntries {
  590. for cacheKey := range s.accountLoadCache {
  591. delete(s.accountLoadCache, cacheKey)
  592. break
  593. }
  594. }
  595. }
  596. s.accountLoadCache[key] = cachedAccountLoadBatch{
  597. loadMap: loadMap,
  598. expiresAt: expiresAt,
  599. }
  600. s.accountLoadCacheMu.Unlock()
  601. }
  602. func accountLoadBatchCacheKey(accounts []AccountWithConcurrency) string {
  603. hash := sha256.New()
  604. var buf [16]byte
  605. for _, account := range accounts {
  606. binary.LittleEndian.PutUint64(buf[:8], uint64(account.ID))
  607. binary.LittleEndian.PutUint64(buf[8:], uint64(int64(account.MaxConcurrency)))
  608. _, _ = hash.Write(buf[:])
  609. }
  610. sum := hash.Sum(nil)
  611. return strconv.Itoa(len(accounts)) + ":" + hex.EncodeToString(sum)
  612. }
  613. func cloneAccountLoadMap(loadMap map[int64]*AccountLoadInfo) map[int64]*AccountLoadInfo {
  614. if len(loadMap) == 0 {
  615. return map[int64]*AccountLoadInfo{}
  616. }
  617. clone := make(map[int64]*AccountLoadInfo, len(loadMap))
  618. for accountID, loadInfo := range loadMap {
  619. if loadInfo == nil {
  620. clone[accountID] = nil
  621. continue
  622. }
  623. copied := *loadInfo
  624. clone[accountID] = &copied
  625. }
  626. return clone
  627. }
  628. // GetUsersLoadBatch returns load info for multiple users.
  629. func (s *ConcurrencyService) GetUsersLoadBatch(ctx context.Context, users []UserWithConcurrency) (map[int64]*UserLoadInfo, error) {
  630. if s.cache == nil {
  631. return map[int64]*UserLoadInfo{}, nil
  632. }
  633. return s.cache.GetUsersLoadBatch(ctx, users)
  634. }
  635. // CleanupExpiredAccountSlots removes expired slots for one account (background task).
  636. func (s *ConcurrencyService) CleanupExpiredAccountSlots(ctx context.Context, accountID int64) error {
  637. if s.cache == nil {
  638. return nil
  639. }
  640. return s.cache.CleanupExpiredAccountSlots(ctx, accountID)
  641. }
  642. // StartSlotCleanupWorker starts a background cleanup worker for expired account slots.
  643. func (s *ConcurrencyService) StartSlotCleanupWorker(_ AccountRepository, interval time.Duration) {
  644. if s == nil || s.cache == nil || interval <= 0 {
  645. return
  646. }
  647. runCleanup := func() {
  648. cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  649. err := s.cache.CleanupExpiredAccountSlotKeys(cleanupCtx)
  650. cancel()
  651. if err != nil {
  652. logger.LegacyPrintf("service.concurrency", "Warning: cleanup expired account slots failed: %v", err)
  653. return
  654. }
  655. }
  656. go func() {
  657. ticker := time.NewTicker(interval)
  658. defer ticker.Stop()
  659. runCleanup()
  660. for range ticker.C {
  661. runCleanup()
  662. }
  663. }()
  664. }
  665. // GetAccountConcurrencyBatch gets current concurrency counts for multiple accounts.
  666. // Uses a detached context with timeout to prevent HTTP request cancellation from
  667. // causing the entire batch to fail (which would show all concurrency as 0).
  668. func (s *ConcurrencyService) GetAccountConcurrencyBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error) {
  669. if len(accountIDs) == 0 {
  670. return map[int64]int{}, nil
  671. }
  672. if s.cache == nil {
  673. result := make(map[int64]int, len(accountIDs))
  674. for _, accountID := range accountIDs {
  675. result[accountID] = 0
  676. }
  677. return result, nil
  678. }
  679. // Use a detached context so that a cancelled HTTP request doesn't cause
  680. // the Redis pipeline to fail and return all-zero concurrency counts.
  681. redisCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  682. defer cancel()
  683. return s.cache.GetAccountConcurrencyBatch(redisCtx, accountIDs)
  684. }