gateway_hotpath_optimization_test.go 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "sync"
  6. "sync/atomic"
  7. "testing"
  8. "time"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
  12. gocache "github.com/patrickmn/go-cache"
  13. "github.com/stretchr/testify/require"
  14. )
  15. type userGroupRateRepoHotpathStub struct {
  16. UserGroupRateRepository
  17. rate *float64
  18. err error
  19. wait <-chan struct{}
  20. calls atomic.Int64
  21. }
  22. func (s *userGroupRateRepoHotpathStub) GetByUserAndGroup(ctx context.Context, userID, groupID int64) (*float64, error) {
  23. s.calls.Add(1)
  24. if s.wait != nil {
  25. <-s.wait
  26. }
  27. if s.err != nil {
  28. return nil, s.err
  29. }
  30. return s.rate, nil
  31. }
  32. type usageLogWindowBatchRepoStub struct {
  33. UsageLogRepository
  34. batchResult map[int64]*usagestats.AccountStats
  35. batchErr error
  36. batchCalls atomic.Int64
  37. singleResult map[int64]*usagestats.AccountStats
  38. singleErr error
  39. singleCalls atomic.Int64
  40. }
  41. func (s *usageLogWindowBatchRepoStub) GetAccountWindowStatsBatch(ctx context.Context, accountIDs []int64, startTime time.Time) (map[int64]*usagestats.AccountStats, error) {
  42. s.batchCalls.Add(1)
  43. if s.batchErr != nil {
  44. return nil, s.batchErr
  45. }
  46. out := make(map[int64]*usagestats.AccountStats, len(accountIDs))
  47. for _, id := range accountIDs {
  48. if stats, ok := s.batchResult[id]; ok {
  49. out[id] = stats
  50. }
  51. }
  52. return out, nil
  53. }
  54. func (s *usageLogWindowBatchRepoStub) GetAccountWindowStats(ctx context.Context, accountID int64, startTime time.Time) (*usagestats.AccountStats, error) {
  55. s.singleCalls.Add(1)
  56. if s.singleErr != nil {
  57. return nil, s.singleErr
  58. }
  59. if stats, ok := s.singleResult[accountID]; ok {
  60. return stats, nil
  61. }
  62. return &usagestats.AccountStats{}, nil
  63. }
  64. type sessionLimitCacheHotpathStub struct {
  65. SessionLimitCache
  66. batchData map[int64]float64
  67. batchErr error
  68. setData map[int64]float64
  69. setErr error
  70. }
  71. func (s *sessionLimitCacheHotpathStub) GetWindowCostBatch(ctx context.Context, accountIDs []int64) (map[int64]float64, error) {
  72. if s.batchErr != nil {
  73. return nil, s.batchErr
  74. }
  75. out := make(map[int64]float64, len(accountIDs))
  76. for _, id := range accountIDs {
  77. if v, ok := s.batchData[id]; ok {
  78. out[id] = v
  79. }
  80. }
  81. return out, nil
  82. }
  83. func (s *sessionLimitCacheHotpathStub) SetWindowCost(ctx context.Context, accountID int64, cost float64) error {
  84. if s.setErr != nil {
  85. return s.setErr
  86. }
  87. if s.setData == nil {
  88. s.setData = make(map[int64]float64)
  89. }
  90. s.setData[accountID] = cost
  91. return nil
  92. }
  93. type modelsListAccountRepoStub struct {
  94. AccountRepository
  95. byGroup map[int64][]Account
  96. all []Account
  97. err error
  98. listByGroupCalls atomic.Int64
  99. listAllCalls atomic.Int64
  100. }
  101. type stickyGatewayCacheHotpathStub struct {
  102. GatewayCache
  103. stickyID int64
  104. getCalls atomic.Int64
  105. }
  106. func (s *stickyGatewayCacheHotpathStub) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
  107. s.getCalls.Add(1)
  108. if s.stickyID > 0 {
  109. return s.stickyID, nil
  110. }
  111. return 0, errors.New("not found")
  112. }
  113. func (s *stickyGatewayCacheHotpathStub) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
  114. return nil
  115. }
  116. func (s *stickyGatewayCacheHotpathStub) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
  117. return nil
  118. }
  119. func (s *stickyGatewayCacheHotpathStub) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
  120. return nil
  121. }
  122. func (s *stickyGatewayCacheHotpathStub) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
  123. return nil
  124. }
  125. func (s *stickyGatewayCacheHotpathStub) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
  126. return nil, nil
  127. }
  128. func (s *stickyGatewayCacheHotpathStub) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
  129. return true, nil
  130. }
  131. func (s *stickyGatewayCacheHotpathStub) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
  132. return nil
  133. }
  134. func (s *stickyGatewayCacheHotpathStub) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
  135. return nil
  136. }
  137. func (s *stickyGatewayCacheHotpathStub) GetReasoningContent(_ context.Context, _ string) (string, error) {
  138. return "", ErrReasoningContentNotFound
  139. }
  140. func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
  141. s.listByGroupCalls.Add(1)
  142. if s.err != nil {
  143. return nil, s.err
  144. }
  145. accounts, ok := s.byGroup[groupID]
  146. if !ok {
  147. return nil, nil
  148. }
  149. out := make([]Account, len(accounts))
  150. copy(out, accounts)
  151. return out, nil
  152. }
  153. func (s *modelsListAccountRepoStub) ListSchedulable(ctx context.Context) ([]Account, error) {
  154. s.listAllCalls.Add(1)
  155. if s.err != nil {
  156. return nil, s.err
  157. }
  158. out := make([]Account, len(s.all))
  159. copy(out, s.all)
  160. return out, nil
  161. }
  162. func resetGatewayHotpathStatsForTest() {
  163. windowCostPrefetchCacheHitTotal.Store(0)
  164. windowCostPrefetchCacheMissTotal.Store(0)
  165. windowCostPrefetchBatchSQLTotal.Store(0)
  166. windowCostPrefetchFallbackTotal.Store(0)
  167. windowCostPrefetchErrorTotal.Store(0)
  168. userGroupRateCacheHitTotal.Store(0)
  169. userGroupRateCacheMissTotal.Store(0)
  170. userGroupRateCacheLoadTotal.Store(0)
  171. userGroupRateCacheSFSharedTotal.Store(0)
  172. userGroupRateCacheFallbackTotal.Store(0)
  173. modelsListCacheHitTotal.Store(0)
  174. modelsListCacheMissTotal.Store(0)
  175. modelsListCacheStoreTotal.Store(0)
  176. }
  177. func TestGetUserGroupRateMultiplier_UsesCacheAndSingleflight(t *testing.T) {
  178. resetGatewayHotpathStatsForTest()
  179. rate := 1.7
  180. unblock := make(chan struct{})
  181. repo := &userGroupRateRepoHotpathStub{
  182. rate: &rate,
  183. wait: unblock,
  184. }
  185. svc := &GatewayService{
  186. userGroupRateRepo: repo,
  187. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  188. cfg: &config.Config{
  189. Gateway: config.GatewayConfig{
  190. UserGroupRateCacheTTLSeconds: 30,
  191. },
  192. },
  193. }
  194. const concurrent = 12
  195. results := make([]float64, concurrent)
  196. start := make(chan struct{})
  197. var wg sync.WaitGroup
  198. wg.Add(concurrent)
  199. for i := 0; i < concurrent; i++ {
  200. go func(idx int) {
  201. defer wg.Done()
  202. <-start
  203. results[idx] = svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.2)
  204. }(i)
  205. }
  206. close(start)
  207. // Wait for every caller to have recorded its cache miss before releasing the
  208. // loader. A fixed sleep raced here: a goroutine that reached the cache after
  209. // the singleflight load had already finished got a hit instead of a miss, and
  210. // the miss assertion below saw 11 of 12. The miss counter is the observable
  211. // that says "all callers are now inside the singleflight group".
  212. require.Eventually(t, func() bool {
  213. _, miss, _, _, _ := GatewayUserGroupRateCacheStats()
  214. return miss == int64(concurrent)
  215. }, 5*time.Second, time.Millisecond, "all callers must miss the cache before the loader is released")
  216. close(unblock)
  217. wg.Wait()
  218. for _, got := range results {
  219. require.Equal(t, rate, got)
  220. }
  221. require.Equal(t, int64(1), repo.calls.Load())
  222. // 再次读取应命中缓存,不再回源。
  223. got := svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.2)
  224. require.Equal(t, rate, got)
  225. require.Equal(t, int64(1), repo.calls.Load())
  226. hit, miss, load, sfShared, fallback := GatewayUserGroupRateCacheStats()
  227. require.GreaterOrEqual(t, hit, int64(1))
  228. require.Equal(t, int64(12), miss)
  229. require.Equal(t, int64(1), load)
  230. require.GreaterOrEqual(t, sfShared, int64(1))
  231. require.Equal(t, int64(0), fallback)
  232. }
  233. func TestGetUserGroupRateMultiplier_FallbackOnRepoError(t *testing.T) {
  234. resetGatewayHotpathStatsForTest()
  235. repo := &userGroupRateRepoHotpathStub{
  236. err: errors.New("db down"),
  237. }
  238. svc := &GatewayService{
  239. userGroupRateRepo: repo,
  240. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  241. cfg: &config.Config{
  242. Gateway: config.GatewayConfig{
  243. UserGroupRateCacheTTLSeconds: 30,
  244. },
  245. },
  246. }
  247. got := svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.25)
  248. require.Equal(t, 1.25, got)
  249. require.Equal(t, int64(1), repo.calls.Load())
  250. _, _, _, _, fallback := GatewayUserGroupRateCacheStats()
  251. require.Equal(t, int64(1), fallback)
  252. }
  253. func TestGetUserGroupRateMultiplier_CacheHitAndNilRepo(t *testing.T) {
  254. resetGatewayHotpathStatsForTest()
  255. repo := &userGroupRateRepoHotpathStub{
  256. err: errors.New("should not be called"),
  257. }
  258. svc := &GatewayService{
  259. userGroupRateRepo: repo,
  260. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  261. }
  262. key := "101:202"
  263. svc.userGroupRateCache.Set(key, 2.3, time.Minute)
  264. got := svc.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.1)
  265. require.Equal(t, 2.3, got)
  266. hit, miss, load, _, fallback := GatewayUserGroupRateCacheStats()
  267. require.Equal(t, int64(1), hit)
  268. require.Equal(t, int64(0), miss)
  269. require.Equal(t, int64(0), load)
  270. require.Equal(t, int64(0), fallback)
  271. require.Equal(t, int64(0), repo.calls.Load())
  272. // 无 repo 时直接返回分组默认倍率
  273. svc2 := &GatewayService{
  274. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  275. }
  276. svc2.userGroupRateCache.Set(key, 1.9, time.Minute)
  277. require.Equal(t, 1.9, svc2.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.4))
  278. require.Equal(t, 1.4, svc2.getUserGroupRateMultiplier(context.Background(), 0, 202, 1.4))
  279. svc2.userGroupRateCache.Delete(key)
  280. require.Equal(t, 1.4, svc2.getUserGroupRateMultiplier(context.Background(), 101, 202, 1.4))
  281. }
  282. func TestWithWindowCostPrefetch_BatchReadAndContextReuse(t *testing.T) {
  283. resetGatewayHotpathStatsForTest()
  284. windowStart := time.Now().Add(-30 * time.Minute).Truncate(time.Hour)
  285. windowEnd := windowStart.Add(5 * time.Hour)
  286. accounts := []Account{
  287. {
  288. ID: 1,
  289. Platform: PlatformAnthropic,
  290. Type: AccountTypeOAuth,
  291. Extra: map[string]any{"window_cost_limit": 100.0},
  292. SessionWindowStart: &windowStart,
  293. SessionWindowEnd: &windowEnd,
  294. },
  295. {
  296. ID: 2,
  297. Platform: PlatformAnthropic,
  298. Type: AccountTypeSetupToken,
  299. Extra: map[string]any{"window_cost_limit": 100.0},
  300. SessionWindowStart: &windowStart,
  301. SessionWindowEnd: &windowEnd,
  302. },
  303. {
  304. ID: 3,
  305. Platform: PlatformAnthropic,
  306. Type: AccountTypeAPIKey,
  307. Extra: map[string]any{"window_cost_limit": 100.0},
  308. },
  309. }
  310. cache := &sessionLimitCacheHotpathStub{
  311. batchData: map[int64]float64{
  312. 1: 11.0,
  313. },
  314. }
  315. repo := &usageLogWindowBatchRepoStub{
  316. batchResult: map[int64]*usagestats.AccountStats{
  317. 2: {StandardCost: 22.0},
  318. },
  319. }
  320. svc := &GatewayService{
  321. sessionLimitCache: cache,
  322. usageLogRepo: repo,
  323. }
  324. outCtx := svc.withWindowCostPrefetch(context.Background(), accounts)
  325. require.NotNil(t, outCtx)
  326. cost1, ok1 := windowCostFromPrefetchContext(outCtx, 1)
  327. require.True(t, ok1)
  328. require.Equal(t, 11.0, cost1)
  329. cost2, ok2 := windowCostFromPrefetchContext(outCtx, 2)
  330. require.True(t, ok2)
  331. require.Equal(t, 22.0, cost2)
  332. _, ok3 := windowCostFromPrefetchContext(outCtx, 3)
  333. require.False(t, ok3)
  334. require.Equal(t, int64(1), repo.batchCalls.Load())
  335. require.Equal(t, 22.0, cache.setData[2])
  336. hit, miss, batchSQL, fallback, errCount := GatewayWindowCostPrefetchStats()
  337. require.Equal(t, int64(1), hit)
  338. require.Equal(t, int64(1), miss)
  339. require.Equal(t, int64(1), batchSQL)
  340. require.Equal(t, int64(0), fallback)
  341. require.Equal(t, int64(0), errCount)
  342. }
  343. func TestWithWindowCostPrefetch_AllHitNoSQL(t *testing.T) {
  344. resetGatewayHotpathStatsForTest()
  345. windowStart := time.Now().Add(-30 * time.Minute).Truncate(time.Hour)
  346. windowEnd := windowStart.Add(5 * time.Hour)
  347. accounts := []Account{
  348. {
  349. ID: 1,
  350. Platform: PlatformAnthropic,
  351. Type: AccountTypeOAuth,
  352. Extra: map[string]any{"window_cost_limit": 100.0},
  353. SessionWindowStart: &windowStart,
  354. SessionWindowEnd: &windowEnd,
  355. },
  356. {
  357. ID: 2,
  358. Platform: PlatformAnthropic,
  359. Type: AccountTypeSetupToken,
  360. Extra: map[string]any{"window_cost_limit": 100.0},
  361. SessionWindowStart: &windowStart,
  362. SessionWindowEnd: &windowEnd,
  363. },
  364. }
  365. cache := &sessionLimitCacheHotpathStub{
  366. batchData: map[int64]float64{
  367. 1: 11.0,
  368. 2: 22.0,
  369. },
  370. }
  371. repo := &usageLogWindowBatchRepoStub{}
  372. svc := &GatewayService{
  373. sessionLimitCache: cache,
  374. usageLogRepo: repo,
  375. }
  376. outCtx := svc.withWindowCostPrefetch(context.Background(), accounts)
  377. cost1, ok1 := windowCostFromPrefetchContext(outCtx, 1)
  378. cost2, ok2 := windowCostFromPrefetchContext(outCtx, 2)
  379. require.True(t, ok1)
  380. require.True(t, ok2)
  381. require.Equal(t, 11.0, cost1)
  382. require.Equal(t, 22.0, cost2)
  383. require.Equal(t, int64(0), repo.batchCalls.Load())
  384. require.Equal(t, int64(0), repo.singleCalls.Load())
  385. hit, miss, batchSQL, fallback, errCount := GatewayWindowCostPrefetchStats()
  386. require.Equal(t, int64(2), hit)
  387. require.Equal(t, int64(0), miss)
  388. require.Equal(t, int64(0), batchSQL)
  389. require.Equal(t, int64(0), fallback)
  390. require.Equal(t, int64(0), errCount)
  391. }
  392. func TestWithWindowCostPrefetch_BatchErrorFallbackSingleQuery(t *testing.T) {
  393. resetGatewayHotpathStatsForTest()
  394. windowStart := time.Now().Add(-30 * time.Minute).Truncate(time.Hour)
  395. windowEnd := windowStart.Add(5 * time.Hour)
  396. accounts := []Account{
  397. {
  398. ID: 2,
  399. Platform: PlatformAnthropic,
  400. Type: AccountTypeSetupToken,
  401. Extra: map[string]any{"window_cost_limit": 100.0},
  402. SessionWindowStart: &windowStart,
  403. SessionWindowEnd: &windowEnd,
  404. },
  405. }
  406. cache := &sessionLimitCacheHotpathStub{}
  407. repo := &usageLogWindowBatchRepoStub{
  408. batchErr: errors.New("batch failed"),
  409. singleResult: map[int64]*usagestats.AccountStats{
  410. 2: {StandardCost: 33.0},
  411. },
  412. }
  413. svc := &GatewayService{
  414. sessionLimitCache: cache,
  415. usageLogRepo: repo,
  416. }
  417. outCtx := svc.withWindowCostPrefetch(context.Background(), accounts)
  418. cost, ok := windowCostFromPrefetchContext(outCtx, 2)
  419. require.True(t, ok)
  420. require.Equal(t, 33.0, cost)
  421. require.Equal(t, int64(1), repo.batchCalls.Load())
  422. require.Equal(t, int64(1), repo.singleCalls.Load())
  423. _, _, _, fallback, errCount := GatewayWindowCostPrefetchStats()
  424. require.Equal(t, int64(1), fallback)
  425. require.Equal(t, int64(1), errCount)
  426. }
  427. func TestGetAvailableModels_UsesShortCacheAndSupportsInvalidation(t *testing.T) {
  428. resetGatewayHotpathStatsForTest()
  429. groupID := int64(9)
  430. repo := &modelsListAccountRepoStub{
  431. byGroup: map[int64][]Account{
  432. groupID: {
  433. {
  434. ID: 1,
  435. Platform: PlatformAnthropic,
  436. Credentials: map[string]any{
  437. "model_mapping": map[string]any{
  438. "claude-3-5-sonnet": "claude-3-5-sonnet",
  439. "claude-3-5-haiku": "claude-3-5-haiku",
  440. },
  441. },
  442. },
  443. {
  444. ID: 2,
  445. Platform: PlatformGemini,
  446. Credentials: map[string]any{
  447. "model_mapping": map[string]any{
  448. "gemini-2.5-pro": "gemini-2.5-pro",
  449. },
  450. },
  451. },
  452. },
  453. },
  454. }
  455. svc := &GatewayService{
  456. accountRepo: repo,
  457. modelsListCache: gocache.New(time.Minute, time.Minute),
  458. modelsListCacheTTL: time.Minute,
  459. }
  460. models1 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
  461. require.Equal(t, []string{"claude-3-5-haiku", "claude-3-5-sonnet"}, models1)
  462. require.Equal(t, int64(1), repo.listByGroupCalls.Load())
  463. // TTL 内再次请求应命中缓存,不回源。
  464. models2 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
  465. require.Equal(t, models1, models2)
  466. require.Equal(t, int64(1), repo.listByGroupCalls.Load())
  467. // 更新仓储数据,但缓存未失效前应继续返回旧值。
  468. repo.byGroup[groupID] = []Account{
  469. {
  470. ID: 3,
  471. Platform: PlatformAnthropic,
  472. Credentials: map[string]any{
  473. "model_mapping": map[string]any{
  474. "claude-3-7-sonnet": "claude-3-7-sonnet",
  475. },
  476. },
  477. },
  478. }
  479. models3 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
  480. require.Equal(t, []string{"claude-3-5-haiku", "claude-3-5-sonnet"}, models3)
  481. require.Equal(t, int64(1), repo.listByGroupCalls.Load())
  482. svc.InvalidateAvailableModelsCache(&groupID, PlatformAnthropic)
  483. models4 := svc.GetAvailableModels(context.Background(), &groupID, PlatformAnthropic)
  484. require.Equal(t, []string{"claude-3-7-sonnet"}, models4)
  485. require.Equal(t, int64(2), repo.listByGroupCalls.Load())
  486. hit, miss, store := GatewayModelsListCacheStats()
  487. require.Equal(t, int64(2), hit)
  488. require.Equal(t, int64(2), miss)
  489. require.Equal(t, int64(2), store)
  490. }
  491. // Scenario: 账号模型变更会失效所属平台缓存
  492. func TestResolveCompositeModelOwnershipUsesModelsCacheInvalidation(t *testing.T) {
  493. groupID := int64(9)
  494. repo := &modelsListAccountRepoStub{
  495. byGroup: map[int64][]Account{
  496. groupID: {{
  497. ID: 1,
  498. Platform: PlatformDeepseek,
  499. Credentials: map[string]any{"model_mapping": map[string]any{"company-model": "deepseek-v4-pro"}},
  500. }},
  501. },
  502. }
  503. svc := &GatewayService{
  504. accountRepo: repo,
  505. modelsListCache: gocache.New(time.Minute, time.Minute),
  506. modelsListCacheTTL: time.Minute,
  507. }
  508. first, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "company-model")
  509. require.NoError(t, err)
  510. require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformDeepseek, Matched: true}, first)
  511. require.Equal(t, int64(1), repo.listByGroupCalls.Load())
  512. repo.byGroup[groupID] = []Account{{
  513. ID: 2,
  514. Platform: PlatformOpenAI,
  515. Credentials: map[string]any{"model_mapping": map[string]any{"company-model": "gpt-5"}},
  516. }}
  517. cached, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "company-model")
  518. require.NoError(t, err)
  519. require.Equal(t, first, cached)
  520. require.Equal(t, int64(1), repo.listByGroupCalls.Load())
  521. svc.InvalidateAvailableModelsCache(&groupID, PlatformDeepseek)
  522. refreshed, err := svc.resolveCompositeModelOwnership(context.Background(), groupID, "company-model")
  523. require.NoError(t, err)
  524. require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, refreshed)
  525. require.Equal(t, int64(2), repo.listByGroupCalls.Load())
  526. }
  527. func TestGetAvailableModels_ErrorAndGlobalListBranches(t *testing.T) {
  528. resetGatewayHotpathStatsForTest()
  529. errRepo := &modelsListAccountRepoStub{
  530. err: errors.New("db error"),
  531. }
  532. svcErr := &GatewayService{
  533. accountRepo: errRepo,
  534. modelsListCache: gocache.New(time.Minute, time.Minute),
  535. modelsListCacheTTL: time.Minute,
  536. }
  537. require.Nil(t, svcErr.GetAvailableModels(context.Background(), nil, ""))
  538. okRepo := &modelsListAccountRepoStub{
  539. all: []Account{
  540. {
  541. ID: 1,
  542. Platform: PlatformAnthropic,
  543. Credentials: map[string]any{
  544. "model_mapping": map[string]any{
  545. "claude-3-5-sonnet": "claude-3-5-sonnet",
  546. },
  547. },
  548. },
  549. {
  550. ID: 2,
  551. Platform: PlatformGemini,
  552. Credentials: map[string]any{
  553. "model_mapping": map[string]any{
  554. "gemini-2.5-pro": "gemini-2.5-pro",
  555. },
  556. },
  557. },
  558. },
  559. }
  560. svcOK := &GatewayService{
  561. accountRepo: okRepo,
  562. modelsListCache: gocache.New(time.Minute, time.Minute),
  563. modelsListCacheTTL: time.Minute,
  564. }
  565. models := svcOK.GetAvailableModels(context.Background(), nil, "")
  566. require.Equal(t, []string{"claude-3-5-sonnet", "gemini-2.5-pro"}, models)
  567. require.Equal(t, int64(1), okRepo.listAllCalls.Load())
  568. }
  569. func TestGetAvailableModels_OpenAIPassthroughUsesDefaultFallback(t *testing.T) {
  570. groupID := int64(10)
  571. tests := []struct {
  572. name string
  573. accounts []Account
  574. want []string
  575. }{
  576. {
  577. name: "passthrough only ignores stale mapping",
  578. accounts: []Account{
  579. {
  580. ID: 1,
  581. Platform: PlatformOpenAI,
  582. Credentials: map[string]any{"model_mapping": map[string]any{"stale-model": "upstream-model"}},
  583. Extra: map[string]any{"openai_passthrough": true},
  584. },
  585. },
  586. want: nil,
  587. },
  588. {
  589. name: "passthrough wins over ordinary account mapping",
  590. accounts: []Account{
  591. {
  592. ID: 2,
  593. Platform: PlatformOpenAI,
  594. Credentials: map[string]any{"model_mapping": map[string]any{"configured-model": "configured-upstream"}},
  595. },
  596. {
  597. ID: 3,
  598. Platform: PlatformOpenAI,
  599. Credentials: map[string]any{"model_mapping": map[string]any{"stale-model": "upstream-model"}},
  600. Extra: map[string]any{"openai_passthrough": true},
  601. },
  602. },
  603. want: nil,
  604. },
  605. {
  606. name: "ordinary accounts preserve mapped whitelist",
  607. accounts: []Account{
  608. {
  609. ID: 4,
  610. Platform: PlatformOpenAI,
  611. Credentials: map[string]any{"model_mapping": map[string]any{"configured-model": "configured-upstream"}},
  612. },
  613. },
  614. want: []string{"configured-model"},
  615. },
  616. }
  617. for _, tt := range tests {
  618. t.Run(tt.name, func(t *testing.T) {
  619. repo := &modelsListAccountRepoStub{byGroup: map[int64][]Account{groupID: tt.accounts}}
  620. svc := &GatewayService{
  621. accountRepo: repo,
  622. modelsListCache: gocache.New(time.Minute, time.Minute),
  623. modelsListCacheTTL: time.Minute,
  624. }
  625. require.Equal(t, tt.want, svc.GetAvailableModels(context.Background(), &groupID, PlatformOpenAI))
  626. })
  627. }
  628. }
  629. func TestGetAvailableModels_GlobalListPreservesMappedModelsWithOpenAIPassthrough(t *testing.T) {
  630. groupID := int64(11)
  631. repo := &modelsListAccountRepoStub{
  632. byGroup: map[int64][]Account{
  633. groupID: {
  634. {
  635. ID: 1,
  636. Platform: PlatformOpenAI,
  637. Extra: map[string]any{"openai_passthrough": true},
  638. },
  639. {
  640. ID: 2,
  641. Platform: PlatformAnthropic,
  642. Credentials: map[string]any{"model_mapping": map[string]any{"claude-mapped": "claude-upstream"}},
  643. },
  644. },
  645. },
  646. }
  647. svc := &GatewayService{
  648. accountRepo: repo,
  649. modelsListCache: gocache.New(time.Minute, time.Minute),
  650. modelsListCacheTTL: time.Minute,
  651. }
  652. require.Equal(t, []string{"claude-mapped"}, svc.GetAvailableModels(context.Background(), &groupID, ""))
  653. }
  654. func TestGatewayHotpathHelpers_CacheTTLAndStickyContext(t *testing.T) {
  655. t.Run("resolve_user_group_rate_cache_ttl", func(t *testing.T) {
  656. require.Equal(t, defaultUserGroupRateCacheTTL, resolveUserGroupRateCacheTTL(nil))
  657. cfg := &config.Config{
  658. Gateway: config.GatewayConfig{
  659. UserGroupRateCacheTTLSeconds: 45,
  660. },
  661. }
  662. require.Equal(t, 45*time.Second, resolveUserGroupRateCacheTTL(cfg))
  663. })
  664. t.Run("resolve_models_list_cache_ttl", func(t *testing.T) {
  665. require.Equal(t, defaultModelsListCacheTTL, resolveModelsListCacheTTL(nil))
  666. cfg := &config.Config{
  667. Gateway: config.GatewayConfig{
  668. ModelsListCacheTTLSeconds: 20,
  669. },
  670. }
  671. require.Equal(t, 20*time.Second, resolveModelsListCacheTTL(cfg))
  672. })
  673. t.Run("prefetched_sticky_account_id_from_context", func(t *testing.T) {
  674. require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(context.TODO(), nil))
  675. require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(context.Background(), nil))
  676. ctx := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, int64(123))
  677. ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(0))
  678. require.Equal(t, int64(123), prefetchedStickyAccountIDFromContext(ctx, nil))
  679. groupID := int64(9)
  680. ctx2 := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, 456)
  681. ctx2 = context.WithValue(ctx2, ctxkey.PrefetchedStickyGroupID, groupID)
  682. require.Equal(t, int64(456), prefetchedStickyAccountIDFromContext(ctx2, &groupID))
  683. ctx3 := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, "invalid")
  684. ctx3 = context.WithValue(ctx3, ctxkey.PrefetchedStickyGroupID, groupID)
  685. require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(ctx3, &groupID))
  686. ctx4 := context.WithValue(context.Background(), ctxkey.PrefetchedStickyAccountID, int64(789))
  687. ctx4 = context.WithValue(ctx4, ctxkey.PrefetchedStickyGroupID, int64(10))
  688. require.Equal(t, int64(0), prefetchedStickyAccountIDFromContext(ctx4, &groupID))
  689. })
  690. t.Run("window_cost_from_prefetch_context", func(t *testing.T) {
  691. require.Equal(t, false, func() bool {
  692. _, ok := windowCostFromPrefetchContext(context.TODO(), 0)
  693. return ok
  694. }())
  695. require.Equal(t, false, func() bool {
  696. _, ok := windowCostFromPrefetchContext(context.Background(), 1)
  697. return ok
  698. }())
  699. ctx := context.WithValue(context.Background(), windowCostPrefetchContextKey, map[int64]float64{
  700. 9: 12.34,
  701. })
  702. cost, ok := windowCostFromPrefetchContext(ctx, 9)
  703. require.True(t, ok)
  704. require.Equal(t, 12.34, cost)
  705. })
  706. }
  707. func TestInvalidateAvailableModelsCache_ByDimensions(t *testing.T) {
  708. svc := &GatewayService{
  709. modelsListCache: gocache.New(time.Minute, time.Minute),
  710. }
  711. group9 := int64(9)
  712. group10 := int64(10)
  713. svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformAnthropic), []string{"a"}, time.Minute)
  714. svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformGemini), []string{"b"}, time.Minute)
  715. svc.modelsListCache.Set(modelsListCacheKey(&group10, PlatformAnthropic), []string{"c"}, time.Minute)
  716. svc.modelsListCache.Set("invalid-key", []string{"d"}, time.Minute)
  717. t.Run("invalidate_group_and_platform", func(t *testing.T) {
  718. svc.InvalidateAvailableModelsCache(&group9, PlatformAnthropic)
  719. _, found := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformAnthropic))
  720. require.False(t, found)
  721. _, stillFound := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformGemini))
  722. require.True(t, stillFound)
  723. })
  724. t.Run("invalidate_group_only", func(t *testing.T) {
  725. svc.InvalidateAvailableModelsCache(&group9, "")
  726. _, foundA := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformAnthropic))
  727. _, foundB := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformGemini))
  728. require.False(t, foundA)
  729. require.False(t, foundB)
  730. _, foundOtherGroup := svc.modelsListCache.Get(modelsListCacheKey(&group10, PlatformAnthropic))
  731. require.True(t, foundOtherGroup)
  732. })
  733. t.Run("invalidate_platform_only", func(t *testing.T) {
  734. // 重建数据后仅按 platform 失效
  735. svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformAnthropic), []string{"a"}, time.Minute)
  736. svc.modelsListCache.Set(modelsListCacheKey(&group9, PlatformGemini), []string{"b"}, time.Minute)
  737. svc.modelsListCache.Set(modelsListCacheKey(&group10, PlatformAnthropic), []string{"c"}, time.Minute)
  738. svc.InvalidateAvailableModelsCache(nil, PlatformAnthropic)
  739. _, found9Anthropic := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformAnthropic))
  740. _, found10Anthropic := svc.modelsListCache.Get(modelsListCacheKey(&group10, PlatformAnthropic))
  741. _, found9Gemini := svc.modelsListCache.Get(modelsListCacheKey(&group9, PlatformGemini))
  742. require.False(t, found9Anthropic)
  743. require.False(t, found10Anthropic)
  744. require.True(t, found9Gemini)
  745. })
  746. }
  747. func TestSelectAccountWithLoadAwareness_StickyReadReuse(t *testing.T) {
  748. now := time.Now().Add(-time.Minute)
  749. account := Account{
  750. ID: 88,
  751. Platform: PlatformAnthropic,
  752. Type: AccountTypeAPIKey,
  753. Status: StatusActive,
  754. Schedulable: true,
  755. Concurrency: 4,
  756. Priority: 1,
  757. LastUsedAt: &now,
  758. }
  759. repo := stubOpenAIAccountRepo{accounts: []Account{account}}
  760. concurrency := NewConcurrencyService(stubConcurrencyCache{})
  761. cfg := &config.Config{
  762. RunMode: config.RunModeStandard,
  763. Gateway: config.GatewayConfig{
  764. Scheduling: config.GatewaySchedulingConfig{
  765. LoadBatchEnabled: true,
  766. StickySessionMaxWaiting: 3,
  767. StickySessionWaitTimeout: time.Second,
  768. FallbackWaitTimeout: time.Second,
  769. FallbackMaxWaiting: 10,
  770. },
  771. },
  772. }
  773. baseCtx := context.WithValue(context.Background(), ctxkey.ForcePlatform, PlatformAnthropic)
  774. t.Run("without_prefetch_reads_cache_once", func(t *testing.T) {
  775. cache := &stickyGatewayCacheHotpathStub{stickyID: account.ID}
  776. svc := &GatewayService{
  777. accountRepo: repo,
  778. cache: cache,
  779. cfg: cfg,
  780. concurrencyService: concurrency,
  781. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  782. modelsListCache: gocache.New(time.Minute, time.Minute),
  783. modelsListCacheTTL: time.Minute,
  784. }
  785. result, err := svc.SelectAccountWithLoadAwareness(baseCtx, nil, "sess-hash", "", nil, "", int64(0))
  786. require.NoError(t, err)
  787. require.NotNil(t, result)
  788. require.NotNil(t, result.Account)
  789. require.Equal(t, account.ID, result.Account.ID)
  790. require.Equal(t, int64(1), cache.getCalls.Load())
  791. })
  792. t.Run("with_prefetch_skips_cache_read", func(t *testing.T) {
  793. cache := &stickyGatewayCacheHotpathStub{stickyID: account.ID}
  794. svc := &GatewayService{
  795. accountRepo: repo,
  796. cache: cache,
  797. cfg: cfg,
  798. concurrencyService: concurrency,
  799. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  800. modelsListCache: gocache.New(time.Minute, time.Minute),
  801. modelsListCacheTTL: time.Minute,
  802. }
  803. ctx := context.WithValue(baseCtx, ctxkey.PrefetchedStickyAccountID, account.ID)
  804. ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(0))
  805. result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "", int64(0))
  806. require.NoError(t, err)
  807. require.NotNil(t, result)
  808. require.NotNil(t, result.Account)
  809. require.Equal(t, account.ID, result.Account.ID)
  810. require.Equal(t, int64(0), cache.getCalls.Load())
  811. })
  812. t.Run("with_prefetch_group_mismatch_reads_cache", func(t *testing.T) {
  813. cache := &stickyGatewayCacheHotpathStub{stickyID: account.ID}
  814. svc := &GatewayService{
  815. accountRepo: repo,
  816. cache: cache,
  817. cfg: cfg,
  818. concurrencyService: concurrency,
  819. userGroupRateCache: gocache.New(time.Minute, time.Minute),
  820. modelsListCache: gocache.New(time.Minute, time.Minute),
  821. modelsListCacheTTL: time.Minute,
  822. }
  823. ctx := context.WithValue(baseCtx, ctxkey.PrefetchedStickyAccountID, int64(999))
  824. ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(77))
  825. result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "", int64(0))
  826. require.NoError(t, err)
  827. require.NotNil(t, result)
  828. require.NotNil(t, result.Account)
  829. require.Equal(t, account.ID, result.Account.ID)
  830. require.Equal(t, int64(1), cache.getCalls.Load())
  831. })
  832. }