token_refresh_pool_health_test.go 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "reflect"
  7. "sort"
  8. "sync"
  9. "sync/atomic"
  10. "testing"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/config"
  13. "github.com/stretchr/testify/require"
  14. )
  15. type poolHealthAccountRepo struct {
  16. AccountRepository
  17. mu sync.Mutex
  18. pages map[int64][]Account
  19. requests []OAuthRefreshPageOptions
  20. updatedCredentialIDs []int64
  21. setErrorCalls int
  22. setTempUnschedCalls int
  23. getByIDErr error
  24. }
  25. func (r *poolHealthAccountRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
  26. if r.getByIDErr != nil {
  27. return nil, r.getByIDErr
  28. }
  29. return nil, ErrAccountNotFound
  30. }
  31. func (r *poolHealthAccountRepo) ListOAuthRefreshCandidatePage(_ context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error) {
  32. r.mu.Lock()
  33. defer r.mu.Unlock()
  34. r.requests = append(r.requests, options)
  35. accounts := append([]Account(nil), r.pages[options.AfterID]...)
  36. page := &OAuthRefreshCandidatePage{Accounts: accounts, HasMore: len(accounts) == options.Limit}
  37. if len(accounts) > 0 {
  38. page.NextAfterID = accounts[len(accounts)-1].ID
  39. }
  40. return page, nil
  41. }
  42. func (r *poolHealthAccountRepo) UpdateCredentials(_ context.Context, id int64, _ map[string]any) error {
  43. r.mu.Lock()
  44. defer r.mu.Unlock()
  45. r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
  46. return nil
  47. }
  48. func (r *poolHealthAccountRepo) UpdateGrokOAuthCredentialsIfUnchanged(
  49. _ context.Context,
  50. id int64,
  51. _ map[string]any,
  52. _ *int64,
  53. _ map[string]any,
  54. ) (bool, error) {
  55. r.mu.Lock()
  56. defer r.mu.Unlock()
  57. r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
  58. return true, nil
  59. }
  60. func (r *poolHealthAccountRepo) SetError(context.Context, int64, string) error {
  61. r.mu.Lock()
  62. defer r.mu.Unlock()
  63. r.setErrorCalls++
  64. return nil
  65. }
  66. func (r *poolHealthAccountRepo) SetGrokOAuthErrorIfCredentialsUnchanged(context.Context, int64, map[string]any, string) (bool, error) {
  67. return false, nil
  68. }
  69. func (r *poolHealthAccountRepo) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, string) (bool, error) {
  70. r.mu.Lock()
  71. defer r.mu.Unlock()
  72. r.setErrorCalls++
  73. return true, nil
  74. }
  75. func (r *poolHealthAccountRepo) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, time.Time, string) (bool, error) {
  76. r.mu.Lock()
  77. defer r.mu.Unlock()
  78. r.setTempUnschedCalls++
  79. return true, nil
  80. }
  81. func (r *poolHealthAccountRepo) SetTempUnschedulable(context.Context, int64, time.Time, string) error {
  82. r.mu.Lock()
  83. defer r.mu.Unlock()
  84. r.setTempUnschedCalls++
  85. return nil
  86. }
  87. func (r *poolHealthAccountRepo) snapshot() ([]OAuthRefreshPageOptions, []int64, int, int) {
  88. r.mu.Lock()
  89. defer r.mu.Unlock()
  90. return append([]OAuthRefreshPageOptions(nil), r.requests...), append([]int64(nil), r.updatedCredentialIDs...), r.setErrorCalls, r.setTempUnschedCalls
  91. }
  92. type poolHealthRefresher struct {
  93. err error
  94. delay time.Duration
  95. startDelays []time.Duration
  96. ignoreContext bool
  97. cancel context.CancelFunc
  98. newCredentials map[string]any
  99. calls atomic.Int64
  100. active atomic.Int64
  101. maxActive atomic.Int64
  102. startMu sync.Mutex
  103. startTimes []time.Time
  104. }
  105. type countingRefreshAttemptGate struct {
  106. calls atomic.Int64
  107. }
  108. type rejectedRefreshAttemptGate struct {
  109. err error
  110. }
  111. type poolHealthTokenCacheStub struct {
  112. GeminiTokenCache
  113. }
  114. type tripBeforeRateAdmissionGate struct {
  115. state *tokenRefreshProviderState
  116. }
  117. func (g *tripBeforeRateAdmissionGate) acquire(ctx context.Context) (func(), error) {
  118. release, err := g.state.acquire(ctx)
  119. if err != nil {
  120. return nil, err
  121. }
  122. g.state.mu.Lock()
  123. g.state.tripped = true
  124. g.state.mu.Unlock()
  125. return release, nil
  126. }
  127. func (g *tripBeforeRateAdmissionGate) acquireRate(ctx context.Context) (func(), error) {
  128. return g.state.acquireRate(ctx)
  129. }
  130. type breakerTripAccountRepo struct {
  131. *productionPathRateRepo
  132. setErrorCalls atomic.Int64
  133. setTempCalls atomic.Int64
  134. }
  135. func (r *breakerTripAccountRepo) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, string) (bool, error) {
  136. r.setErrorCalls.Add(1)
  137. return true, nil
  138. }
  139. func (r *breakerTripAccountRepo) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, time.Time, string) (bool, error) {
  140. r.setTempCalls.Add(1)
  141. return true, nil
  142. }
  143. func (g *rejectedRefreshAttemptGate) acquire(context.Context) (func(), error) {
  144. return nil, g.err
  145. }
  146. type productionPathRateRepo struct {
  147. AccountRepository
  148. mu sync.Mutex
  149. accounts map[int64]*Account
  150. }
  151. func (r *productionPathRateRepo) GetByID(_ context.Context, id int64) (*Account, error) {
  152. r.mu.Lock()
  153. defer r.mu.Unlock()
  154. account := r.accounts[id]
  155. if account == nil {
  156. return nil, ErrAccountNotFound
  157. }
  158. return snapshotOAuthRefreshAccount(account), nil
  159. }
  160. func (r *productionPathRateRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
  161. r.mu.Lock()
  162. defer r.mu.Unlock()
  163. account := r.accounts[id]
  164. if account == nil {
  165. return ErrAccountNotFound
  166. }
  167. account.Credentials = shallowCopyMap(credentials)
  168. return nil
  169. }
  170. func (r *productionPathRateRepo) UpdateGrokOAuthCredentialsIfUnchanged(
  171. _ context.Context,
  172. id int64,
  173. expectedCredentials map[string]any,
  174. expectedProxyID *int64,
  175. credentials map[string]any,
  176. ) (bool, error) {
  177. r.mu.Lock()
  178. defer r.mu.Unlock()
  179. account := r.accounts[id]
  180. if account == nil || !reflect.DeepEqual(account.Credentials, expectedCredentials) ||
  181. !reflect.DeepEqual(account.ProxyID, expectedProxyID) {
  182. return false, nil
  183. }
  184. account.Credentials = shallowCopyMap(credentials)
  185. return true, nil
  186. }
  187. type productionPathRefreshStart struct {
  188. accountID int64
  189. at time.Time
  190. }
  191. type productionPathRateExecutor struct {
  192. firstStarted chan struct{}
  193. releaseFirst chan struct{}
  194. calls atomic.Int64
  195. startMu sync.Mutex
  196. starts []productionPathRefreshStart
  197. }
  198. func (e *productionPathRateExecutor) CacheKey(account *Account) string {
  199. return fmt.Sprintf("production-path-rate:%d", account.ID)
  200. }
  201. func (e *productionPathRateExecutor) CanRefresh(account *Account) bool {
  202. return account != nil && account.IsGrokOAuth()
  203. }
  204. func (e *productionPathRateExecutor) NeedsRefresh(account *Account, _ time.Duration) bool {
  205. needsRefresh, _ := account.Credentials["needs_refresh"].(bool)
  206. return needsRefresh
  207. }
  208. func (e *productionPathRateExecutor) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
  209. call := e.calls.Add(1)
  210. e.startMu.Lock()
  211. e.starts = append(e.starts, productionPathRefreshStart{accountID: account.ID, at: time.Now()})
  212. e.startMu.Unlock()
  213. if call == 1 {
  214. close(e.firstStarted)
  215. select {
  216. case <-ctx.Done():
  217. return nil, ctx.Err()
  218. case <-e.releaseFirst:
  219. }
  220. }
  221. return map[string]any{
  222. "access_token": fmt.Sprintf("fresh-access-%d", account.ID),
  223. "refresh_token": fmt.Sprintf("fresh-refresh-%d", account.ID),
  224. "needs_refresh": false,
  225. }, nil
  226. }
  227. func (e *productionPathRateExecutor) startsSnapshot() []productionPathRefreshStart {
  228. e.startMu.Lock()
  229. defer e.startMu.Unlock()
  230. return append([]productionPathRefreshStart(nil), e.starts...)
  231. }
  232. func (g *countingRefreshAttemptGate) acquire(ctx context.Context) (func(), error) {
  233. if err := ctx.Err(); err != nil {
  234. return nil, err
  235. }
  236. g.calls.Add(1)
  237. return func() {}, nil
  238. }
  239. func (r *poolHealthRefresher) CacheKey(account *Account) string {
  240. return fmt.Sprintf("pool-health:%d", account.ID)
  241. }
  242. func (r *poolHealthRefresher) CanRefresh(account *Account) bool {
  243. return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth
  244. }
  245. func (r *poolHealthRefresher) NeedsRefresh(*Account, time.Duration) bool { return true }
  246. func (r *poolHealthRefresher) Refresh(ctx context.Context, _ *Account) (map[string]any, error) {
  247. r.calls.Add(1)
  248. active := r.active.Add(1)
  249. defer r.active.Add(-1)
  250. r.startMu.Lock()
  251. startIndex := len(r.startTimes)
  252. r.startTimes = append(r.startTimes, time.Now())
  253. delay := r.delay
  254. if startIndex < len(r.startDelays) {
  255. delay = r.startDelays[startIndex]
  256. }
  257. r.startMu.Unlock()
  258. for {
  259. maxActive := r.maxActive.Load()
  260. if active <= maxActive || r.maxActive.CompareAndSwap(maxActive, active) {
  261. break
  262. }
  263. }
  264. if r.cancel != nil {
  265. r.cancel()
  266. }
  267. if delay > 0 {
  268. if r.ignoreContext {
  269. time.Sleep(delay)
  270. } else {
  271. timer := time.NewTimer(delay)
  272. defer timer.Stop()
  273. select {
  274. case <-ctx.Done():
  275. return nil, ctx.Err()
  276. case <-timer.C:
  277. }
  278. }
  279. }
  280. if r.err != nil {
  281. return nil, r.err
  282. }
  283. if r.newCredentials != nil {
  284. credentials := make(map[string]any, len(r.newCredentials))
  285. for key, value := range r.newCredentials {
  286. credentials[key] = value
  287. }
  288. return credentials, nil
  289. }
  290. return map[string]any{"access_token": "new-token", "refresh_token": "new-refresh-token"}, nil
  291. }
  292. func (r *poolHealthRefresher) startsSnapshot() []time.Time {
  293. r.startMu.Lock()
  294. defer r.startMu.Unlock()
  295. return append([]time.Time(nil), r.startTimes...)
  296. }
  297. func grokPoolAccount(id int64) Account {
  298. return Account{
  299. ID: id,
  300. Platform: PlatformGrok,
  301. Type: AccountTypeOAuth,
  302. Status: StatusActive,
  303. Credentials: map[string]any{
  304. "access_token": "old-token",
  305. "refresh_token": "refresh-token",
  306. },
  307. }
  308. }
  309. func newPoolHealthService(repo *poolHealthAccountRepo, refresher *poolHealthRefresher, cfg config.TokenRefreshConfig) *TokenRefreshService {
  310. return &TokenRefreshService{
  311. accountRepo: repo,
  312. candidatePager: repo,
  313. registrations: []tokenRefreshRegistration{{
  314. platform: PlatformGrok,
  315. refresher: refresher,
  316. executor: refresher,
  317. }},
  318. refreshPolicy: DefaultBackgroundRefreshPolicy(),
  319. cfg: &cfg,
  320. }
  321. }
  322. func TestTokenRefreshService_RegistrationsAreCandidateEligibilitySource(t *testing.T) {
  323. cfg := &config.Config{}
  324. svc := NewTokenRefreshService(nil, nil, nil, nil, nil, nil, nil, cfg, nil)
  325. require.Equal(t, []string{
  326. PlatformAnthropic,
  327. PlatformOpenAI,
  328. PlatformGemini,
  329. PlatformAntigravity,
  330. PlatformGrok,
  331. }, svc.eligiblePlatforms())
  332. require.Len(t, svc.registrations, 5)
  333. for _, registration := range svc.registrations {
  334. require.NotNil(t, registration.refresher)
  335. require.NotNil(t, registration.executor)
  336. }
  337. }
  338. func TestTokenRefreshService_ProcessRefreshPagesByStableCursor(t *testing.T) {
  339. repo := &poolHealthAccountRepo{pages: map[int64][]Account{
  340. 0: {grokPoolAccount(1), grokPoolAccount(2)},
  341. 2: {grokPoolAccount(3)},
  342. }}
  343. refresher := &poolHealthRefresher{}
  344. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  345. RefreshBeforeExpiryHours: 1,
  346. MaxRetries: 1,
  347. CandidatePageSize: 2,
  348. ProviderConcurrency: 4,
  349. ProviderQPS: 10000,
  350. AttemptTimeoutSeconds: 1,
  351. CycleTimeoutSeconds: 2,
  352. })
  353. svc.processRefreshContext(context.Background())
  354. requests, updatedIDs, _, _ := repo.snapshot()
  355. require.Len(t, requests, 2)
  356. require.Equal(t, int64(0), requests[0].AfterID)
  357. require.Equal(t, int64(2), requests[1].AfterID)
  358. require.Equal(t, []string{PlatformGrok}, requests[0].Platforms)
  359. require.True(t, requests[0].ActiveOnly)
  360. require.True(t, requests[0].RequireRefreshToken)
  361. require.True(t, requests[0].ExcludeRetryCooldown)
  362. sort.Slice(updatedIDs, func(i, j int) bool { return updatedIDs[i] < updatedIDs[j] })
  363. require.Equal(t, []int64{1, 2, 3}, updatedIDs)
  364. require.Zero(t, svc.candidateAfterID(), "a short final page must wrap the next cycle to the beginning")
  365. }
  366. func TestTokenRefreshService_BoundsPerProviderConcurrency(t *testing.T) {
  367. accounts := make([]Account, 0, 8)
  368. for id := int64(1); id <= 8; id++ {
  369. accounts = append(accounts, grokPoolAccount(id))
  370. }
  371. repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
  372. refresher := &poolHealthRefresher{delay: 20 * time.Millisecond}
  373. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  374. MaxRetries: 1,
  375. CandidatePageSize: 20,
  376. ProviderConcurrency: 2,
  377. ProviderQPS: 10000,
  378. AttemptTimeoutSeconds: 1,
  379. CycleTimeoutSeconds: 2,
  380. })
  381. svc.processRefreshContext(context.Background())
  382. require.Equal(t, int64(8), refresher.calls.Load())
  383. require.Equal(t, int64(2), refresher.maxActive.Load())
  384. }
  385. func TestTokenRefreshRateGate_ReservesSpacedSlotsAndHonorsCancellation(t *testing.T) {
  386. const interval = 25 * time.Millisecond
  387. gate := newTokenRefreshRateGateWithInterval(interval)
  388. base := time.Unix(1_700_000_000, 0)
  389. require.Equal(t, base, gate.reserveSlot(base))
  390. require.Equal(t, base.Add(interval), gate.reserveSlot(base))
  391. require.Equal(t, base.Add(2*interval), gate.reserveSlot(base))
  392. jumped := base.Add(time.Second)
  393. require.Equal(t, jumped, gate.reserveSlot(jumped), "an idle gate should not retain stale delay")
  394. cancelGate := newTokenRefreshRateGateWithInterval(time.Hour)
  395. require.NoError(t, cancelGate.wait(context.Background()), "the first slot is immediately available")
  396. ctx, cancel := context.WithCancel(context.Background())
  397. cancel()
  398. started := time.Now()
  399. require.ErrorIs(t, cancelGate.wait(ctx), context.Canceled)
  400. require.Less(t, time.Since(started), 100*time.Millisecond, "cancellation must not wait for the reserved slot")
  401. }
  402. func TestTokenRefreshService_RetriesAcquireRateSlotPerAttempt(t *testing.T) {
  403. repo := &poolHealthAccountRepo{}
  404. refresher := &poolHealthRefresher{err: errors.New("temporary provider failure")}
  405. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{MaxRetries: 3})
  406. gate := &countingRefreshAttemptGate{}
  407. account := grokPoolAccount(44)
  408. err := svc.refreshWithRetryWithRateGate(context.Background(), &account, refresher, nil, time.Hour, gate)
  409. require.Error(t, err)
  410. require.Equal(t, int64(3), refresher.calls.Load())
  411. require.Equal(t, int64(3), gate.calls.Load(), "every upstream retry must consume a provider rate slot")
  412. }
  413. func TestTokenRefreshService_ProcessProviderAccountsLegacyNilReleaseGateIsSafe(t *testing.T) {
  414. repo := &poolHealthAccountRepo{}
  415. refresher := &poolHealthRefresher{}
  416. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  417. MaxRetries: 1,
  418. ProviderConcurrency: 1,
  419. })
  420. state := &tokenRefreshProviderState{
  421. service: svc,
  422. registration: tokenRefreshRegistration{
  423. platform: PlatformGrok,
  424. refresher: refresher,
  425. // nil executor deliberately exercises the legacy/direct fallback.
  426. executor: nil,
  427. },
  428. // Admission rejection validly returns no release callback. The direct
  429. // fallback must propagate the skip without dereferencing that nil handle.
  430. rateGate: &rejectedRefreshAttemptGate{err: errRefreshSkipped},
  431. poolGate: nil,
  432. }
  433. account := grokPoolAccount(45)
  434. refreshed, skipped, failed := svc.processProviderAccounts(
  435. context.Background(),
  436. state,
  437. []*Account{&account},
  438. time.Hour,
  439. )
  440. require.Zero(t, refreshed)
  441. require.Equal(t, 1, skipped)
  442. require.Zero(t, failed)
  443. require.Zero(t, refresher.calls.Load(), "rejected rate admission must not reach the legacy upstream refresher")
  444. }
  445. func TestTokenRefreshService_ProviderRateGateIsSharedAcrossRuns(t *testing.T) {
  446. svc := &TokenRefreshService{cfg: &config.TokenRefreshConfig{ProviderQPS: 40}}
  447. first := svc.providerRateGate(PlatformGrok)
  448. second := svc.providerRateGate(PlatformGrok)
  449. require.Same(t, first, second, "background cycles and reconciliation must share the process-local provider limiter")
  450. base := time.Unix(1_700_000_000, 0)
  451. require.Equal(t, base, first.reserveSlot(base))
  452. require.Equal(t, base.Add(25*time.Millisecond), second.reserveSlot(base))
  453. }
  454. func TestTokenRefreshService_ProviderConcurrencyGateIsSharedAcrossBackgroundAndConcurrentAdminReconciliation(t *testing.T) {
  455. accounts := []Account{
  456. grokPoolAccount(1),
  457. grokPoolAccount(2),
  458. grokPoolAccount(3),
  459. grokPoolAccount(4),
  460. }
  461. repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
  462. refresher := &poolHealthRefresher{delay: 80 * time.Millisecond}
  463. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  464. RefreshBeforeExpiryHours: 1,
  465. MaxRetries: 1,
  466. CandidatePageSize: 20,
  467. ProviderConcurrency: 2,
  468. ProviderQPS: 100,
  469. ProviderFailureThreshold: 20,
  470. AttemptTimeoutSeconds: 1,
  471. CycleTimeoutSeconds: 3,
  472. })
  473. firstGate := svc.providerConcurrencyGate(PlatformGrok)
  474. require.Same(t, firstGate, svc.providerConcurrencyGate(PlatformGrok))
  475. start := make(chan struct{})
  476. adminErrors := make(chan error, 2)
  477. var wg sync.WaitGroup
  478. wg.Add(3)
  479. go func() {
  480. defer wg.Done()
  481. <-start
  482. svc.processRefreshContext(context.Background())
  483. }()
  484. for i := 0; i < 2; i++ {
  485. go func() {
  486. defer wg.Done()
  487. <-start
  488. _, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 20})
  489. adminErrors <- err
  490. }()
  491. }
  492. close(start)
  493. wg.Wait()
  494. close(adminErrors)
  495. for err := range adminErrors {
  496. require.NoError(t, err)
  497. }
  498. require.Equal(t, int64(12), refresher.calls.Load(), "background and both admin calls must all execute")
  499. require.Equal(t, int64(2), refresher.maxActive.Load(),
  500. "all entry points must share the configured per-provider upstream concurrency cap")
  501. }
  502. func TestTokenRefreshService_SaturatedProviderPreservesConcurrencyAndActualQPSStartSpacing(t *testing.T) {
  503. const (
  504. providerConcurrency = 2
  505. providerQPS = 20
  506. attemptCount = 8
  507. )
  508. repo := &poolHealthAccountRepo{}
  509. refresher := &poolHealthRefresher{
  510. // The first two QPS-spaced attempts finish together. If queued callers
  511. // reserve QPS slots before acquiring provider capacity, two expired
  512. // reservations can then burst upstream at the same time.
  513. startDelays: []time.Duration{
  514. 220 * time.Millisecond,
  515. 170 * time.Millisecond,
  516. 20 * time.Millisecond,
  517. 20 * time.Millisecond,
  518. 20 * time.Millisecond,
  519. 20 * time.Millisecond,
  520. 20 * time.Millisecond,
  521. 20 * time.Millisecond,
  522. },
  523. }
  524. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  525. MaxRetries: 1,
  526. ProviderConcurrency: providerConcurrency,
  527. ProviderQPS: providerQPS,
  528. AttemptTimeoutSeconds: 1,
  529. })
  530. registration := svc.registrations[0]
  531. sharedRateGate := svc.providerRateGate(PlatformGrok)
  532. sharedPoolGate := svc.providerConcurrencyGate(PlatformGrok)
  533. start := make(chan struct{})
  534. errorsCh := make(chan error, attemptCount)
  535. var wg sync.WaitGroup
  536. for i := 0; i < attemptCount; i++ {
  537. account := grokPoolAccount(int64(i + 1))
  538. state := &tokenRefreshProviderState{
  539. service: svc,
  540. registration: registration,
  541. rateGate: sharedRateGate,
  542. poolGate: sharedPoolGate,
  543. }
  544. wg.Add(1)
  545. go func() {
  546. defer wg.Done()
  547. <-start
  548. errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), &account, refresher, nil, time.Hour, state)
  549. }()
  550. }
  551. close(start)
  552. wg.Wait()
  553. close(errorsCh)
  554. for err := range errorsCh {
  555. require.NoError(t, err)
  556. }
  557. require.Equal(t, int64(providerConcurrency), refresher.maxActive.Load(),
  558. "the scripted attempts must actually saturate the provider semaphore")
  559. starts := refresher.startsSnapshot()
  560. require.Len(t, starts, attemptCount)
  561. configuredSpacing := time.Second / time.Duration(providerQPS)
  562. // The floor is a tenth of the configured spacing, not spacing minus a few
  563. // milliseconds. Each start timestamp is taken after the rate gate releases
  564. // the goroutine, so scheduler delay can compress one observed gap without
  565. // the gate having done anything wrong: this assertion failed on a loaded
  566. // machine at 13ms and again at 37ms against a 40ms floor, both times with
  567. // the gate pacing correctly.
  568. //
  569. // A tenth still separates the two cases by a wide margin. Measured with
  570. // providerQPS=20 (50ms spacing), 8 attempts, providerConcurrency=2:
  571. //
  572. // gate at 50ms -> minimum observed gap 49.97ms
  573. // gate at 0 -> minimum observed gap 22µs
  574. //
  575. // So an unpaced gate lands three orders of magnitude below the 5ms floor
  576. // and is still caught, while jitter on a busy machine has room to move.
  577. //
  578. // Do not swap this for an assertion on the total span of the starts: the
  579. // span is dominated by how long each attempt takes under
  580. // providerConcurrency, not by the gate. Same measurement — 471ms paced
  581. // against 241ms unpaced — so a span check passes with the gate disabled and
  582. // tests nothing.
  583. minimumObservedSpacing := configuredSpacing / 10
  584. actualMinimumSpacing := starts[1].Sub(starts[0])
  585. for i := 1; i < len(starts); i++ {
  586. spacing := starts[i].Sub(starts[i-1])
  587. if spacing < actualMinimumSpacing {
  588. actualMinimumSpacing = spacing
  589. }
  590. require.GreaterOrEqualf(t, spacing, minimumObservedSpacing,
  591. "upstream starts %d and %d violated configured QPS spacing", i-1, i)
  592. }
  593. t.Logf("max_active=%d configured_concurrency=%d minimum_start_spacing=%s configured_spacing=%s",
  594. refresher.maxActive.Load(), providerConcurrency, actualMinimumSpacing, configuredSpacing)
  595. }
  596. func TestTokenRefreshService_ProductionPathRatesOnlyActualRefreshAfterSameAccountContention(t *testing.T) {
  597. const interval = 200 * time.Millisecond
  598. accountOne := grokPoolAccount(71)
  599. accountOne.Credentials["needs_refresh"] = true
  600. accountTwo := grokPoolAccount(72)
  601. accountTwo.Credentials["needs_refresh"] = true
  602. firstSelection := snapshotOAuthRefreshAccount(&accountOne)
  603. contendingSelection := snapshotOAuthRefreshAccount(&accountOne)
  604. differentSelection := snapshotOAuthRefreshAccount(&accountTwo)
  605. repo := &productionPathRateRepo{accounts: map[int64]*Account{
  606. accountOne.ID: snapshotOAuthRefreshAccount(&accountOne),
  607. accountTwo.ID: snapshotOAuthRefreshAccount(&accountTwo),
  608. }}
  609. executor := &productionPathRateExecutor{
  610. firstStarted: make(chan struct{}),
  611. releaseFirst: make(chan struct{}),
  612. }
  613. svc := &TokenRefreshService{
  614. accountRepo: repo,
  615. refreshAPI: NewOAuthRefreshAPI(repo, nil),
  616. refreshPolicy: DefaultBackgroundRefreshPolicy(),
  617. cfg: &config.TokenRefreshConfig{MaxRetries: 1},
  618. attemptTimeoutOverride: 2 * time.Second,
  619. }
  620. state := &tokenRefreshProviderState{
  621. service: svc,
  622. rateGate: newTokenRefreshRateGateWithInterval(interval),
  623. poolGate: newTokenRefreshConcurrencyGate(2),
  624. }
  625. errorsCh := make(chan error, 3)
  626. go func() {
  627. errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), firstSelection, executor, executor, time.Hour, state)
  628. }()
  629. select {
  630. case <-executor.firstStarted:
  631. case <-time.After(time.Second):
  632. require.FailNow(t, "first production-path refresh did not reach the upstream executor")
  633. }
  634. go func() {
  635. errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), contendingSelection, executor, executor, time.Hour, state)
  636. }()
  637. require.Eventually(t, func() bool {
  638. return len(state.poolGate.slots) == 2
  639. }, time.Second, time.Millisecond, "same-account contender must hold the second provider slot while waiting on the local refresh lock")
  640. go func() {
  641. errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), differentSelection, executor, executor, time.Hour, state)
  642. }()
  643. close(executor.releaseFirst)
  644. skipped := 0
  645. for i := 0; i < 3; i++ {
  646. err := <-errorsCh
  647. if errors.Is(err, errRefreshSkipped) {
  648. skipped++
  649. continue
  650. }
  651. require.NoError(t, err)
  652. }
  653. require.Equal(t, 1, skipped, "the same-account contender must reread the refreshed row and skip without upstream admission")
  654. starts := executor.startsSnapshot()
  655. require.Len(t, starts, 2, "only the two accounts that actually refresh may consume QPS admission")
  656. require.Equal(t, int64(71), starts[0].accountID)
  657. require.Equal(t, int64(72), starts[1].accountID)
  658. spacing := starts[1].at.Sub(starts[0].at)
  659. require.GreaterOrEqual(t, spacing, interval-30*time.Millisecond)
  660. require.Less(t, spacing, 350*time.Millisecond,
  661. "a same-account lock waiter must not consume a rate slot and push the different-account refresh to the second interval")
  662. t.Logf("actual_refresh_calls=%d actual_start_spacing=%s configured_spacing=%s", executor.calls.Load(), spacing, interval)
  663. }
  664. func TestTokenRefreshService_ProviderTripBeforeRateAdmissionSkipsWithoutAccountMutation(t *testing.T) {
  665. account := grokPoolAccount(73)
  666. stored := snapshotOAuthRefreshAccount(&account)
  667. repo := &breakerTripAccountRepo{productionPathRateRepo: &productionPathRateRepo{
  668. accounts: map[int64]*Account{account.ID: stored},
  669. }}
  670. refresher := &poolHealthRefresher{}
  671. svc := &TokenRefreshService{
  672. accountRepo: repo,
  673. refreshAPI: NewOAuthRefreshAPI(repo, nil),
  674. refreshPolicy: DefaultBackgroundRefreshPolicy(),
  675. cfg: &config.TokenRefreshConfig{MaxRetries: 1},
  676. }
  677. state := &tokenRefreshProviderState{
  678. service: svc,
  679. rateGate: newTokenRefreshRateGate(1),
  680. poolGate: newTokenRefreshConcurrencyGate(1),
  681. }
  682. gate := &tripBeforeRateAdmissionGate{state: state}
  683. err := svc.refreshWithRetryWithRateGate(context.Background(), &account, refresher, refresher, time.Hour, gate)
  684. require.ErrorIs(t, err, errRefreshSkipped)
  685. require.Zero(t, refresher.calls.Load(), "a tripped provider must not reach upstream rate admission")
  686. require.Zero(t, repo.setErrorCalls.Load())
  687. require.Zero(t, repo.setTempCalls.Load(), "provider skip must never fall through to per-account cooldown")
  688. }
  689. func TestTokenRefreshService_ConfigBounds(t *testing.T) {
  690. maxInt := int(^uint(0) >> 1)
  691. svc := &TokenRefreshService{cfg: &config.TokenRefreshConfig{
  692. MaxRetries: maxInt,
  693. RetryBackoffSeconds: maxInt,
  694. ProviderFailureThreshold: maxInt,
  695. AttemptTimeoutSeconds: maxInt,
  696. CycleTimeoutSeconds: maxInt,
  697. }}
  698. require.Equal(t, maxTokenRefreshMaxRetries, svc.maxRetries())
  699. require.Equal(t, maxTokenRefreshProviderFailureThreshold, svc.providerFailureThreshold())
  700. require.Equal(t, maxTokenRefreshAttemptTimeout, svc.attemptTimeout())
  701. require.Equal(t, maxTokenRefreshCycleTimeout, svc.cycleTimeout())
  702. require.LessOrEqual(t, svc.retryBackoff(1, maxTokenRefreshMaxRetries), maxTokenRefreshRetryBackoff)
  703. require.Equal(t, maxGrokOAuthReconcilePageSize, svc.grokOAuthReconcileMaxPageSize())
  704. }
  705. func TestTokenRefreshService_AttemptTimeoutStaysInsideDistributedLockLease(t *testing.T) {
  706. cache := &poolHealthTokenCacheStub{}
  707. svc := &TokenRefreshService{
  708. cfg: &config.TokenRefreshConfig{AttemptTimeoutSeconds: int(maxTokenRefreshAttemptTimeout / time.Second)},
  709. refreshAPI: NewOAuthRefreshAPI(&poolHealthAccountRepo{}, cache),
  710. }
  711. require.Equal(t, 55*time.Second, svc.attemptTimeout())
  712. require.Less(t, svc.attemptTimeout(), defaultRefreshLockTTL)
  713. }
  714. func TestTokenRefreshService_SharedProviderFailureContainsCycleWithoutAccountMutation(t *testing.T) {
  715. accounts := make([]Account, 0, 5)
  716. for id := int64(1); id <= 5; id++ {
  717. accounts = append(accounts, grokPoolAccount(id))
  718. }
  719. repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
  720. refresher := &poolHealthRefresher{err: errors.New("invalid_client: provider configuration rejected")}
  721. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  722. MaxRetries: 1,
  723. CandidatePageSize: 10,
  724. ProviderConcurrency: 4,
  725. ProviderQPS: 10000,
  726. ProviderFailureThreshold: 3,
  727. AttemptTimeoutSeconds: 1,
  728. CycleTimeoutSeconds: 2,
  729. })
  730. svc.processRefreshContext(context.Background())
  731. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  732. require.Equal(t, int64(1), refresher.calls.Load(), "shared provider configuration failures must open the in-cycle breaker immediately")
  733. require.Zero(t, setErrorCalls, "shared provider failures must not mass-disable accounts")
  734. require.Zero(t, setTempUnschedCalls, "shared provider failures must not mutate per-account scheduling state")
  735. }
  736. func TestTokenRefreshService_SharedDBRereadFailureContainsCycleWithoutAccountMutation(t *testing.T) {
  737. accounts := []Account{grokPoolAccount(1), grokPoolAccount(2), grokPoolAccount(3)}
  738. repo := &poolHealthAccountRepo{
  739. pages: map[int64][]Account{0: accounts},
  740. getByIDErr: errors.New("database unavailable"),
  741. }
  742. refresher := &poolHealthRefresher{}
  743. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  744. MaxRetries: 3,
  745. CandidatePageSize: 10,
  746. ProviderConcurrency: 4,
  747. ProviderQPS: 10000,
  748. AttemptTimeoutSeconds: 1,
  749. CycleTimeoutSeconds: 2,
  750. })
  751. svc.refreshAPI = NewOAuthRefreshAPI(repo, nil)
  752. svc.processRefreshContext(context.Background())
  753. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  754. require.Zero(t, refresher.calls.Load(), "refresh must fail closed before using stale account credentials")
  755. require.Zero(t, setErrorCalls)
  756. require.Zero(t, setTempUnschedCalls, "a shared DB outage must not mutate the selected account")
  757. }
  758. func TestTokenRefreshService_GenericGrokForbiddenContainsCycleWithoutAccountMutation(t *testing.T) {
  759. accounts := []Account{grokPoolAccount(1), grokPoolAccount(2), grokPoolAccount(3)}
  760. repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
  761. refresher := &poolHealthRefresher{err: errors.New(`GROK_OAUTH_ENTITLEMENT_DENIED: token refresh failed: status 403, body: <html>request blocked</html>`)}
  762. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  763. MaxRetries: 1,
  764. CandidatePageSize: 10,
  765. ProviderConcurrency: 4,
  766. ProviderQPS: 10000,
  767. ProviderFailureThreshold: 3,
  768. AttemptTimeoutSeconds: 1,
  769. CycleTimeoutSeconds: 2,
  770. })
  771. svc.processRefreshContext(context.Background())
  772. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  773. require.Equal(t, int64(1), refresher.calls.Load(), "an ambiguous Grok 403 must contain the provider immediately")
  774. require.Zero(t, setErrorCalls, "a generic 403 is not evidence that an account credential is permanently invalid")
  775. require.Zero(t, setTempUnschedCalls, "provider containment must not mutate account scheduling state")
  776. }
  777. func TestTokenRefreshService_ExplicitGrokEntitlementDenialIsPermanent(t *testing.T) {
  778. repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: {grokPoolAccount(1)}}}
  779. refresher := &poolHealthRefresher{err: errors.New(`GROK_OAUTH_ENTITLEMENT_DENIED: token refresh failed: status 403, body: {"error":"subscription required"}`)}
  780. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  781. MaxRetries: 1,
  782. CandidatePageSize: 10,
  783. ProviderConcurrency: 1,
  784. ProviderQPS: 10000,
  785. AttemptTimeoutSeconds: 1,
  786. CycleTimeoutSeconds: 2,
  787. })
  788. svc.processRefreshContext(context.Background())
  789. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  790. require.Equal(t, int64(1), refresher.calls.Load())
  791. require.Equal(t, 1, setErrorCalls, "explicit entitlement evidence is an account-permanent failure")
  792. require.Zero(t, setTempUnschedCalls)
  793. }
  794. func TestTokenRefreshService_AttemptTimeoutTripsRetryableProviderThreshold(t *testing.T) {
  795. accounts := []Account{grokPoolAccount(1), grokPoolAccount(2), grokPoolAccount(3)}
  796. repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
  797. refresher := &poolHealthRefresher{delay: time.Second}
  798. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  799. MaxRetries: 1,
  800. CandidatePageSize: 10,
  801. ProviderConcurrency: 1,
  802. ProviderQPS: 10000,
  803. ProviderFailureThreshold: 2,
  804. CycleTimeoutSeconds: 2,
  805. })
  806. svc.attemptTimeoutOverride = 20 * time.Millisecond
  807. svc.processRefreshContext(context.Background())
  808. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  809. require.Equal(t, int64(2), refresher.calls.Load(), "two attempt timeouts should trip the retryable provider threshold")
  810. require.Zero(t, setErrorCalls)
  811. require.Equal(t, 2, setTempUnschedCalls, "attempt timeouts remain account-transient failures before containment opens")
  812. }
  813. func TestTokenRefreshService_ParentCancellationStopsRetryWithoutAccountMutation(t *testing.T) {
  814. repo := &poolHealthAccountRepo{}
  815. ctx, cancel := context.WithCancel(context.Background())
  816. refresher := &poolHealthRefresher{
  817. err: errors.New("temporary provider failure"),
  818. cancel: cancel,
  819. }
  820. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
  821. MaxRetries: 3,
  822. RetryBackoffSeconds: 1,
  823. AttemptTimeoutSeconds: 1,
  824. })
  825. account := grokPoolAccount(42)
  826. err := svc.refreshWithRetry(ctx, &account, refresher, nil, time.Hour)
  827. require.ErrorIs(t, err, context.Canceled)
  828. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  829. require.Zero(t, setErrorCalls)
  830. require.Zero(t, setTempUnschedCalls)
  831. }
  832. func TestTokenRefreshService_LateSuccessPastAttemptDeadlineIsRejected(t *testing.T) {
  833. repo := &poolHealthAccountRepo{}
  834. refresher := &poolHealthRefresher{
  835. delay: 30 * time.Millisecond,
  836. ignoreContext: true,
  837. }
  838. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{MaxRetries: 1})
  839. svc.attemptTimeoutOverride = 10 * time.Millisecond
  840. account := grokPoolAccount(43)
  841. err := svc.refreshWithRetry(context.Background(), &account, refresher, nil, time.Hour)
  842. var timeoutErr *refreshAttemptTimeoutError
  843. require.ErrorAs(t, err, &timeoutErr)
  844. _, updatedIDs, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  845. require.Empty(t, updatedIDs, "credentials returned after the deadline must not be persisted")
  846. require.Zero(t, setErrorCalls)
  847. require.Equal(t, 1, setTempUnschedCalls)
  848. }
  849. func TestTokenRefreshService_NonRetryableGrokFailureInvalidatesTokenCache(t *testing.T) {
  850. repo := &poolHealthAccountRepo{}
  851. invalidator := &reconcileInvalidator{}
  852. refresher := &poolHealthRefresher{err: errors.New("invalid_grant: revoked")}
  853. svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{MaxRetries: 1})
  854. svc.cacheInvalidator = invalidator
  855. account := grokPoolAccount(77)
  856. err := svc.refreshWithRetry(context.Background(), &account, refresher, nil, time.Hour)
  857. require.Error(t, err)
  858. _, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
  859. require.Equal(t, 1, setErrorCalls)
  860. require.Zero(t, setTempUnschedCalls)
  861. require.Equal(t, 1, invalidator.count())
  862. }