mirror of
https://mirror.skon.top/github.com/router-for-me/CLIProxyAPI
synced 2026-04-21 01:10:14 +08:00
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.
This commit is contained in:
@@ -90,6 +90,10 @@ max-retry-interval: 30
|
||||
# When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states).
|
||||
disable-cooling: false
|
||||
|
||||
# Core auth auto-refresh worker pool size (OAuth/file-based auth token refresh).
|
||||
# When > 0, overrides the default worker count (16).
|
||||
# auth-auto-refresh-workers: 16
|
||||
|
||||
# Quota exceeded behavior
|
||||
quota-exceeded:
|
||||
switch-project: true # Whether to automatically switch to another project when a quota is exceeded
|
||||
|
||||
@@ -68,6 +68,10 @@ type Config struct {
|
||||
// DisableCooling disables quota cooldown scheduling when true.
|
||||
DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
|
||||
|
||||
// AuthAutoRefreshWorkers overrides the size of the core auth auto-refresh worker pool.
|
||||
// When <= 0, the default worker count is used.
|
||||
AuthAutoRefreshWorkers int `yaml:"auth-auto-refresh-workers" json:"auth-auto-refresh-workers"`
|
||||
|
||||
// RequestRetry defines the retry times when the request failed.
|
||||
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
||||
// MaxRetryCredentials defines the maximum number of credentials to try for a failed request.
|
||||
|
||||
@@ -11,8 +11,9 @@ import (
|
||||
)
|
||||
|
||||
type authAutoRefreshLoop struct {
|
||||
manager *Manager
|
||||
interval time.Duration
|
||||
manager *Manager
|
||||
interval time.Duration
|
||||
concurrency int
|
||||
|
||||
mu sync.Mutex
|
||||
queue refreshMinHeap
|
||||
@@ -23,21 +24,25 @@ type authAutoRefreshLoop struct {
|
||||
jobs chan string
|
||||
}
|
||||
|
||||
func newAuthAutoRefreshLoop(manager *Manager, interval time.Duration) *authAutoRefreshLoop {
|
||||
func newAuthAutoRefreshLoop(manager *Manager, interval time.Duration, concurrency int) *authAutoRefreshLoop {
|
||||
if interval <= 0 {
|
||||
interval = refreshCheckInterval
|
||||
}
|
||||
jobBuffer := refreshMaxConcurrency * 4
|
||||
if concurrency <= 0 {
|
||||
concurrency = refreshMaxConcurrency
|
||||
}
|
||||
jobBuffer := concurrency * 4
|
||||
if jobBuffer < 64 {
|
||||
jobBuffer = 64
|
||||
}
|
||||
return &authAutoRefreshLoop{
|
||||
manager: manager,
|
||||
interval: interval,
|
||||
index: make(map[string]*refreshHeapItem),
|
||||
dirty: make(map[string]struct{}),
|
||||
wakeCh: make(chan struct{}, 1),
|
||||
jobs: make(chan string, jobBuffer),
|
||||
manager: manager,
|
||||
interval: interval,
|
||||
concurrency: concurrency,
|
||||
index: make(map[string]*refreshHeapItem),
|
||||
dirty: make(map[string]struct{}),
|
||||
wakeCh: make(chan struct{}, 1),
|
||||
jobs: make(chan string, jobBuffer),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -59,7 +64,11 @@ func (l *authAutoRefreshLoop) run(ctx context.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < refreshMaxConcurrency; i++ {
|
||||
workers := l.concurrency
|
||||
if workers <= 0 {
|
||||
workers = refreshMaxConcurrency
|
||||
}
|
||||
for i := 0; i < workers; i++ {
|
||||
go l.worker(ctx)
|
||||
}
|
||||
|
||||
|
||||
@@ -2912,7 +2912,11 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio
|
||||
}
|
||||
|
||||
ctx, cancelCtx := context.WithCancel(parent)
|
||||
loop := newAuthAutoRefreshLoop(m, interval)
|
||||
workers := refreshMaxConcurrency
|
||||
if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 {
|
||||
workers = cfg.AuthAutoRefreshWorkers
|
||||
}
|
||||
loop := newAuthAutoRefreshLoop(m, interval, workers)
|
||||
|
||||
m.mu.Lock()
|
||||
m.refreshCancel = cancelCtx
|
||||
|
||||
Reference in New Issue
Block a user