History

feat(auth): add configurable worker pool size for auto-refresh loop

- Introduced `auth-auto-refresh-workers` config option to override default concurrency. - Updated `authAutoRefreshLoop` to support customizable worker counts. - Enhanced token refresh scheduling flexibility by aligning worker pool with runtime configurations.

Luis Pater committed 7 months ago

Commit 5bfaf801 parent 6c0a1ef4 files · +33 −12Browse files

Unified · Split

4 files
  1. config.example.yaml +4 −0
  2. internal/config/config.go +4 −0
  3. sdk/cliproxy/auth/auto_refresh_loop.go +20 −11
  4. sdk/cliproxy/auth/conductor.go +5 −1
config.example.yaml+4−0
@@ -90,6 +90,10 @@ max-retry-interval: 30 · 89 unmodified lines
90# When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states).90# When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states).
91disable-cooling: false91disable-cooling: false
92 92
93# Core auth auto-refresh worker pool size (OAuth/file-based auth token refresh).
94# When > 0, overrides the default worker count (16).
95# auth-auto-refresh-workers: 16
96
93# Quota exceeded behavior97# Quota exceeded behavior
94quota-exceeded:98quota-exceeded:
95 switch-project: true # Whether to automatically switch to another project when a quota is exceeded99 switch-project: true # Whether to automatically switch to another project when a quota is exceeded
internal/config/config.go+4−0
@@ -68,6 +68,10 @@ type Config struct { · 67 unmodified lines
68 // DisableCooling disables quota cooldown scheduling when true.68 // DisableCooling disables quota cooldown scheduling when true.
69 DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`69 DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
70 70
71 // AuthAutoRefreshWorkers overrides the size of the core auth auto-refresh worker pool.
72 // When <= 0, the default worker count is used.
73 AuthAutoRefreshWorkers int `yaml:"auth-auto-refresh-workers" json:"auth-auto-refresh-workers"`
74
71 // RequestRetry defines the retry times when the request failed.75 // RequestRetry defines the retry times when the request failed.
72 RequestRetry int `yaml:"request-retry" json:"request-retry"`76 RequestRetry int `yaml:"request-retry" json:"request-retry"`
73 // MaxRetryCredentials defines the maximum number of credentials to try for a failed request.77 // MaxRetryCredentials defines the maximum number of credentials to try for a failed request.
sdk/cliproxy/auth/auto_refresh_loop.go+20−11
@@ -11,8 +11,9 @@ import ( · 10 unmodified lines
11)11)
12 12
13type authAutoRefreshLoop struct {13type authAutoRefreshLoop struct {
14 manager *Manager14 manager *Manager
15 interval time.Duration15 interval time.Duration
16 concurrency int
16 17
17 mu sync.Mutex18 mu sync.Mutex
18 queue refreshMinHeap19 queue refreshMinHeap
@@ -23,21 +24,25 @@ type authAutoRefreshLoop struct { · 4 unmodified lines
23 jobs chan string24 jobs chan string
24}25}
25 26
26func newAuthAutoRefreshLoop(manager *Manager, interval time.Duration) *authAutoRefreshLoop {27func newAuthAutoRefreshLoop(manager *Manager, interval time.Duration, concurrency int) *authAutoRefreshLoop {
27 if interval <= 0 {28 if interval <= 0 {
28 interval = refreshCheckInterval29 interval = refreshCheckInterval
29 }30 }
30 jobBuffer := refreshMaxConcurrency * 431 if concurrency <= 0 {
32 concurrency = refreshMaxConcurrency
33 }
34 jobBuffer := concurrency * 4
31 if jobBuffer < 64 {35 if jobBuffer < 64 {
32 jobBuffer = 6436 jobBuffer = 64
33 }37 }
34 return &authAutoRefreshLoop{38 return &authAutoRefreshLoop{
35 manager: manager,39 manager: manager,
36 interval: interval,40 interval: interval,
37 index: make(map[string]*refreshHeapItem),41 concurrency: concurrency,
38 dirty: make(map[string]struct{}),42 index: make(map[string]*refreshHeapItem),
39 wakeCh: make(chan struct{}, 1),43 dirty: make(map[string]struct{}),
40 jobs: make(chan string, jobBuffer),44 wakeCh: make(chan struct{}, 1),
45 jobs: make(chan string, jobBuffer),
41 }46 }
42}47}
43 48
@@ -59,7 +64,11 @@ func (l *authAutoRefreshLoop) run(ctx context.Context) { · 15 unmodified lines
59 return64 return
60 }65 }
61 66
62 for i := 0; i < refreshMaxConcurrency; i++ {67 workers := l.concurrency
68 if workers <= 0 {
69 workers = refreshMaxConcurrency
70 }
71 for i := 0; i < workers; i++ {
63 go l.worker(ctx)72 go l.worker(ctx)
64 }73 }
65 74
sdk/cliproxy/auth/conductor.go+5−1
@@ -2912,7 +2912,11 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio · 2911 unmodified lines
2912 }2912 }
2913 2913
2914 ctx, cancelCtx := context.WithCancel(parent)2914 ctx, cancelCtx := context.WithCancel(parent)
2915 loop := newAuthAutoRefreshLoop(m, interval)2915 workers := refreshMaxConcurrency
2916 if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 {
2917 workers = cfg.AuthAutoRefreshWorkers
2918 }
2919 loop := newAuthAutoRefreshLoop(m, interval, workers)
2916 2920
2917 m.mu.Lock()2921 m.mu.Lock()
2918 m.refreshCancel = cancelCtx2922 m.refreshCancel = cancelCtx