mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-30 16:23:31 +00:00
Merge pull request #7064 from xlplbo/codex/openai-ws-v2-pool-capacity
fix(openai-ws): 对齐新版 ctx_pool 连接容量计算
This commit is contained in:
@@ -1274,7 +1274,9 @@ type GatewayOpenAIWSConfig struct {
|
||||
MaxConnsPerAccount int `mapstructure:"max_conns_per_account"`
|
||||
MinIdlePerAccount int `mapstructure:"min_idle_per_account"`
|
||||
MaxIdlePerAccount int `mapstructure:"max_idle_per_account"`
|
||||
// DynamicMaxConnsByAccountConcurrencyEnabled: 是否按账号并发动态计算连接池上限
|
||||
// DynamicMaxConnsByAccountConcurrencyEnabled: 是否按账号并发动态计算连接池上限。
|
||||
// 旧版及 mode_router_v2 的 ctx_pool 共用此开关和类型系数;关闭后使用 max_conns_per_account。
|
||||
// mode_router_v2 下并发数 <= 0 的账号仍不可调度。
|
||||
DynamicMaxConnsByAccountConcurrencyEnabled bool `mapstructure:"dynamic_max_conns_by_account_concurrency_enabled"`
|
||||
// OAuthMaxConnsFactor: OAuth 账号连接池系数(effective=ceil(concurrency*factor),再受 max_conns_per_account 封顶)。
|
||||
// ctx_pool 接入下每个客户端会话在整个生命周期(含轮次之间)持有一条上游连接,此上限限制的是同时持有连接的会话数,
|
||||
|
||||
@@ -2221,14 +2221,8 @@ func (p *openAIWSConnPool) effectiveMaxConnsByAccount(account *Account) int {
|
||||
if hardCap <= 0 {
|
||||
return 0
|
||||
}
|
||||
if p.modeRouterV2Enabled() {
|
||||
if account == nil {
|
||||
return hardCap
|
||||
}
|
||||
if account.Concurrency <= 0 {
|
||||
return 0
|
||||
}
|
||||
return min(account.Concurrency, hardCap)
|
||||
if p.modeRouterV2Enabled() && account != nil && account.Concurrency <= 0 {
|
||||
return 0
|
||||
}
|
||||
if account == nil || !p.dynamicMaxConnsEnabled() {
|
||||
return hardCap
|
||||
|
||||
@@ -1338,21 +1338,95 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T)
|
||||
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(account), "关闭动态模式后应保持旧行为")
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2RespectsHardCap(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
||||
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 0.3
|
||||
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 0.6
|
||||
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
accountType string
|
||||
concurrency int
|
||||
factor float64
|
||||
dynamic bool
|
||||
want int
|
||||
}{
|
||||
{name: "oauth expansion", accountType: AccountTypeOAuth, concurrency: 1, factor: 5, dynamic: true, want: 5},
|
||||
{name: "apikey expansion", accountType: AccountTypeAPIKey, concurrency: 1, factor: 5, dynamic: true, want: 5},
|
||||
{name: "oauth fraction", accountType: AccountTypeOAuth, concurrency: 20, factor: 0.3, dynamic: true, want: 6},
|
||||
{name: "apikey rounding", accountType: AccountTypeAPIKey, concurrency: 3, factor: 0.6, dynamic: true, want: 2},
|
||||
{name: "minimum one", accountType: AccountTypeOAuth, concurrency: 1, factor: 0.3, dynamic: true, want: 1},
|
||||
{name: "hard cap", accountType: AccountTypeOAuth, concurrency: 3, factor: 5, dynamic: true, want: 8},
|
||||
{name: "dynamic disabled", accountType: AccountTypeAPIKey, concurrency: 2, factor: 0.6, want: 8},
|
||||
{name: "zero concurrency", accountType: AccountTypeOAuth, factor: 5, dynamic: true, want: 0},
|
||||
{name: "negative concurrency", accountType: AccountTypeAPIKey, concurrency: -1, factor: 5, dynamic: true, want: 0},
|
||||
{name: "zero concurrency dynamic disabled", accountType: AccountTypeOAuth, factor: 5, want: 0},
|
||||
{name: "negative concurrency dynamic disabled", accountType: AccountTypeAPIKey, concurrency: -1, factor: 5, want: 0},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
||||
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = tt.dynamic
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1
|
||||
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 1
|
||||
if tt.accountType == AccountTypeOAuth {
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = tt.factor
|
||||
} else {
|
||||
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = tt.factor
|
||||
}
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
t.Cleanup(pool.Close)
|
||||
account := &Account{Platform: PlatformOpenAI, Type: tt.accountType, Concurrency: tt.concurrency}
|
||||
require.Equal(t, tt.want, pool.effectiveMaxConnsByAccount(account))
|
||||
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(nil))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
func TestOpenAIWSConnPool_AcquireRetainedSessionsUsesScaledCapacity(t *testing.T) {
|
||||
for _, mode := range []struct {
|
||||
name string
|
||||
v2 bool
|
||||
}{
|
||||
{name: "legacy"},
|
||||
{name: "mode router v2", v2: true},
|
||||
} {
|
||||
t.Run(mode.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = mode.v2
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 3
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 3
|
||||
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 1
|
||||
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 2
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
pool.setClientDialerForTest(&openAIWSFakeDialer{})
|
||||
t.Cleanup(pool.Close)
|
||||
req := openAIWSAcquireRequest{
|
||||
Account: &Account{ID: 902, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 2},
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
leases := make([]*openAIWSConnLease, 0, 3)
|
||||
for range 3 {
|
||||
lease, err := pool.Acquire(ctx, req)
|
||||
require.NoError(t, err, "轮次之间仍持有连接的会话应可使用系数扩出的容量")
|
||||
t.Cleanup(lease.Release)
|
||||
leases = append(leases, lease)
|
||||
}
|
||||
|
||||
high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20}
|
||||
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束")
|
||||
fullCtx, fullCancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer fullCancel()
|
||||
_, err := pool.Acquire(fullCtx, req)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded, "扩容后仍必须受全局硬上限约束")
|
||||
|
||||
nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0}
|
||||
require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度")
|
||||
leases[0].Release()
|
||||
reused, err := pool.Acquire(ctx, req)
|
||||
require.NoError(t, err)
|
||||
defer reused.Release()
|
||||
require.Equal(t, leases[0].ConnID(), reused.ConnID())
|
||||
require.True(t, reused.Reused())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_AcquireRejectsWhenEffectiveMaxConnsIsZero(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user