diff --git a/controller/channel-test.go b/controller/channel-test.go index 4ba3698bd54c..3a5dd8f13123 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -12,6 +12,7 @@ import ( "net/http/httptest" "strconv" "strings" + "sync" "time" "github.com/QuantumNous/new-api/common" @@ -908,31 +909,56 @@ func performChannelTests(ctx context.Context, channels []*model.Channel, testUse summary := channelTestSummary{} var disableThreshold = int64(common.ChannelDisableThreshold * 1000) if disableThreshold == 0 { - disableThreshold = 10000000 // a impossible value + disableThreshold = 10000000 // an impossible value } total := len(channels) - for index, channel := range channels { - if ctx != nil && ctx.Err() != nil { - break - } - if report != nil { - report(index, total) // channels completed before this one - } - if channel.Status == common.ChannelStatusManuallyDisabled { - continue + + concurrency := operation_setting.GetMonitorSetting().ChannelTestConcurrency + if concurrency < 1 { + concurrency = 1 + } + if concurrency > total && total > 0 { + concurrency = total + } + + var ( + mu sync.Mutex // guards summary + progressMu sync.Mutex // serializes report so progress stays monotonic + completed int + wg sync.WaitGroup + sem = make(chan struct{}, concurrency) + ) + + // advanceProgress reports one more processed channel. report may persist + // progress to storage and is not assumed to be concurrency-safe, so calls are + // serialized here to keep the reported count ordered and monotonic. + advanceProgress := func() { + if report == nil { + return } + progressMu.Lock() + completed++ + report(completed, total) + progressMu.Unlock() + } + + if report != nil { + report(0, total) + } + + // testOne runs a single channel test and folds its outcome into summary. It is + // safe to call from multiple goroutines: every testChannel call builds its own + // gin context and recorder, and the shared summary is mutex guarded. + testOne := func(channel *model.Channel) { isChannelEnabled := channel.Status == common.ChannelStatusEnabled tik := time.Now() result := testChannel(ctx, channel, testUserID, "", "", shouldUseStreamForAutomaticChannelTest(channel)) - tok := time.Now() - milliseconds := tok.Sub(tik).Milliseconds() + milliseconds := time.Since(tik).Milliseconds() if ctx != nil && ctx.Err() != nil { - break + return } - summary.Tested++ - shouldBanChannel := false newAPIError := result.newAPIError // request error disables the channel @@ -941,45 +967,91 @@ func performChannelTests(ctx context.Context, channels []*model.Channel, testUse } // 当错误检查通过，才检查响应时间 - if common.AutomaticDisableChannelEnabled && !shouldBanChannel { - if milliseconds > disableThreshold { - err := fmt.Errorf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0) - newAPIError = types.NewOpenAIError(err, types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout) - shouldBanChannel = true - } + if common.AutomaticDisableChannelEnabled && !shouldBanChannel && milliseconds > disableThreshold { + err := fmt.Errorf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0) + newAPIError = types.NewOpenAIError(err, types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout) + shouldBanChannel = true + } + + banned := allowDisable && isChannelEnabled && shouldBanChannel && channel.GetAutoBan() + if banned { + processChannelError(result.context, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) + } + enabled := result.localErr == nil && !isChannelEnabled && service.ShouldEnableChannel(newAPIError, channel.Status) + if enabled { + service.EnableChannel(channel.Id, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.Name) } + channel.UpdateResponseTime(milliseconds) + mu.Lock() + summary.Tested++ if newAPIError == nil { summary.Succeeded++ } else { summary.Failed++ } - - // disable channel - if allowDisable && isChannelEnabled && shouldBanChannel && channel.GetAutoBan() { - processChannelError(result.context, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) + if banned { summary.Disabled++ } - - // enable channel - if result.localErr == nil && !isChannelEnabled && service.ShouldEnableChannel(newAPIError, channel.Status) { - service.EnableChannel(channel.Id, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.Name) + if enabled { summary.Enabled++ } + mu.Unlock() - channel.UpdateResponseTime(milliseconds) - if common.RequestInterval > 0 { + advanceProgress() + } + + dispatched := false +dispatch: + for _, channel := range channels { + if ctx != nil && ctx.Err() != nil { + break + } + if channel.Status == common.ChannelStatusManuallyDisabled { + advanceProgress() + continue + } + + // RequestInterval throttles how fast new tests are dispatched, spacing out + // upstream load even when tests run concurrently. Only throttle between + // dispatches, never before the first or after the last one. + if dispatched && common.RequestInterval > 0 { if ctx == nil { time.Sleep(common.RequestInterval) } else { + timer := time.NewTimer(common.RequestInterval) select { case <-ctx.Done(): - return summary - case <-time.After(common.RequestInterval): + timer.Stop() + break dispatch + case <-timer.C: } } } + + // Acquire a worker slot, honoring cancellation so a runner that loses its + // lease stops promptly. + if ctx == nil { + sem <- struct{}{} + } else { + select { + case <-ctx.Done(): + break dispatch + case sem <- struct{}{}: + } + } + + wg.Add(1) + go func(ch *model.Channel) { + defer wg.Done() + defer func() { <-sem }() + testOne(ch) + }(channel) + dispatched = true } + + wg.Wait() + if report != nil && (ctx == nil || ctx.Err() == nil) { report(total, total) // mark complete only when the full set was tested } diff --git a/setting/operation_setting/monitor_setting.go b/setting/operation_setting/monitor_setting.go index 8593d8349a4c..2a3f80b0238b 100644 --- a/setting/operation_setting/monitor_setting.go +++ b/setting/operation_setting/monitor_setting.go @@ -11,6 +11,9 @@ type MonitorSetting struct { AutoTestChannelEnabled bool `json:"auto_test_channel_enabled"` AutoTestChannelMinutes float64 `json:"auto_test_channel_minutes"` ChannelTestMode string `json:"channel_test_mode"` + // ChannelTestConcurrency is the number of channels tested in parallel during + // a batch test run. 1 preserves the original fully sequential behavior. + ChannelTestConcurrency int `json:"test_concurrency"` } const ( @@ -23,6 +26,7 @@ var monitorSetting = MonitorSetting{ AutoTestChannelEnabled: false, AutoTestChannelMinutes: 10, ChannelTestMode: ChannelTestModeScheduledAll, + ChannelTestConcurrency: 1, } func init() { @@ -48,5 +52,14 @@ func GetMonitorSetting() *MonitorSetting { if monitorSetting.ChannelTestMode != ChannelTestModePassiveRecovery { monitorSetting.ChannelTestMode = ChannelTestModeScheduledAll } + if v, ok := os.LookupEnv("CHANNEL_TEST_CONCURRENCY"); ok { + if parsed, err := strconv.Atoi(v); err == nil { + // A non-positive override is normalized to 1 by the clamp below. + monitorSetting.ChannelTestConcurrency = parsed + } + } + if monitorSetting.ChannelTestConcurrency < 1 { + monitorSetting.ChannelTestConcurrency = 1 + } return &monitorSetting } diff --git a/setting/operation_setting/monitor_setting_test.go b/setting/operation_setting/monitor_setting_test.go index 7aef7eaaa9cc..7aaee71019db 100644 --- a/setting/operation_setting/monitor_setting_test.go +++ b/setting/operation_setting/monitor_setting_test.go @@ -41,3 +41,50 @@ func TestGetMonitorSetting_ChannelTestEnabledEnvCanEnableDisabledConfig(t *testi assert.True(t, setting.AutoTestChannelEnabled) assert.Equal(t, float64(12), setting.AutoTestChannelMinutes) } + +func TestGetMonitorSetting_ChannelTestConcurrencyNormalizedToAtLeastOne(t *testing.T) { + orig := monitorSetting + t.Cleanup(func() { monitorSetting = orig }) + + // Isolate from any ambient env override so the stored value is what is tested. + // An empty value makes GetMonitorSetting skip the env branch and fall through + // to the clamp. + t.Setenv("CHANNEL_TEST_CONCURRENCY", "") + + // A non-positive stored value (e.g. an unset legacy option) must normalize to 1 + // so the batch test never runs with a zero-sized worker pool. + monitorSetting = MonitorSetting{ChannelTestConcurrency: 0} + + setting := GetMonitorSetting() + + require.NotNil(t, setting) + assert.Equal(t, 1, setting.ChannelTestConcurrency) +} + +func TestGetMonitorSetting_ChannelTestConcurrencyEnvOverride(t *testing.T) { + orig := monitorSetting + t.Cleanup(func() { monitorSetting = orig }) + + t.Setenv("CHANNEL_TEST_CONCURRENCY", "8") + monitorSetting = MonitorSetting{ChannelTestConcurrency: 1} + + setting := GetMonitorSetting() + + require.NotNil(t, setting) + assert.Equal(t, 8, setting.ChannelTestConcurrency) +} + +func TestGetMonitorSetting_ChannelTestConcurrencyEnvNonPositiveNormalizedToOne(t *testing.T) { + orig := monitorSetting + t.Cleanup(func() { monitorSetting = orig }) + + // An explicit non-positive env override must be honored (over the stored + // value) and then normalized to 1 rather than silently ignored. + t.Setenv("CHANNEL_TEST_CONCURRENCY", "0") + monitorSetting = MonitorSetting{ChannelTestConcurrency: 5} + + setting := GetMonitorSetting() + + require.NotNil(t, setting) + assert.Equal(t, 1, setting.ChannelTestConcurrency) +} diff --git a/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx b/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx index db8330ef20b4..426154b75795 100644 --- a/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx +++ b/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx @@ -212,6 +212,7 @@ export function ModelMutateDrawer({ '100-199,300-399,401-407,409-499,500-503,505-523,525-599', 'monitor_setting.auto_test_channel_enabled': false, 'monitor_setting.auto_test_channel_minutes': 10, + 'monitor_setting.test_concurrency': 1, 'monitor_setting.channel_test_mode': 'scheduled_all', 'channel_affinity_setting.enabled': false, 'channel_affinity_setting.switch_on_success': true, diff --git a/web/default/src/features/system-settings/models/index.tsx b/web/default/src/features/system-settings/models/index.tsx index e7a1ee1c490c..574804c0b4f4 100644 --- a/web/default/src/features/system-settings/models/index.tsx +++ b/web/default/src/features/system-settings/models/index.tsx @@ -72,6 +72,7 @@ const defaultModelSettings: ModelSettings = { '100-199,300-399,401-407,409-499,500-503,505-523,525-599', 'monitor_setting.auto_test_channel_enabled': false, 'monitor_setting.auto_test_channel_minutes': 10, + 'monitor_setting.test_concurrency': 1, 'monitor_setting.channel_test_mode': 'scheduled_all', 'channel_affinity_setting.enabled': false, 'channel_affinity_setting.switch_on_success': true, diff --git a/web/default/src/features/system-settings/models/routing-reliability-section.tsx b/web/default/src/features/system-settings/models/routing-reliability-section.tsx index efc8092a1c43..6678f882bd7a 100644 --- a/web/default/src/features/system-settings/models/routing-reliability-section.tsx +++ b/web/default/src/features/system-settings/models/routing-reliability-section.tsx @@ -81,6 +81,10 @@ const routingReliabilitySchema = z .number() .int() .min(1, 'Interval must be at least 1 minute'), + test_concurrency: z.coerce + .number() + .int() + .min(1, 'Concurrency must be at least 1'), channel_test_mode: z.enum(channelTestModes), }), }) @@ -126,6 +130,7 @@ type RoutingReliabilitySectionProps = { AutomaticRetryStatusCodes: string 'monitor_setting.auto_test_channel_enabled': boolean 'monitor_setting.auto_test_channel_minutes': number + 'monitor_setting.test_concurrency': number 'monitor_setting.channel_test_mode': ChannelTestMode } } @@ -144,6 +149,7 @@ type NormalizedRoutingReliabilityValues = { AutomaticRetryStatusCodes: string 'monitor_setting.auto_test_channel_enabled': boolean 'monitor_setting.auto_test_channel_minutes': number + 'monitor_setting.test_concurrency': number 'monitor_setting.channel_test_mode': ChannelTestMode } @@ -168,6 +174,7 @@ const buildFormDefaults = ( defaults['monitor_setting.auto_test_channel_enabled'], auto_test_channel_minutes: defaults['monitor_setting.auto_test_channel_minutes'], + test_concurrency: defaults['monitor_setting.test_concurrency'], channel_test_mode: normalizeChannelTestMode( defaults['monitor_setting.channel_test_mode'] ), @@ -194,6 +201,8 @@ const normalizeDefaults = ( defaults['monitor_setting.auto_test_channel_enabled'], 'monitor_setting.auto_test_channel_minutes': defaults['monitor_setting.auto_test_channel_minutes'], + 'monitor_setting.test_concurrency': + defaults['monitor_setting.test_concurrency'], 'monitor_setting.channel_test_mode': normalizeChannelTestMode( defaults['monitor_setting.channel_test_mode'] ), @@ -219,6 +228,7 @@ const normalizeFormValues = ( values.monitor_setting.auto_test_channel_enabled, 'monitor_setting.auto_test_channel_minutes': values.monitor_setting.auto_test_channel_minutes, + 'monitor_setting.test_concurrency': values.monitor_setting.test_concurrency, 'monitor_setting.channel_test_mode': values.monitor_setting.channel_test_mode, }) @@ -453,6 +463,30 @@ export function RoutingReliabilitySection({ )} /> + ( + + {t('Test concurrency')} + + + + + {t( + 'Number of channels tested in parallel during a batch test (1 = one at a time).' + )} + + + + )} + /> +