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

config.example.yaml +4 −0
internal/config/config.go +4 −0
sdk/cliproxy/auth/auto_refresh_loop.go +20 −11
sdk/cliproxy/auth/conductor.go +5 −1

config.example.yaml+4−0

@@ -90,6 +90,10 @@ max-retry-interval: 30 · 89 unmodified lines
90	90	# When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states).
91	91	disable-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	97	# Quota exceeded behavior
94	98	quota-exceeded:
95	99	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	68	// DisableCooling disables quota cooldown scheduling when true.
69	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	75	// RequestRetry defines the retry times when the request failed.
72	76	RequestRetry int `yaml:"request-retry" json:"request-retry"`
73	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
13	13	type authAutoRefreshLoop struct {
14		− manager *Manager
15		− interval time.Duration
	14	+ manager *Manager
	15	+ interval time.Duration
	16	+ concurrency int
16	17
17	18	mu sync.Mutex
18	19	queue refreshMinHeap
@@ -23,21 +24,25 @@ type authAutoRefreshLoop struct { · 4 unmodified lines
23	24	jobs chan string
24	25	}
25	26
26		−func newAuthAutoRefreshLoop(manager Manager, interval time.Duration) authAutoRefreshLoop {
	27	+func newAuthAutoRefreshLoop(manager Manager, interval time.Duration, concurrency int) authAutoRefreshLoop {
27	28	if interval <= 0 {
28	29	interval = refreshCheckInterval
29	30	}
30		− jobBuffer := refreshMaxConcurrency * 4
	31	+ if concurrency <= 0 {
	32	+ concurrency = refreshMaxConcurrency
	33	+ }
	34	+ jobBuffer := concurrency * 4
31	35	if jobBuffer < 64 {
32	36	jobBuffer = 64
33	37	}
34	38	return &authAutoRefreshLoop{
35		− manager: manager,
36		− interval: interval,
37		− index: make(map[string]*refreshHeapItem),
38		− dirty: make(map[string]struct{}),
39		− wakeCh: make(chan struct{}, 1),
40		− jobs: make(chan string, jobBuffer),
	39	+ manager: manager,
	40	+ interval: interval,
	41	+ concurrency: concurrency,
	42	+ index: make(map[string]*refreshHeapItem),
	43	+ dirty: make(map[string]struct{}),
	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	64	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	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	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	2921	m.mu.Lock()
2918	2922	m.refreshCancel = cancelCtx