李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
2371 lines
76 KiB
Go
2371 lines
76 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestOpenAIWSConnPool_CleanupStaleAndTrimIdle(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
accountID := int64(10)
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
|
|
stale := newOpenAIWSConn("stale", accountID, nil, nil)
|
|
stale.createdAtNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
|
|
stale.lastUsedNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
|
|
|
|
idleOld := newOpenAIWSConn("idle_old", accountID, nil, nil)
|
|
idleOld.lastUsedNano.Store(time.Now().Add(-10 * time.Minute).UnixNano())
|
|
|
|
idleNew := newOpenAIWSConn("idle_new", accountID, nil, nil)
|
|
idleNew.lastUsedNano.Store(time.Now().Add(-1 * time.Minute).UnixNano())
|
|
|
|
ap.conns[stale.id] = stale
|
|
ap.conns[idleOld.id] = idleOld
|
|
ap.conns[idleNew.id] = idleNew
|
|
|
|
evicted := pool.cleanupAccountLocked(ap, time.Now(), pool.maxConnsHardCap())
|
|
closeOpenAIWSConns(evicted)
|
|
|
|
require.Nil(t, ap.conns["stale"], "stale connection should be rotated")
|
|
require.Nil(t, ap.conns["idle_old"], "old idle should be trimmed by max_idle")
|
|
require.NotNil(t, ap.conns["idle_new"], "newer idle should be kept")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_NextConnIDFormat(t *testing.T) {
|
|
pool := newOpenAIWSConnPool(&config.Config{})
|
|
id1 := pool.nextConnID(42)
|
|
id2 := pool.nextConnID(42)
|
|
|
|
require.True(t, strings.HasPrefix(id1, "oa_ws_42_"))
|
|
require.True(t, strings.HasPrefix(id2, "oa_ws_42_"))
|
|
require.NotEqual(t, id1, id2)
|
|
require.Equal(t, "oa_ws_42_1", id1)
|
|
require.Equal(t, "oa_ws_42_2", id2)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireCleanupInterval(t *testing.T) {
|
|
require.Equal(t, 3*time.Second, openAIWSAcquireCleanupInterval)
|
|
require.Less(t, openAIWSAcquireCleanupInterval, openAIWSBackgroundSweepTicker)
|
|
}
|
|
|
|
func TestNormalizeOpenAIWSRoutingAffinityPrefersCanonicalAndSortsVariants(t *testing.T) {
|
|
headers := http.Header{
|
|
"X-CODEX-ROUTING-HINT": []string{" variant-uppercase "},
|
|
"X-Codex-Routing-Hint": []string{" ", " canonical "},
|
|
}
|
|
|
|
require.Equal(t, "canonical", normalizeOpenAIWSRoutingAffinity(headers))
|
|
|
|
delete(headers, "X-Codex-Routing-Hint")
|
|
require.Equal(t, "variant-uppercase", normalizeOpenAIWSRoutingAffinity(headers))
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_WriteJSONAndGuards(t *testing.T) {
|
|
conn := newOpenAIWSConn("lease_write", 1, &openAIWSFakeConn{}, nil)
|
|
lease := &openAIWSConnLease{conn: conn}
|
|
require.NoError(t, lease.WriteJSON(map[string]any{"type": "response.create"}, 0))
|
|
|
|
var nilLease *openAIWSConnLease
|
|
err := nilLease.WriteJSONWithContextTimeout(context.Background(), map[string]any{"type": "response.create"}, time.Second)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
|
|
err = (&openAIWSConnLease{}).WriteJSONWithContextTimeout(context.Background(), map[string]any{"type": "response.create"}, time.Second)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
}
|
|
|
|
func TestOpenAIWSConn_WriteJSONWithTimeout_NilParentContextUsesBackground(t *testing.T) {
|
|
probe := &openAIWSContextProbeConn{}
|
|
conn := newOpenAIWSConn("ctx_probe", 1, probe, nil)
|
|
require.NoError(t, conn.writeJSONWithTimeout(context.Background(), map[string]any{"type": "response.create"}, 0))
|
|
require.NotNil(t, probe.lastWriteCtx)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_TargetConnCountAdaptive(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 6
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
|
|
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.5
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
ap := pool.getOrCreateAccountPool(88)
|
|
|
|
conn1 := newOpenAIWSConn("c1", 88, nil, nil)
|
|
conn2 := newOpenAIWSConn("c2", 88, nil, nil)
|
|
require.True(t, conn1.tryAcquire())
|
|
require.True(t, conn2.tryAcquire())
|
|
conn1.waiters.Store(1)
|
|
conn2.waiters.Store(1)
|
|
|
|
ap.conns[conn1.id] = conn1
|
|
ap.conns[conn2.id] = conn2
|
|
|
|
target := pool.targetConnCountLocked(ap, pool.maxConnsHardCap())
|
|
require.Equal(t, 6, target, "应按 inflight+waiters 与 target_utilization 自适应扩容到上限")
|
|
|
|
conn1.release()
|
|
conn2.release()
|
|
conn1.waiters.Store(0)
|
|
conn2.waiters.Store(0)
|
|
target = pool.targetConnCountLocked(ap, pool.maxConnsHardCap())
|
|
require.Equal(t, 1, target, "低负载时应缩回到最小空闲连接")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_TargetConnCountMinIdleZero(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
ap := pool.getOrCreateAccountPool(66)
|
|
|
|
target := pool.targetConnCountLocked(ap, pool.maxConnsHardCap())
|
|
require.Equal(t, 0, target, "min_idle=0 且无负载时应允许缩容到 0")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_EnsureTargetIdleAsync(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 2
|
|
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
|
|
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
pool.setClientDialerForTest(&openAIWSFakeDialer{})
|
|
|
|
accountID := int64(77)
|
|
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
ap.mu.Lock()
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
}
|
|
ap.mu.Unlock()
|
|
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
|
|
require.Eventually(t, func() bool {
|
|
ap, ok := pool.getAccountPool(accountID)
|
|
if !ok || ap == nil {
|
|
return false
|
|
}
|
|
ap.mu.Lock()
|
|
defer ap.mu.Unlock()
|
|
return len(ap.conns) >= 2
|
|
}, 2*time.Second, 20*time.Millisecond)
|
|
|
|
metrics := pool.SnapshotMetrics()
|
|
require.GreaterOrEqual(t, metrics.ScaleUpTotal, int64(2))
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_EnsureTargetIdleAsyncCooldown(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 2
|
|
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
|
|
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
|
|
cfg.Gateway.OpenAIWS.PrewarmCooldownMS = 500
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
|
|
accountID := int64(178)
|
|
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
ap.mu.Lock()
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
}
|
|
ap.mu.Unlock()
|
|
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
require.Eventually(t, func() bool {
|
|
ap, ok := pool.getAccountPool(accountID)
|
|
if !ok || ap == nil {
|
|
return false
|
|
}
|
|
ap.mu.Lock()
|
|
defer ap.mu.Unlock()
|
|
return len(ap.conns) >= 2 && !ap.prewarmActive
|
|
}, 2*time.Second, 20*time.Millisecond)
|
|
firstDialCount := dialer.DialCount()
|
|
require.GreaterOrEqual(t, firstDialCount, 2)
|
|
|
|
// 人工制造缺口触发新一轮预热需求。
|
|
ap, ok := pool.getAccountPool(accountID)
|
|
require.True(t, ok)
|
|
require.NotNil(t, ap)
|
|
ap.mu.Lock()
|
|
for id := range ap.conns {
|
|
delete(ap.conns, id)
|
|
break
|
|
}
|
|
ap.mu.Unlock()
|
|
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
time.Sleep(120 * time.Millisecond)
|
|
require.Equal(t, firstDialCount, dialer.DialCount(), "cooldown 窗口内不应再次触发预热")
|
|
|
|
time.Sleep(450 * time.Millisecond)
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
require.Eventually(t, func() bool {
|
|
return dialer.DialCount() > firstDialCount
|
|
}, 2*time.Second, 20*time.Millisecond)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_EnsureTargetIdleAsyncFailureSuppress(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
|
|
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
|
|
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
|
|
cfg.Gateway.OpenAIWS.PrewarmCooldownMS = 0
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSAlwaysFailDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
|
|
accountID := int64(279)
|
|
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
ap.mu.Lock()
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
}
|
|
ap.mu.Unlock()
|
|
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
require.Eventually(t, func() bool {
|
|
ap, ok := pool.getAccountPool(accountID)
|
|
if !ok || ap == nil {
|
|
return false
|
|
}
|
|
ap.mu.Lock()
|
|
defer ap.mu.Unlock()
|
|
return !ap.prewarmActive
|
|
}, 2*time.Second, 20*time.Millisecond)
|
|
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
require.Eventually(t, func() bool {
|
|
ap, ok := pool.getAccountPool(accountID)
|
|
if !ok || ap == nil {
|
|
return false
|
|
}
|
|
ap.mu.Lock()
|
|
defer ap.mu.Unlock()
|
|
return !ap.prewarmActive
|
|
}, 2*time.Second, 20*time.Millisecond)
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
|
|
// 连续失败达到阈值后，新的预热触发应被抑制，不再继续拨号。
|
|
pool.ensureTargetIdleAsync(accountID)
|
|
time.Sleep(120 * time.Millisecond)
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireQueueWaitMetrics(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
accountID := int64(99)
|
|
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
conn := newOpenAIWSConn("busy", accountID, &openAIWSFakeConn{}, nil)
|
|
require.True(t, conn.tryAcquire()) // 占用连接，触发后续排队
|
|
|
|
ap := pool.ensureAccountPoolLocked(accountID)
|
|
ap.mu.Lock()
|
|
ap.conns[conn.id] = conn
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
}
|
|
ap.mu.Unlock()
|
|
|
|
go func() {
|
|
time.Sleep(60 * time.Millisecond)
|
|
conn.release()
|
|
}()
|
|
|
|
lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, lease)
|
|
require.True(t, lease.Reused())
|
|
require.GreaterOrEqual(t, lease.QueueWaitDuration(), 50*time.Millisecond)
|
|
lease.Release()
|
|
|
|
metrics := pool.SnapshotMetrics()
|
|
require.GreaterOrEqual(t, metrics.AcquireQueueWaitTotal, int64(1))
|
|
require.Greater(t, metrics.AcquireQueueWaitMsTotal, int64(0))
|
|
require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1))
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_DialSuccessWakesTopologyWaiterAndCanceledWaiterDoesNotLoseLease(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
|
|
pool.setClientDialerForTest(dialer)
|
|
account := &Account{ID: 991, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
|
|
|
type result struct {
|
|
lease *openAIWSConnLease
|
|
err error
|
|
}
|
|
firstCh := make(chan result, 1)
|
|
go func() {
|
|
lease, err := pool.Acquire(context.Background(), req)
|
|
firstCh <- result{lease: lease, err: err}
|
|
}()
|
|
<-dialer.firstStarted
|
|
|
|
waitCtx, cancelWait := context.WithCancel(context.Background())
|
|
secondCh := make(chan result, 1)
|
|
waitReq := req
|
|
waitReq.Headers = http.Header{openAICodexRoutingHintHeader: {"model=gpt-5.6-codex;tier=priority"}}
|
|
go func() {
|
|
lease, err := pool.Acquire(waitCtx, waitReq)
|
|
secondCh <- result{lease: lease, err: err}
|
|
}()
|
|
|
|
close(dialer.releaseFirst)
|
|
first := <-firstCh
|
|
require.NoError(t, first.err)
|
|
require.NotNil(t, first.lease)
|
|
|
|
// The second acquire initially waits on the account topology channel while
|
|
// the first dial is in flight. Dial success must wake it immediately so it
|
|
// can queue on the newly-created (still leased) connection.
|
|
require.Eventually(t, func() bool {
|
|
ap, ok := pool.getAccountPool(account.ID)
|
|
if !ok || ap == nil {
|
|
return false
|
|
}
|
|
ap.mu.Lock()
|
|
defer ap.mu.Unlock()
|
|
for _, conn := range ap.conns {
|
|
if conn != nil && conn.waiters.Load() == 1 {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}, time.Second, 5*time.Millisecond)
|
|
|
|
cancelWait()
|
|
second := <-secondCh
|
|
require.ErrorIs(t, second.err, context.Canceled)
|
|
require.Nil(t, second.lease)
|
|
ap, ok := pool.getAccountPool(account.ID)
|
|
require.True(t, ok)
|
|
ap.mu.Lock()
|
|
require.NotNil(t, ap.lastAcquire)
|
|
require.Empty(t, normalizeOpenAIWSRoutingAffinity(ap.lastAcquire.Headers), "a canceled acquire must not replace the successful prewarm target")
|
|
ap.mu.Unlock()
|
|
first.lease.Release()
|
|
|
|
third, err := pool.Acquire(context.Background(), req)
|
|
require.NoError(t, err)
|
|
require.True(t, third.Reused(), "a canceled waiter must not consume the released semaphore token")
|
|
require.Equal(t, first.lease.ConnID(), third.ConnID())
|
|
third.Release()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_PrewarmHintChangeDoesNotInvalidateHealthyDial(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
|
|
pool.setClientDialerForTest(dialer)
|
|
account := &Account{ID: 992, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
oldHeaders := make(http.Header)
|
|
oldHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex")
|
|
newHeaders := make(http.Header)
|
|
newHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
|
|
ap := pool.getOrCreateAccountPool(account.ID)
|
|
ap.mu.Lock()
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: oldHeaders,
|
|
}
|
|
ap.mu.Unlock()
|
|
|
|
pool.ensureTargetIdleAsync(account.ID)
|
|
<-dialer.firstStarted
|
|
|
|
// Simulate a newer priority target arriving while the old model-only
|
|
// prewarm dial is in flight. Routing hints are advisory, so this alone must
|
|
// not discard an otherwise compatible connection.
|
|
ap.mu.Lock()
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: newHeaders,
|
|
}
|
|
ap.mu.Unlock()
|
|
close(dialer.releaseFirst)
|
|
|
|
require.Eventually(t, func() bool {
|
|
ap.mu.Lock()
|
|
defer ap.mu.Unlock()
|
|
if ap.prewarmActive || len(ap.conns) != 1 {
|
|
return false
|
|
}
|
|
for _, conn := range ap.conns {
|
|
return conn != nil && conn.routingAffinity == "model=gpt-5.6-codex"
|
|
}
|
|
return false
|
|
}, 2*time.Second, 10*time.Millisecond)
|
|
require.Equal(t, 1, dialer.DialCount(), "routing-hint-only changes must not turn advisory metadata into hard reconnects")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_ClearAccountWakesIncompatibleTopologyWaiter(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := &Account{ID: 993, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
|
betaAReq := baseReq
|
|
betaAReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_a"}}
|
|
betaBReq := baseReq
|
|
betaBReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_b"}}
|
|
|
|
busy, err := pool.Acquire(context.Background(), betaAReq)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, dialer.DialCount())
|
|
|
|
type result struct {
|
|
lease *openAIWSConnLease
|
|
err error
|
|
}
|
|
resultCh := make(chan result, 1)
|
|
waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancelWait()
|
|
go func() {
|
|
lease, acquireErr := pool.Acquire(waitCtx, betaBReq)
|
|
resultCh <- result{lease: lease, err: acquireErr}
|
|
}()
|
|
|
|
require.Never(t, func() bool { return dialer.DialCount() > 1 }, 50*time.Millisecond, 5*time.Millisecond)
|
|
pool.ClearAccount(account.ID)
|
|
|
|
resultB := <-resultCh
|
|
require.NoError(t, resultB.err)
|
|
require.NotNil(t, resultB.lease)
|
|
require.False(t, resultB.lease.Reused())
|
|
require.Equal(t, 2, dialer.DialCount(), "ClearAccount must wake the waiter to redial immediately")
|
|
resultB.lease.Release()
|
|
busy.Release()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_ClearAccountDoesNotReviveInFlightDialGeneration(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
|
|
pool.setClientDialerForTest(dialer)
|
|
account := &Account{ID: 994, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
|
|
|
type result struct {
|
|
lease *openAIWSConnLease
|
|
err error
|
|
}
|
|
resultCh := make(chan result, 1)
|
|
go func() {
|
|
lease, err := pool.Acquire(context.Background(), req)
|
|
resultCh <- result{lease: lease, err: err}
|
|
}()
|
|
<-dialer.firstStarted
|
|
|
|
pool.ClearAccount(account.ID)
|
|
close(dialer.releaseFirst)
|
|
got := <-resultCh
|
|
require.NoError(t, got.err)
|
|
require.NotNil(t, got.lease)
|
|
require.Equal(t, 2, dialer.DialCount(), "the pre-clear dial must be discarded and retried in the new generation")
|
|
require.True(t, strings.HasSuffix(got.lease.ConnID(), "_2"))
|
|
|
|
ap, ok := pool.getAccountPool(account.ID)
|
|
require.True(t, ok)
|
|
ap.mu.Lock()
|
|
require.Equal(t, uint64(1), ap.generation)
|
|
require.Len(t, ap.conns, 1)
|
|
require.NotNil(t, ap.lastAcquire, "only the post-clear successful acquire may restore the prewarm target")
|
|
ap.mu.Unlock()
|
|
got.lease.Release()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_ForceNewConnSkipsReuse(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
|
|
account := &Account{ID: 123, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
|
|
lease1, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, lease1)
|
|
lease1.Release()
|
|
|
|
lease2, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
ForceNewConn: true,
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, lease2)
|
|
lease2.Release()
|
|
|
|
require.Equal(t, 2, dialer.DialCount(), "ForceNewConn=true 时应跳过空闲连接复用并新建连接")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireReusesOnlyMatchingBetaFeatures(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
|
|
account := &Account{ID: 128, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
baseReq := openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
}
|
|
|
|
plainLease, err := pool.Acquire(context.Background(), baseReq)
|
|
require.NoError(t, err)
|
|
plainConnID := plainLease.ConnID()
|
|
plainLease.Release()
|
|
|
|
betaReq := baseReq
|
|
betaReq.Headers = http.Header{"X-Codex-Beta-Features": {" remote_compaction_v2 ", " responses_websockets_v2 "}}
|
|
betaLease, err := pool.Acquire(context.Background(), betaReq)
|
|
require.NoError(t, err)
|
|
require.False(t, betaLease.Reused())
|
|
require.NotEqual(t, plainConnID, betaLease.ConnID())
|
|
betaConnID := betaLease.ConnID()
|
|
betaLease.Release()
|
|
|
|
reorderedReq := baseReq
|
|
reorderedReq.Headers = http.Header{"X-Codex-Beta-Features": {"responses_websockets_v2,remote_compaction_v2"}}
|
|
reorderedLease, err := pool.Acquire(context.Background(), reorderedReq)
|
|
require.NoError(t, err)
|
|
require.True(t, reorderedLease.Reused())
|
|
require.Equal(t, betaConnID, reorderedLease.ConnID())
|
|
reorderedLease.Release()
|
|
|
|
_, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: baseReq.WSURL,
|
|
Headers: betaReq.Headers,
|
|
PreferredConnID: plainConnID,
|
|
ForcePreferredConn: true,
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable)
|
|
|
|
plainLease, err = pool.Acquire(context.Background(), baseReq)
|
|
require.NoError(t, err)
|
|
require.True(t, plainLease.Reused())
|
|
require.Equal(t, plainConnID, plainLease.ConnID())
|
|
plainLease.Release()
|
|
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
}
|
|
|
|
func activeCodexFingerprintPoolAccountForTest(id int64) *Account {
|
|
return &Account{
|
|
ID: id,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
Extra: map[string]any{
|
|
codexFingerprintModeExtraKey: "session",
|
|
codexFingerprintSeedExtraKey: "11111111-1111-4111-8111-111111111111",
|
|
},
|
|
}
|
|
}
|
|
|
|
func stableOpenAIWSIdentityHeadersForTest() http.Header {
|
|
headers := make(http.Header)
|
|
headers.Set("X-Codex-Beta-Features", "remote_compaction_v2,responses_websockets_v2")
|
|
headers.Set("X-Codex-Installation-ID", "install-a")
|
|
headers.Set("session-id", "session-hyphen-a")
|
|
headers.Set("session_id", "session-underscore-a")
|
|
headers.Set("thread-id", "thread-a")
|
|
headers.Set("x-client-request-id", "client-request-a")
|
|
headers.Set("x-codex-window-id", "window-a")
|
|
return headers
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireReusesSameStableIdentityWithDifferentTurnMetadata(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := activeCodexFingerprintPoolAccountForTest(132)
|
|
headers := stableOpenAIWSIdentityHeadersForTest()
|
|
headers.Set("Authorization", "Bearer token-a")
|
|
headers.Set("x-codex-turn-metadata", `{"turn_id":"turn-a"}`)
|
|
|
|
first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: headers,
|
|
})
|
|
require.NoError(t, err)
|
|
firstConnID := first.ConnID()
|
|
first.Release()
|
|
|
|
nextHeaders := stableOpenAIWSIdentityHeadersForTest()
|
|
nextHeaders.Set("Authorization", "Bearer token-b")
|
|
nextHeaders.Set("x-codex-turn-metadata", `{"turn_id":"turn-b"}`)
|
|
nextHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
|
|
second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: nextHeaders,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, second.Reused())
|
|
require.Equal(t, firstConnID, second.ConnID())
|
|
second.Release()
|
|
require.Equal(t, 1, dialer.DialCount(), "stable identity match should ignore auth, turn metadata, and soft routing hints")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireDoesNotReuseDifferentStableIdentity(t *testing.T) {
|
|
for _, tt := range []struct {
|
|
name string
|
|
header string
|
|
value string
|
|
}{
|
|
{name: "installation", header: "x-codex-installation-id", value: "install-b"},
|
|
{name: "session hyphen", header: "session-id", value: "session-hyphen-b"},
|
|
{name: "session underscore", header: "session_id", value: "session-underscore-b"},
|
|
{name: "thread", header: "thread-id", value: "thread-b"},
|
|
{name: "client request", header: "x-client-request-id", value: "client-request-b"},
|
|
{name: "window", header: "x-codex-window-id", value: "window-b"},
|
|
} {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := activeCodexFingerprintPoolAccountForTest(133)
|
|
|
|
first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: stableOpenAIWSIdentityHeadersForTest(),
|
|
})
|
|
require.NoError(t, err)
|
|
firstConnID := first.ConnID()
|
|
first.Release()
|
|
|
|
nextHeaders := stableOpenAIWSIdentityHeadersForTest()
|
|
nextHeaders.Set(tt.header, tt.value)
|
|
second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: nextHeaders,
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, second.Reused())
|
|
require.NotEqual(t, firstConnID, second.ConnID())
|
|
second.Release()
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireRoutingHintRemainsSoftAffinity(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := activeCodexFingerprintPoolAccountForTest(134)
|
|
|
|
firstHeaders := stableOpenAIWSIdentityHeadersForTest()
|
|
firstHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex")
|
|
first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: firstHeaders,
|
|
})
|
|
require.NoError(t, err)
|
|
firstConnID := first.ConnID()
|
|
first.Release()
|
|
|
|
secondHeaders := stableOpenAIWSIdentityHeadersForTest()
|
|
secondHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
|
|
second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: secondHeaders,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, second.Reused())
|
|
require.Equal(t, firstConnID, second.ConnID())
|
|
second.Release()
|
|
require.Equal(t, 1, dialer.DialCount())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_DeviceModeKeysOnlyInstallationIdentity(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := activeCodexFingerprintPoolAccountForTest(135)
|
|
account.Extra[codexFingerprintModeExtraKey] = "device"
|
|
|
|
firstHeaders := stableOpenAIWSIdentityHeadersForTest()
|
|
first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: firstHeaders,
|
|
})
|
|
require.NoError(t, err)
|
|
firstConnID := first.ConnID()
|
|
first.Release()
|
|
|
|
sessionChanged := stableOpenAIWSIdentityHeadersForTest()
|
|
sessionChanged.Set("session-id", "session-hyphen-b")
|
|
sessionChanged.Set("session_id", "session-underscore-b")
|
|
sessionChanged.Set("thread-id", "thread-b")
|
|
sessionChanged.Set("x-client-request-id", "client-request-b")
|
|
sessionChanged.Set("x-codex-window-id", "window-b")
|
|
second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: sessionChanged,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, second.Reused())
|
|
require.Equal(t, firstConnID, second.ConnID())
|
|
second.Release()
|
|
|
|
installationChanged := sessionChanged.Clone()
|
|
installationChanged.Set("x-codex-installation-id", "install-b")
|
|
third, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: installationChanged,
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, third.Reused())
|
|
require.NotEqual(t, firstConnID, third.ConnID())
|
|
third.Release()
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireReplacesIdleConnWithDifferentBetaFeatures(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
|
|
account := &Account{ID: 129, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
plainLease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.NoError(t, err)
|
|
plainConnID := plainLease.ConnID()
|
|
plainLease.Release()
|
|
|
|
betaLease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
Headers: http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}},
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, betaLease.Reused())
|
|
require.NotEqual(t, plainConnID, betaLease.ConnID())
|
|
betaLease.Release()
|
|
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireWaitsForBusyIncompatibleConnection(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := &Account{ID: 130, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
|
|
|
plainLease, err := pool.Acquire(context.Background(), baseReq)
|
|
require.NoError(t, err)
|
|
plainConnID := plainLease.ConnID()
|
|
|
|
type acquireResult struct {
|
|
lease *openAIWSConnLease
|
|
err error
|
|
}
|
|
resultCh := make(chan acquireResult, 1)
|
|
var done atomic.Bool
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
go func() {
|
|
betaReq := baseReq
|
|
betaReq.Headers = http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}}
|
|
lease, acquireErr := pool.Acquire(ctx, betaReq)
|
|
resultCh <- acquireResult{lease: lease, err: acquireErr}
|
|
done.Store(true)
|
|
}()
|
|
|
|
require.Never(t, done.Load, 50*time.Millisecond, 5*time.Millisecond)
|
|
plainLease.Release()
|
|
|
|
result := <-resultCh
|
|
require.NoError(t, result.err)
|
|
require.NotNil(t, result.lease)
|
|
require.False(t, result.lease.Reused())
|
|
require.NotEqual(t, plainConnID, result.lease.ConnID())
|
|
result.lease.Release()
|
|
require.Equal(t, 2, dialer.DialCount())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireReplacesIncompatibleIdleWhenMatchingBusy(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
dialer := &openAIWSCountingDialer{}
|
|
pool.setClientDialerForTest(dialer)
|
|
account := &Account{ID: 131, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
|
|
|
plainLease, err := pool.Acquire(context.Background(), baseReq)
|
|
require.NoError(t, err)
|
|
plainConnID := plainLease.ConnID()
|
|
plainLease.Release()
|
|
|
|
betaReq := baseReq
|
|
betaReq.Headers = http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}}
|
|
busyBetaLease, err := pool.Acquire(context.Background(), betaReq)
|
|
require.NoError(t, err)
|
|
|
|
secondBetaLease, err := pool.Acquire(context.Background(), betaReq)
|
|
require.NoError(t, err)
|
|
require.False(t, secondBetaLease.Reused())
|
|
require.NotEqual(t, plainConnID, secondBetaLease.ConnID())
|
|
require.NotEqual(t, busyBetaLease.ConnID(), secondBetaLease.ConnID())
|
|
|
|
secondBetaLease.Release()
|
|
busyBetaLease.Release()
|
|
require.Equal(t, 3, dialer.DialCount())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireForcePreferredConnUnavailable(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
account := &Account{ID: 124, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := pool.getOrCreateAccountPool(account.ID)
|
|
otherConn := newOpenAIWSConn("other_conn", account.ID, &openAIWSFakeConn{}, nil)
|
|
ap.mu.Lock()
|
|
ap.conns[otherConn.id] = otherConn
|
|
ap.mu.Unlock()
|
|
|
|
_, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
ForcePreferredConn: true,
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable)
|
|
|
|
_, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
PreferredConnID: "missing_conn",
|
|
ForcePreferredConn: true,
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireForcePreferredConnQueuesOnPreferredOnly(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
account := &Account{ID: 125, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := pool.getOrCreateAccountPool(account.ID)
|
|
preferredConn := newOpenAIWSConn("preferred_conn", account.ID, &openAIWSFakeConn{}, nil)
|
|
otherConn := newOpenAIWSConn("other_conn_idle", account.ID, &openAIWSFakeConn{}, nil)
|
|
require.True(t, preferredConn.tryAcquire(), "先占用 preferred 连接，触发排队获取")
|
|
ap.mu.Lock()
|
|
ap.conns[preferredConn.id] = preferredConn
|
|
ap.conns[otherConn.id] = otherConn
|
|
ap.lastCleanupAt = time.Now()
|
|
ap.mu.Unlock()
|
|
|
|
go func() {
|
|
time.Sleep(60 * time.Millisecond)
|
|
preferredConn.release()
|
|
}()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
lease, err := pool.Acquire(ctx, openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
PreferredConnID: preferredConn.id,
|
|
ForcePreferredConn: true,
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, lease)
|
|
require.Equal(t, preferredConn.id, lease.ConnID(), "严格模式应只等待并复用 preferred 连接，不可漂移")
|
|
require.GreaterOrEqual(t, lease.QueueWaitDuration(), 40*time.Millisecond)
|
|
lease.Release()
|
|
require.True(t, otherConn.tryAcquire(), "other 连接不应被严格模式抢占")
|
|
otherConn.release()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireForcePreferredConnDirectAndQueueFull(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
account := &Account{ID: 127, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := pool.getOrCreateAccountPool(account.ID)
|
|
preferredConn := newOpenAIWSConn("preferred_conn_direct", account.ID, &openAIWSFakeConn{}, nil)
|
|
otherConn := newOpenAIWSConn("other_conn_direct", account.ID, &openAIWSFakeConn{}, nil)
|
|
ap.mu.Lock()
|
|
ap.conns[preferredConn.id] = preferredConn
|
|
ap.conns[otherConn.id] = otherConn
|
|
ap.lastCleanupAt = time.Now()
|
|
ap.mu.Unlock()
|
|
|
|
lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
PreferredConnID: preferredConn.id,
|
|
ForcePreferredConn: true,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, preferredConn.id, lease.ConnID(), "preferred 空闲时应直接命中")
|
|
lease.Release()
|
|
|
|
require.True(t, preferredConn.tryAcquire())
|
|
preferredConn.waiters.Store(1)
|
|
_, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
PreferredConnID: preferredConn.id,
|
|
ForcePreferredConn: true,
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSConnQueueFull, "严格模式下队列满应直接失败，不得漂移")
|
|
preferredConn.waiters.Store(0)
|
|
preferredConn.release()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_CleanupSkipsPinnedConn(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 0
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
accountID := int64(126)
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
pinnedConn := newOpenAIWSConn("pinned_conn", accountID, &openAIWSFakeConn{}, nil)
|
|
idleConn := newOpenAIWSConn("idle_conn", accountID, &openAIWSFakeConn{}, nil)
|
|
ap.mu.Lock()
|
|
ap.conns[pinnedConn.id] = pinnedConn
|
|
ap.conns[idleConn.id] = idleConn
|
|
ap.mu.Unlock()
|
|
|
|
require.True(t, pool.PinConn(accountID, pinnedConn.id))
|
|
evicted := pool.cleanupAccountLocked(ap, time.Now(), pool.maxConnsHardCap())
|
|
closeOpenAIWSConns(evicted)
|
|
|
|
ap.mu.Lock()
|
|
_, pinnedExists := ap.conns[pinnedConn.id]
|
|
_, idleExists := ap.conns[idleConn.id]
|
|
ap.mu.Unlock()
|
|
require.True(t, pinnedExists, "被 active ingress 绑定的连接不应被 cleanup 回收")
|
|
require.False(t, idleExists, "非绑定的空闲连接应被回收")
|
|
|
|
pool.UnpinConn(accountID, pinnedConn.id)
|
|
evicted = pool.cleanupAccountLocked(ap, time.Now(), pool.maxConnsHardCap())
|
|
closeOpenAIWSConns(evicted)
|
|
ap.mu.Lock()
|
|
_, pinnedExists = ap.conns[pinnedConn.id]
|
|
ap.mu.Unlock()
|
|
require.False(t, pinnedExists, "解绑后连接应可被正常回收")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_PinUnpinConnBranches(t *testing.T) {
|
|
var nilPool *openAIWSConnPool
|
|
require.False(t, nilPool.PinConn(1, "x"))
|
|
nilPool.UnpinConn(1, "x")
|
|
|
|
cfg := &config.Config{}
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
accountID := int64(128)
|
|
ap := &openAIWSAccountPool{
|
|
conns: map[string]*openAIWSConn{},
|
|
}
|
|
pool.accounts.Store(accountID, ap)
|
|
|
|
require.False(t, pool.PinConn(0, "x"))
|
|
require.False(t, pool.PinConn(999, "x"))
|
|
require.False(t, pool.PinConn(accountID, ""))
|
|
require.False(t, pool.PinConn(accountID, "missing"))
|
|
|
|
conn := newOpenAIWSConn("pin_refcount", accountID, &openAIWSFakeConn{}, nil)
|
|
ap.mu.Lock()
|
|
ap.conns[conn.id] = conn
|
|
ap.mu.Unlock()
|
|
require.True(t, pool.PinConn(accountID, conn.id))
|
|
require.True(t, pool.PinConn(accountID, conn.id))
|
|
|
|
ap.mu.Lock()
|
|
require.Equal(t, 2, ap.pinnedConns[conn.id])
|
|
ap.mu.Unlock()
|
|
|
|
pool.UnpinConn(accountID, conn.id)
|
|
ap.mu.Lock()
|
|
require.Equal(t, 1, ap.pinnedConns[conn.id])
|
|
ap.mu.Unlock()
|
|
|
|
pool.UnpinConn(accountID, conn.id)
|
|
ap.mu.Lock()
|
|
_, exists := ap.pinnedConns[conn.id]
|
|
ap.mu.Unlock()
|
|
require.False(t, exists)
|
|
|
|
pool.UnpinConn(accountID, conn.id)
|
|
pool.UnpinConn(accountID, "")
|
|
pool.UnpinConn(0, conn.id)
|
|
pool.UnpinConn(999, conn.id)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
|
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
|
|
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1.0
|
|
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 0.6
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
oauthHigh := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 10}
|
|
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(oauthHigh), "应受全局硬上限约束")
|
|
|
|
oauthLow := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 3}
|
|
require.Equal(t, 3, pool.effectiveMaxConnsByAccount(oauthLow))
|
|
|
|
apiKeyHigh := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 10}
|
|
require.Equal(t, 6, pool.effectiveMaxConnsByAccount(apiKeyHigh), "API Key 应按系数缩放")
|
|
|
|
apiKeyLow := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
|
|
require.Equal(t, 1, pool.effectiveMaxConnsByAccount(apiKeyLow), "最小值应保持为 1")
|
|
|
|
unlimited := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 0}
|
|
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(unlimited), "无限并发应回退到全局硬上限")
|
|
|
|
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(nil), "缺少账号上下文应回退到全局硬上限")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
|
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = false
|
|
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1.0
|
|
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 1.0
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 2}
|
|
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
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20}
|
|
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束")
|
|
|
|
nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0}
|
|
require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_AcquireRejectsWhenEffectiveMaxConnsIsZero(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
account := &Account{ID: 901, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 0}
|
|
_, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSConnQueueFull)
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_ReadMessageWithContextTimeout_PerRead(t *testing.T) {
|
|
conn := newOpenAIWSConn("timeout", 1, &openAIWSBlockingConn{readDelay: 80 * time.Millisecond}, nil)
|
|
lease := &openAIWSConnLease{conn: conn}
|
|
|
|
_, err := lease.ReadMessageWithContextTimeout(context.Background(), 20*time.Millisecond)
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, context.DeadlineExceeded)
|
|
|
|
payload, err := lease.ReadMessageWithContextTimeout(context.Background(), 150*time.Millisecond)
|
|
require.NoError(t, err)
|
|
require.Contains(t, string(payload), "response.completed")
|
|
|
|
parentCtx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
_, err = lease.ReadMessageWithContextTimeout(parentCtx, 150*time.Millisecond)
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_WriteJSONWithContextTimeout_RespectsParentContext(t *testing.T) {
|
|
conn := newOpenAIWSConn("write_timeout_ctx", 1, &openAIWSWriteBlockingConn{}, nil)
|
|
lease := &openAIWSConnLease{conn: conn}
|
|
|
|
parentCtx, cancel := context.WithCancel(context.Background())
|
|
go func() {
|
|
time.Sleep(20 * time.Millisecond)
|
|
cancel()
|
|
}()
|
|
|
|
start := time.Now()
|
|
err := lease.WriteJSONWithContextTimeout(parentCtx, map[string]any{"type": "response.create"}, 2*time.Minute)
|
|
elapsed := time.Since(start)
|
|
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
require.Less(t, elapsed, 200*time.Millisecond)
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_PingWithTimeout(t *testing.T) {
|
|
conn := newOpenAIWSConn("ping_ok", 1, &openAIWSFakeConn{}, nil)
|
|
lease := &openAIWSConnLease{conn: conn}
|
|
require.NoError(t, lease.PingWithTimeout(50*time.Millisecond))
|
|
|
|
var nilLease *openAIWSConnLease
|
|
err := nilLease.PingWithTimeout(50 * time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
}
|
|
|
|
func TestOpenAIWSConn_ReadAndWriteCanProceedConcurrently(t *testing.T) {
|
|
conn := newOpenAIWSConn("full_duplex", 1, &openAIWSBlockingConn{readDelay: 120 * time.Millisecond}, nil)
|
|
|
|
readDone := make(chan error, 1)
|
|
go func() {
|
|
_, err := conn.readMessageWithContextTimeout(context.Background(), 200*time.Millisecond)
|
|
readDone <- err
|
|
}()
|
|
|
|
// 让读取先占用 readMu。
|
|
time.Sleep(20 * time.Millisecond)
|
|
|
|
start := time.Now()
|
|
err := conn.pingWithTimeout(50 * time.Millisecond)
|
|
elapsed := time.Since(start)
|
|
|
|
require.NoError(t, err)
|
|
require.Less(t, elapsed, 80*time.Millisecond, "写路径不应被读锁长期阻塞")
|
|
require.NoError(t, <-readDone)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_BackgroundPingSweep_EvictsDeadIdleConn(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
accountID := int64(301)
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
conn := newOpenAIWSConn("dead_idle", accountID, &openAIWSPingFailConn{}, nil)
|
|
ap.mu.Lock()
|
|
ap.conns[conn.id] = conn
|
|
ap.mu.Unlock()
|
|
|
|
pool.runBackgroundPingSweep()
|
|
|
|
ap.mu.Lock()
|
|
_, exists := ap.conns[conn.id]
|
|
ap.mu.Unlock()
|
|
require.False(t, exists, "后台 ping 失败的空闲连接应被回收")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_BackgroundCleanupSweep_WithoutAcquire(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
accountID := int64(302)
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
stale := newOpenAIWSConn("stale_bg", accountID, &openAIWSFakeConn{}, nil)
|
|
stale.createdAtNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
|
|
stale.lastUsedNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
|
|
ap.mu.Lock()
|
|
ap.conns[stale.id] = stale
|
|
ap.mu.Unlock()
|
|
|
|
pool.runBackgroundCleanupSweep(time.Now())
|
|
|
|
ap.mu.Lock()
|
|
_, exists := ap.conns[stale.id]
|
|
ap.mu.Unlock()
|
|
require.False(t, exists, "后台清理应在无新 acquire 时也回收过期连接")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_BackgroundWorkerGuardBranches(t *testing.T) {
|
|
var nilPool *openAIWSConnPool
|
|
require.NotPanics(t, func() {
|
|
nilPool.startBackgroundWorkers()
|
|
nilPool.runBackgroundPingWorker()
|
|
nilPool.runBackgroundPingSweep()
|
|
_ = nilPool.snapshotIdleConnsForPing()
|
|
nilPool.runBackgroundCleanupWorker()
|
|
nilPool.runBackgroundCleanupSweep(time.Now())
|
|
})
|
|
|
|
poolNoStop := &openAIWSConnPool{}
|
|
require.NotPanics(t, func() {
|
|
poolNoStop.startBackgroundWorkers()
|
|
})
|
|
|
|
poolStopPing := &openAIWSConnPool{workerStopCh: make(chan struct{})}
|
|
pingDone := make(chan struct{})
|
|
go func() {
|
|
poolStopPing.runBackgroundPingWorker()
|
|
close(pingDone)
|
|
}()
|
|
close(poolStopPing.workerStopCh)
|
|
select {
|
|
case <-pingDone:
|
|
case <-time.After(500 * time.Millisecond):
|
|
t.Fatal("runBackgroundPingWorker 未在 stop 信号后退出")
|
|
}
|
|
|
|
poolStopCleanup := &openAIWSConnPool{workerStopCh: make(chan struct{})}
|
|
cleanupDone := make(chan struct{})
|
|
go func() {
|
|
poolStopCleanup.runBackgroundCleanupWorker()
|
|
close(cleanupDone)
|
|
}()
|
|
close(poolStopCleanup.workerStopCh)
|
|
select {
|
|
case <-cleanupDone:
|
|
case <-time.After(500 * time.Millisecond):
|
|
t.Fatal("runBackgroundCleanupWorker 未在 stop 信号后退出")
|
|
}
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_SnapshotIdleConnsForPing_SkipsInvalidEntries(t *testing.T) {
|
|
pool := &openAIWSConnPool{}
|
|
pool.accounts.Store("invalid-key", &openAIWSAccountPool{})
|
|
pool.accounts.Store(int64(123), "invalid-value")
|
|
|
|
accountID := int64(123)
|
|
ap := &openAIWSAccountPool{
|
|
conns: make(map[string]*openAIWSConn),
|
|
}
|
|
ap.conns["nil_conn"] = nil
|
|
|
|
leased := newOpenAIWSConn("leased", accountID, &openAIWSFakeConn{}, nil)
|
|
require.True(t, leased.tryAcquire())
|
|
ap.conns[leased.id] = leased
|
|
|
|
waiting := newOpenAIWSConn("waiting", accountID, &openAIWSFakeConn{}, nil)
|
|
waiting.waiters.Store(1)
|
|
ap.conns[waiting.id] = waiting
|
|
|
|
idle := newOpenAIWSConn("idle", accountID, &openAIWSFakeConn{}, nil)
|
|
ap.conns[idle.id] = idle
|
|
|
|
pool.accounts.Store(accountID, ap)
|
|
candidates := pool.snapshotIdleConnsForPing()
|
|
require.Len(t, candidates, 1)
|
|
require.Equal(t, idle.id, candidates[0].conn.id)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_RunBackgroundCleanupSweep_SkipsInvalidAndUsesAccountCap(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
|
|
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
|
|
|
|
pool := &openAIWSConnPool{cfg: cfg}
|
|
pool.accounts.Store("bad-key", "bad-value")
|
|
|
|
accountID := int64(2026)
|
|
ap := &openAIWSAccountPool{
|
|
conns: make(map[string]*openAIWSConn),
|
|
}
|
|
ap.conns["nil_conn"] = nil
|
|
stale := newOpenAIWSConn("stale_bg_cleanup", accountID, &openAIWSFakeConn{}, nil)
|
|
stale.createdAtNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
|
|
stale.lastUsedNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
|
|
ap.conns[stale.id] = stale
|
|
ap.lastAcquire = &openAIWSAcquireRequest{
|
|
Account: &Account{
|
|
ID: accountID,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeAPIKey,
|
|
Concurrency: 1,
|
|
},
|
|
}
|
|
pool.accounts.Store(accountID, ap)
|
|
|
|
now := time.Now()
|
|
require.NotPanics(t, func() {
|
|
pool.runBackgroundCleanupSweep(now)
|
|
})
|
|
|
|
ap.mu.Lock()
|
|
_, nilConnExists := ap.conns["nil_conn"]
|
|
_, exists := ap.conns[stale.id]
|
|
lastCleanupAt := ap.lastCleanupAt
|
|
ap.mu.Unlock()
|
|
|
|
require.False(t, nilConnExists, "后台清理应移除无效 nil 连接条目")
|
|
require.False(t, exists, "后台清理应清理过期连接")
|
|
require.Equal(t, now, lastCleanupAt)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_QueueLimitPerConn_DefaultAndConfigured(t *testing.T) {
|
|
var nilPool *openAIWSConnPool
|
|
require.Equal(t, 256, nilPool.queueLimitPerConn())
|
|
|
|
pool := &openAIWSConnPool{cfg: &config.Config{}}
|
|
require.Equal(t, 256, pool.queueLimitPerConn())
|
|
|
|
pool.cfg.Gateway.OpenAIWS.QueueLimitPerConn = 9
|
|
require.Equal(t, 9, pool.queueLimitPerConn())
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_Close(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
// Close 应该可以安全调用
|
|
pool.Close()
|
|
|
|
// workerStopCh 应已关闭
|
|
select {
|
|
case <-pool.workerStopCh:
|
|
// 预期：channel 已关闭
|
|
default:
|
|
t.Fatal("Close 后 workerStopCh 应已关闭")
|
|
}
|
|
|
|
// 多次调用 Close 不应 panic
|
|
pool.Close()
|
|
|
|
// nil pool 调用 Close 不应 panic
|
|
var nilPool *openAIWSConnPool
|
|
nilPool.Close()
|
|
}
|
|
|
|
func TestOpenAIWSDialError_ErrorAndUnwrap(t *testing.T) {
|
|
baseErr := errors.New("boom")
|
|
dialErr := &openAIWSDialError{StatusCode: 502, Err: baseErr}
|
|
require.Contains(t, dialErr.Error(), "status=502")
|
|
require.ErrorIs(t, dialErr.Unwrap(), baseErr)
|
|
|
|
noStatus := &openAIWSDialError{Err: baseErr}
|
|
require.Contains(t, noStatus.Error(), "boom")
|
|
|
|
var nilDialErr *openAIWSDialError
|
|
require.Equal(t, "", nilDialErr.Error())
|
|
require.NoError(t, nilDialErr.Unwrap())
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_ReadWriteHelpersAndConnStats(t *testing.T) {
|
|
conn := newOpenAIWSConn("helper_conn", 1, &openAIWSFakeConn{}, http.Header{
|
|
"X-Test": []string{" value "},
|
|
})
|
|
lease := &openAIWSConnLease{conn: conn}
|
|
|
|
require.NoError(t, lease.WriteJSONContext(context.Background(), map[string]any{"type": "response.create"}))
|
|
payload, err := lease.ReadMessage(100 * time.Millisecond)
|
|
require.NoError(t, err)
|
|
require.Contains(t, string(payload), "response.completed")
|
|
|
|
payload, err = lease.ReadMessageContext(context.Background())
|
|
require.NoError(t, err)
|
|
require.Contains(t, string(payload), "response.completed")
|
|
|
|
payload, err = conn.readMessageWithTimeout(100 * time.Millisecond)
|
|
require.NoError(t, err)
|
|
require.Contains(t, string(payload), "response.completed")
|
|
|
|
require.Equal(t, "value", conn.handshakeHeader(" X-Test "))
|
|
require.NotZero(t, conn.createdAt())
|
|
require.NotZero(t, conn.lastUsedAt())
|
|
require.GreaterOrEqual(t, conn.age(time.Now()), time.Duration(0))
|
|
require.GreaterOrEqual(t, conn.idleDuration(time.Now()), time.Duration(0))
|
|
require.False(t, conn.isLeased())
|
|
|
|
// 覆盖空上下文路径
|
|
_, err = conn.readMessage(context.Background())
|
|
require.NoError(t, err)
|
|
|
|
// 覆盖 nil 保护分支
|
|
var nilConn *openAIWSConn
|
|
require.ErrorIs(t, nilConn.writeJSONWithTimeout(context.Background(), map[string]any{}, time.Second), errOpenAIWSConnClosed)
|
|
_, err = nilConn.readMessageWithTimeout(10 * time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
_, err = nilConn.readMessageWithContextTimeout(context.Background(), 10*time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_PickOldestIdleAndAccountPoolLoad(t *testing.T) {
|
|
pool := &openAIWSConnPool{}
|
|
accountID := int64(404)
|
|
ap := &openAIWSAccountPool{conns: map[string]*openAIWSConn{}}
|
|
|
|
idleOld := newOpenAIWSConn("idle_old", accountID, &openAIWSFakeConn{}, nil)
|
|
idleOld.lastUsedNano.Store(time.Now().Add(-10 * time.Minute).UnixNano())
|
|
idleNew := newOpenAIWSConn("idle_new", accountID, &openAIWSFakeConn{}, nil)
|
|
idleNew.lastUsedNano.Store(time.Now().Add(-1 * time.Minute).UnixNano())
|
|
leased := newOpenAIWSConn("leased", accountID, &openAIWSFakeConn{}, nil)
|
|
require.True(t, leased.tryAcquire())
|
|
leased.waiters.Store(2)
|
|
|
|
ap.conns[idleOld.id] = idleOld
|
|
ap.conns[idleNew.id] = idleNew
|
|
ap.conns[leased.id] = leased
|
|
|
|
oldest := pool.pickOldestIdleConnLocked(ap)
|
|
require.NotNil(t, oldest)
|
|
require.Equal(t, idleOld.id, oldest.id)
|
|
|
|
inflight, waiters := accountPoolLoadLocked(ap)
|
|
require.Equal(t, 1, inflight)
|
|
require.Equal(t, 2, waiters)
|
|
|
|
pool.accounts.Store(accountID, ap)
|
|
loadInflight, loadWaiters, conns := pool.AccountPoolLoad(accountID)
|
|
require.Equal(t, 1, loadInflight)
|
|
require.Equal(t, 2, loadWaiters)
|
|
require.Equal(t, 3, conns)
|
|
|
|
zeroInflight, zeroWaiters, zeroConns := pool.AccountPoolLoad(0)
|
|
require.Equal(t, 0, zeroInflight)
|
|
require.Equal(t, 0, zeroWaiters)
|
|
require.Equal(t, 0, zeroConns)
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_Close_WaitsWorkerGroupAndNilStopChannel(t *testing.T) {
|
|
pool := &openAIWSConnPool{}
|
|
release := make(chan struct{})
|
|
pool.workerWg.Add(1)
|
|
go func() {
|
|
defer pool.workerWg.Done()
|
|
<-release
|
|
}()
|
|
|
|
closed := make(chan struct{})
|
|
go func() {
|
|
pool.Close()
|
|
close(closed)
|
|
}()
|
|
|
|
select {
|
|
case <-closed:
|
|
t.Fatal("Close 不应在 WaitGroup 未完成时提前返回")
|
|
case <-time.After(30 * time.Millisecond):
|
|
}
|
|
|
|
close(release)
|
|
select {
|
|
case <-closed:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("Close 未等待 workerWg 完成")
|
|
}
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_Close_ClosesOnlyIdleConnections(t *testing.T) {
|
|
pool := &openAIWSConnPool{
|
|
workerStopCh: make(chan struct{}),
|
|
}
|
|
|
|
accountID := int64(606)
|
|
ap := &openAIWSAccountPool{
|
|
conns: map[string]*openAIWSConn{},
|
|
}
|
|
idle := newOpenAIWSConn("idle_conn", accountID, &openAIWSFakeConn{}, nil)
|
|
leased := newOpenAIWSConn("leased_conn", accountID, &openAIWSFakeConn{}, nil)
|
|
require.True(t, leased.tryAcquire())
|
|
|
|
ap.conns[idle.id] = idle
|
|
ap.conns[leased.id] = leased
|
|
pool.accounts.Store(accountID, ap)
|
|
pool.accounts.Store("invalid-key", "invalid-value")
|
|
|
|
pool.Close()
|
|
|
|
select {
|
|
case <-idle.closedCh:
|
|
// idle should be closed
|
|
default:
|
|
t.Fatal("空闲连接应在 Close 时被关闭")
|
|
}
|
|
|
|
select {
|
|
case <-leased.closedCh:
|
|
t.Fatal("已租赁连接不应在 Close 时被关闭")
|
|
default:
|
|
}
|
|
|
|
leased.release()
|
|
pool.Close()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_RunBackgroundPingSweep_ConcurrencyLimit(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
accountID := int64(505)
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
|
|
var current atomic.Int32
|
|
var maxConcurrent atomic.Int32
|
|
release := make(chan struct{})
|
|
for i := 0; i < 25; i++ {
|
|
conn := newOpenAIWSConn(pool.nextConnID(accountID), accountID, &openAIWSPingBlockingConn{
|
|
current: ¤t,
|
|
maxConcurrent: &maxConcurrent,
|
|
release: release,
|
|
}, nil)
|
|
ap.mu.Lock()
|
|
ap.conns[conn.id] = conn
|
|
ap.mu.Unlock()
|
|
}
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
pool.runBackgroundPingSweep()
|
|
close(done)
|
|
}()
|
|
|
|
require.Eventually(t, func() bool {
|
|
return maxConcurrent.Load() >= 10
|
|
}, time.Second, 10*time.Millisecond)
|
|
|
|
close(release)
|
|
select {
|
|
case <-done:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("runBackgroundPingSweep 未在释放后完成")
|
|
}
|
|
|
|
require.LessOrEqual(t, maxConcurrent.Load(), int32(10))
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_BasicGetterBranches(t *testing.T) {
|
|
var nilLease *openAIWSConnLease
|
|
require.Equal(t, "", nilLease.ConnID())
|
|
require.Equal(t, time.Duration(0), nilLease.QueueWaitDuration())
|
|
require.Equal(t, time.Duration(0), nilLease.ConnPickDuration())
|
|
require.False(t, nilLease.Reused())
|
|
require.Equal(t, "", nilLease.HandshakeHeader("x-test"))
|
|
require.False(t, nilLease.IsPrewarmed())
|
|
nilLease.MarkPrewarmed()
|
|
nilLease.Release()
|
|
|
|
conn := newOpenAIWSConn("getter_conn", 1, &openAIWSFakeConn{}, http.Header{"X-Test": []string{"ok"}})
|
|
lease := &openAIWSConnLease{
|
|
conn: conn,
|
|
queueWait: 3 * time.Millisecond,
|
|
connPick: 4 * time.Millisecond,
|
|
reused: true,
|
|
}
|
|
require.Equal(t, "getter_conn", lease.ConnID())
|
|
require.Equal(t, 3*time.Millisecond, lease.QueueWaitDuration())
|
|
require.Equal(t, 4*time.Millisecond, lease.ConnPickDuration())
|
|
require.True(t, lease.Reused())
|
|
require.Equal(t, "ok", lease.HandshakeHeader("x-test"))
|
|
require.False(t, lease.IsPrewarmed())
|
|
lease.MarkPrewarmed()
|
|
require.True(t, lease.IsPrewarmed())
|
|
lease.Release()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_UtilityBranches(t *testing.T) {
|
|
var nilPool *openAIWSConnPool
|
|
require.Equal(t, OpenAIWSPoolMetricsSnapshot{}, nilPool.SnapshotMetrics())
|
|
require.Equal(t, OpenAIWSTransportMetricsSnapshot{}, nilPool.SnapshotTransportMetrics())
|
|
|
|
pool := &openAIWSConnPool{cfg: &config.Config{}}
|
|
pool.metrics.acquireTotal.Store(7)
|
|
pool.metrics.acquireReuseTotal.Store(3)
|
|
metrics := pool.SnapshotMetrics()
|
|
require.Equal(t, int64(7), metrics.AcquireTotal)
|
|
require.Equal(t, int64(3), metrics.AcquireReuseTotal)
|
|
|
|
// 非 transport metrics dialer 路径
|
|
pool.clientDialer = &openAIWSFakeDialer{}
|
|
require.Equal(t, OpenAIWSTransportMetricsSnapshot{}, pool.SnapshotTransportMetrics())
|
|
pool.setClientDialerForTest(nil)
|
|
require.NotNil(t, pool.clientDialer)
|
|
|
|
require.Equal(t, 8, nilPool.maxConnsHardCap())
|
|
require.False(t, nilPool.dynamicMaxConnsEnabled())
|
|
require.Equal(t, 1.0, nilPool.maxConnsFactorByAccount(nil))
|
|
require.Equal(t, 0, nilPool.minIdlePerAccount())
|
|
require.Equal(t, 4, nilPool.maxIdlePerAccount())
|
|
require.Equal(t, 256, nilPool.queueLimitPerConn())
|
|
require.Equal(t, 0.7, nilPool.targetUtilization())
|
|
require.Equal(t, time.Duration(0), nilPool.prewarmCooldown())
|
|
require.Equal(t, 10*time.Second, nilPool.dialTimeout())
|
|
|
|
// shouldSuppressPrewarmLocked 覆盖 3 条分支
|
|
now := time.Now()
|
|
apNilFail := &openAIWSAccountPool{prewarmFails: 1}
|
|
require.False(t, pool.shouldSuppressPrewarmLocked(apNilFail, now))
|
|
apZeroTime := &openAIWSAccountPool{prewarmFails: 2}
|
|
require.False(t, pool.shouldSuppressPrewarmLocked(apZeroTime, now))
|
|
require.Equal(t, 0, apZeroTime.prewarmFails)
|
|
apOldFail := &openAIWSAccountPool{prewarmFails: 2, prewarmFailAt: now.Add(-openAIWSPrewarmFailureWindow - time.Second)}
|
|
require.False(t, pool.shouldSuppressPrewarmLocked(apOldFail, now))
|
|
apRecentFail := &openAIWSAccountPool{prewarmFails: openAIWSPrewarmFailureSuppress, prewarmFailAt: now}
|
|
require.True(t, pool.shouldSuppressPrewarmLocked(apRecentFail, now))
|
|
|
|
// recordConnPickDuration 的保护分支
|
|
nilPool.recordConnPickDuration(10 * time.Millisecond)
|
|
pool.recordConnPickDuration(-10 * time.Millisecond)
|
|
require.Equal(t, int64(1), pool.metrics.connPickTotal.Load())
|
|
|
|
// account pool 读写分支
|
|
require.Nil(t, nilPool.getOrCreateAccountPool(1))
|
|
require.Nil(t, pool.getOrCreateAccountPool(0))
|
|
pool.accounts.Store(int64(7), "invalid")
|
|
ap := pool.getOrCreateAccountPool(7)
|
|
require.NotNil(t, ap)
|
|
_, ok := pool.getAccountPool(0)
|
|
require.False(t, ok)
|
|
_, ok = pool.getAccountPool(12345)
|
|
require.False(t, ok)
|
|
pool.accounts.Store(int64(8), "bad-type")
|
|
_, ok = pool.getAccountPool(8)
|
|
require.False(t, ok)
|
|
|
|
// health check 条件
|
|
require.False(t, pool.shouldHealthCheckConn(nil))
|
|
conn := newOpenAIWSConn("health", 1, &openAIWSFakeConn{}, nil)
|
|
conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano())
|
|
require.True(t, pool.shouldHealthCheckConn(conn))
|
|
unsafeConn := newOpenAIWSConn("unsafe_health", 1, &openAIWSIdlePingUnsupportedConn{}, nil)
|
|
unsafeConn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano())
|
|
require.False(t, pool.shouldHealthCheckConn(unsafeConn))
|
|
}
|
|
|
|
func TestOpenAIWSConn_LeaseAndTimeHelpers_NilAndClosedBranches(t *testing.T) {
|
|
var nilConn *openAIWSConn
|
|
nilConn.touch()
|
|
require.Equal(t, time.Time{}, nilConn.createdAt())
|
|
require.Equal(t, time.Time{}, nilConn.lastUsedAt())
|
|
require.Equal(t, time.Duration(0), nilConn.idleDuration(time.Now()))
|
|
require.Equal(t, time.Duration(0), nilConn.age(time.Now()))
|
|
require.False(t, nilConn.isLeased())
|
|
require.False(t, nilConn.isPrewarmed())
|
|
nilConn.markPrewarmed()
|
|
|
|
conn := newOpenAIWSConn("lease_state", 1, &openAIWSFakeConn{}, nil)
|
|
require.True(t, conn.tryAcquire())
|
|
require.True(t, conn.isLeased())
|
|
conn.release()
|
|
require.False(t, conn.isLeased())
|
|
conn.close()
|
|
require.False(t, conn.tryAcquire())
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
err := conn.acquire(ctx)
|
|
require.Error(t, err)
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_ReadWriteNilConnBranches(t *testing.T) {
|
|
lease := &openAIWSConnLease{}
|
|
require.ErrorIs(t, lease.WriteJSON(map[string]any{"k": "v"}, time.Second), errOpenAIWSConnClosed)
|
|
require.ErrorIs(t, lease.WriteJSONContext(context.Background(), map[string]any{"k": "v"}), errOpenAIWSConnClosed)
|
|
_, err := lease.ReadMessage(10 * time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
_, err = lease.ReadMessageContext(context.Background())
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
_, err = lease.ReadMessageWithContextTimeout(context.Background(), 10*time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_ReleasedLeaseGuards(t *testing.T) {
|
|
conn := newOpenAIWSConn("released_guard", 1, &openAIWSFakeConn{}, nil)
|
|
lease := &openAIWSConnLease{conn: conn}
|
|
|
|
require.NoError(t, lease.PingWithTimeout(50*time.Millisecond))
|
|
|
|
lease.Release()
|
|
lease.Release() // idempotent
|
|
|
|
require.ErrorIs(t, lease.WriteJSON(map[string]any{"k": "v"}, time.Second), errOpenAIWSConnClosed)
|
|
require.ErrorIs(t, lease.WriteJSONContext(context.Background(), map[string]any{"k": "v"}), errOpenAIWSConnClosed)
|
|
require.ErrorIs(t, lease.WriteJSONWithContextTimeout(context.Background(), map[string]any{"k": "v"}, time.Second), errOpenAIWSConnClosed)
|
|
|
|
_, err := lease.ReadMessage(10 * time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
_, err = lease.ReadMessageContext(context.Background())
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
_, err = lease.ReadMessageWithContextTimeout(context.Background(), 10*time.Millisecond)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
|
|
require.ErrorIs(t, lease.PingWithTimeout(50*time.Millisecond), errOpenAIWSConnClosed)
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_MarkBrokenAfterRelease_NoEviction(t *testing.T) {
|
|
conn := newOpenAIWSConn("released_markbroken", 7, &openAIWSFakeConn{}, nil)
|
|
ap := &openAIWSAccountPool{
|
|
conns: map[string]*openAIWSConn{
|
|
conn.id: conn,
|
|
},
|
|
}
|
|
pool := &openAIWSConnPool{}
|
|
pool.accounts.Store(int64(7), ap)
|
|
|
|
lease := &openAIWSConnLease{
|
|
pool: pool,
|
|
accountID: 7,
|
|
conn: conn,
|
|
}
|
|
|
|
lease.Release()
|
|
lease.MarkBroken()
|
|
|
|
ap.mu.Lock()
|
|
_, exists := ap.conns[conn.id]
|
|
ap.mu.Unlock()
|
|
require.True(t, exists, "released lease should not evict active pool connection")
|
|
}
|
|
|
|
func TestOpenAIWSConn_AdditionalGuardBranches(t *testing.T) {
|
|
var nilConn *openAIWSConn
|
|
require.False(t, nilConn.tryAcquire())
|
|
require.ErrorIs(t, nilConn.acquire(context.Background()), errOpenAIWSConnClosed)
|
|
nilConn.release()
|
|
nilConn.close()
|
|
require.Equal(t, "", nilConn.handshakeHeader("x-test"))
|
|
|
|
connBusy := newOpenAIWSConn("busy_ctx", 1, &openAIWSFakeConn{}, nil)
|
|
require.True(t, connBusy.tryAcquire())
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
require.ErrorIs(t, connBusy.acquire(ctx), context.Canceled)
|
|
connBusy.release()
|
|
|
|
connClosed := newOpenAIWSConn("closed_guard", 1, &openAIWSFakeConn{}, nil)
|
|
connClosed.close()
|
|
require.ErrorIs(
|
|
t,
|
|
connClosed.writeJSONWithTimeout(context.Background(), map[string]any{"k": "v"}, time.Second),
|
|
errOpenAIWSConnClosed,
|
|
)
|
|
_, err := connClosed.readMessageWithContextTimeout(context.Background(), time.Second)
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
require.ErrorIs(t, connClosed.pingWithTimeout(time.Second), errOpenAIWSConnClosed)
|
|
|
|
connNoWS := newOpenAIWSConn("no_ws", 1, nil, nil)
|
|
require.ErrorIs(t, connNoWS.writeJSON(map[string]any{"k": "v"}, context.Background()), errOpenAIWSConnClosed)
|
|
_, err = connNoWS.readMessage(context.Background())
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
require.ErrorIs(t, connNoWS.pingWithTimeout(time.Second), errOpenAIWSConnClosed)
|
|
require.Equal(t, "", connNoWS.handshakeHeader("x-test"))
|
|
|
|
connOK := newOpenAIWSConn("ok", 1, &openAIWSFakeConn{}, nil)
|
|
require.NoError(t, connOK.writeJSON(map[string]any{"k": "v"}, nil))
|
|
_, err = connOK.readMessageWithContextTimeout(context.Background(), 0)
|
|
require.NoError(t, err)
|
|
require.NoError(t, connOK.pingWithTimeout(0))
|
|
|
|
connZero := newOpenAIWSConn("zero_ts", 1, &openAIWSFakeConn{}, nil)
|
|
connZero.createdAtNano.Store(0)
|
|
connZero.lastUsedNano.Store(0)
|
|
require.True(t, connZero.createdAt().IsZero())
|
|
require.True(t, connZero.lastUsedAt().IsZero())
|
|
require.Equal(t, time.Duration(0), connZero.idleDuration(time.Now()))
|
|
require.Equal(t, time.Duration(0), connZero.age(time.Now()))
|
|
|
|
require.Nil(t, cloneOpenAIWSAcquireRequestPtr(nil))
|
|
copied := cloneHeader(http.Header{
|
|
"X-Empty": []string{},
|
|
"X-Test": []string{"v1"},
|
|
})
|
|
require.Contains(t, copied, "X-Empty")
|
|
require.Nil(t, copied["X-Empty"])
|
|
require.Equal(t, "v1", copied.Get("X-Test"))
|
|
|
|
closeOpenAIWSConns([]*openAIWSConn{nil, connOK})
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_CanceledWaiterReturnsDeliveredLease(t *testing.T) {
|
|
conn := newOpenAIWSConn("cancelled_delivery", 1, &openAIWSFakeConn{}, nil)
|
|
|
|
// Both branches of acquire's select are ready. Before the post-delivery
|
|
// cancellation check this intermittently returned nil after consuming the
|
|
// only lease token, which made the next pool acquire block forever.
|
|
for range 64 {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
require.ErrorIs(t, conn.acquire(ctx), context.Canceled)
|
|
require.True(t, conn.tryAcquire(), "a canceled waiter must return a delivered lease token")
|
|
conn.release()
|
|
}
|
|
}
|
|
|
|
func TestOpenAIWSConnLease_MarkBrokenEvictsConn(t *testing.T) {
|
|
pool := newOpenAIWSConnPool(&config.Config{})
|
|
accountID := int64(5001)
|
|
conn := newOpenAIWSConn("broken_me", accountID, &openAIWSFakeConn{}, nil)
|
|
ap := pool.getOrCreateAccountPool(accountID)
|
|
ap.mu.Lock()
|
|
ap.conns[conn.id] = conn
|
|
ap.mu.Unlock()
|
|
|
|
lease := &openAIWSConnLease{
|
|
pool: pool,
|
|
accountID: accountID,
|
|
conn: conn,
|
|
}
|
|
lease.MarkBroken()
|
|
|
|
ap.mu.Lock()
|
|
_, exists := ap.conns[conn.id]
|
|
ap.mu.Unlock()
|
|
require.False(t, exists)
|
|
require.False(t, conn.tryAcquire(), "被标记为 broken 的连接应被关闭")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_TargetConnCountAndPrewarmBranches(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
require.Equal(t, 0, pool.targetConnCountLocked(nil, 1))
|
|
ap := &openAIWSAccountPool{conns: map[string]*openAIWSConn{}}
|
|
require.Equal(t, 0, pool.targetConnCountLocked(ap, 0))
|
|
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 3
|
|
require.Equal(t, 1, pool.targetConnCountLocked(ap, 1), "minIdle 应被 maxConns 截断")
|
|
|
|
// 覆盖 waiters>0 且 target 需要至少 len(conns)+1 的分支
|
|
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
|
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.9
|
|
busy := newOpenAIWSConn("busy_target", 2, &openAIWSFakeConn{}, nil)
|
|
require.True(t, busy.tryAcquire())
|
|
busy.waiters.Store(1)
|
|
ap.conns[busy.id] = busy
|
|
target := pool.targetConnCountLocked(ap, 4)
|
|
require.GreaterOrEqual(t, target, len(ap.conns)+1)
|
|
|
|
// prewarm: account pool 缺失时，拨号后的连接应被关闭并提前返回
|
|
req := openAIWSAcquireRequest{
|
|
Account: &Account{ID: 999, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
|
|
WSURL: "wss://example.com/v1/responses",
|
|
}
|
|
pool.prewarmConns(999, req, 1)
|
|
|
|
// prewarm: 拨号失败分支（prewarmFails 累加）
|
|
accountID := int64(1000)
|
|
failPool := newOpenAIWSConnPool(cfg)
|
|
failPool.setClientDialerForTest(&openAIWSAlwaysFailDialer{})
|
|
apFail := failPool.getOrCreateAccountPool(accountID)
|
|
apFail.mu.Lock()
|
|
apFail.creating = 1
|
|
apFail.mu.Unlock()
|
|
req.Account.ID = accountID
|
|
failPool.prewarmConns(accountID, req, 1)
|
|
apFail.mu.Lock()
|
|
require.GreaterOrEqual(t, apFail.prewarmFails, 1)
|
|
apFail.mu.Unlock()
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_Acquire_ErrorBranches(t *testing.T) {
|
|
var nilPool *openAIWSConnPool
|
|
_, err := nilPool.Acquire(context.Background(), openAIWSAcquireRequest{})
|
|
require.Error(t, err)
|
|
|
|
pool := newOpenAIWSConnPool(&config.Config{})
|
|
_, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: &Account{ID: 1},
|
|
WSURL: " ",
|
|
})
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "ws url is empty")
|
|
|
|
// target=nil 分支：池满且仅有 nil 连接
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
|
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 1
|
|
fullPool := newOpenAIWSConnPool(cfg)
|
|
account := &Account{ID: 2001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap := fullPool.getOrCreateAccountPool(account.ID)
|
|
ap.mu.Lock()
|
|
ap.conns["nil"] = nil
|
|
ap.lastCleanupAt = time.Now()
|
|
ap.mu.Unlock()
|
|
_, err = fullPool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSConnClosed)
|
|
|
|
// queue full 分支：waiters 达上限
|
|
account2 := &Account{ID: 2002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
ap2 := fullPool.getOrCreateAccountPool(account2.ID)
|
|
conn := newOpenAIWSConn("queue_full", account2.ID, &openAIWSFakeConn{}, nil)
|
|
require.True(t, conn.tryAcquire())
|
|
conn.waiters.Store(1)
|
|
ap2.mu.Lock()
|
|
ap2.conns[conn.id] = conn
|
|
ap2.lastCleanupAt = time.Now()
|
|
ap2.mu.Unlock()
|
|
_, err = fullPool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account2,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.ErrorIs(t, err, errOpenAIWSConnQueueFull)
|
|
}
|
|
|
|
type openAIWSFakeDialer struct{}
|
|
|
|
func (d *openAIWSFakeDialer) Dial(
|
|
ctx context.Context,
|
|
wsURL string,
|
|
headers http.Header,
|
|
proxyURL string,
|
|
) (openAIWSClientConn, int, http.Header, error) {
|
|
_ = ctx
|
|
_ = wsURL
|
|
_ = headers
|
|
_ = proxyURL
|
|
return &openAIWSFakeConn{}, 0, nil, nil
|
|
}
|
|
|
|
type openAIWSCountingDialer struct {
|
|
mu sync.Mutex
|
|
dialCount int
|
|
}
|
|
|
|
type openAIWSFirstDialBlockingCaptureDialer struct {
|
|
mu sync.Mutex
|
|
dialCount int
|
|
headers []http.Header
|
|
firstStarted chan struct{}
|
|
releaseFirst chan struct{}
|
|
}
|
|
|
|
func newOpenAIWSFirstDialBlockingCaptureDialer() *openAIWSFirstDialBlockingCaptureDialer {
|
|
return &openAIWSFirstDialBlockingCaptureDialer{
|
|
firstStarted: make(chan struct{}),
|
|
releaseFirst: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
type openAIWSAlwaysFailDialer struct {
|
|
mu sync.Mutex
|
|
dialCount int
|
|
}
|
|
|
|
type openAIWSPingBlockingConn struct {
|
|
current *atomic.Int32
|
|
maxConcurrent *atomic.Int32
|
|
release <-chan struct{}
|
|
}
|
|
|
|
type openAIWSIdlePingUnsupportedConn struct {
|
|
openAIWSFakeConn
|
|
}
|
|
|
|
func (c *openAIWSIdlePingUnsupportedConn) SupportsIdlePingWithoutReader() bool {
|
|
return false
|
|
}
|
|
|
|
func (c *openAIWSPingBlockingConn) WriteJSON(context.Context, any) error {
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSPingBlockingConn) ReadMessage(context.Context) ([]byte, error) {
|
|
return []byte(`{"type":"response.completed","response":{"id":"resp_blocking_ping"}}`), nil
|
|
}
|
|
|
|
func (c *openAIWSPingBlockingConn) Ping(ctx context.Context) error {
|
|
if c.current == nil || c.maxConcurrent == nil {
|
|
return nil
|
|
}
|
|
|
|
now := c.current.Add(1)
|
|
for {
|
|
prev := c.maxConcurrent.Load()
|
|
if now <= prev || c.maxConcurrent.CompareAndSwap(prev, now) {
|
|
break
|
|
}
|
|
}
|
|
defer c.current.Add(-1)
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-c.release:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (c *openAIWSPingBlockingConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func (d *openAIWSCountingDialer) Dial(
|
|
ctx context.Context,
|
|
wsURL string,
|
|
headers http.Header,
|
|
proxyURL string,
|
|
) (openAIWSClientConn, int, http.Header, error) {
|
|
_ = ctx
|
|
_ = wsURL
|
|
_ = headers
|
|
_ = proxyURL
|
|
d.mu.Lock()
|
|
d.dialCount++
|
|
d.mu.Unlock()
|
|
return &openAIWSFakeConn{}, 0, nil, nil
|
|
}
|
|
|
|
func (d *openAIWSFirstDialBlockingCaptureDialer) Dial(
|
|
ctx context.Context,
|
|
wsURL string,
|
|
headers http.Header,
|
|
proxyURL string,
|
|
) (openAIWSClientConn, int, http.Header, error) {
|
|
_ = wsURL
|
|
_ = proxyURL
|
|
d.mu.Lock()
|
|
d.dialCount++
|
|
dialNumber := d.dialCount
|
|
d.headers = append(d.headers, cloneHeader(headers))
|
|
d.mu.Unlock()
|
|
if dialNumber == 1 {
|
|
close(d.firstStarted)
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, 0, nil, ctx.Err()
|
|
case <-d.releaseFirst:
|
|
}
|
|
}
|
|
return &openAIWSFakeConn{}, 0, nil, nil
|
|
}
|
|
|
|
func (d *openAIWSFirstDialBlockingCaptureDialer) DialCount() int {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
return d.dialCount
|
|
}
|
|
|
|
func (d *openAIWSCountingDialer) DialCount() int {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
return d.dialCount
|
|
}
|
|
|
|
func (d *openAIWSAlwaysFailDialer) Dial(
|
|
ctx context.Context,
|
|
wsURL string,
|
|
headers http.Header,
|
|
proxyURL string,
|
|
) (openAIWSClientConn, int, http.Header, error) {
|
|
_ = ctx
|
|
_ = wsURL
|
|
_ = headers
|
|
_ = proxyURL
|
|
d.mu.Lock()
|
|
d.dialCount++
|
|
d.mu.Unlock()
|
|
return nil, 503, nil, errors.New("dial failed")
|
|
}
|
|
|
|
func (d *openAIWSAlwaysFailDialer) DialCount() int {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
return d.dialCount
|
|
}
|
|
|
|
type openAIWSFakeConn struct {
|
|
mu sync.Mutex
|
|
closed bool
|
|
payload [][]byte
|
|
}
|
|
|
|
func (c *openAIWSFakeConn) WriteJSON(ctx context.Context, value any) error {
|
|
_ = ctx
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if c.closed {
|
|
return errors.New("closed")
|
|
}
|
|
c.payload = append(c.payload, []byte("ok"))
|
|
_ = value
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSFakeConn) ReadMessage(ctx context.Context) ([]byte, error) {
|
|
_ = ctx
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if c.closed {
|
|
return nil, errors.New("closed")
|
|
}
|
|
return []byte(`{"type":"response.completed","response":{"id":"resp_fake"}}`), nil
|
|
}
|
|
|
|
func (c *openAIWSFakeConn) Ping(ctx context.Context) error {
|
|
_ = ctx
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSFakeConn) Close() error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.closed = true
|
|
return nil
|
|
}
|
|
|
|
type openAIWSBlockingConn struct {
|
|
readDelay time.Duration
|
|
}
|
|
|
|
func (c *openAIWSBlockingConn) WriteJSON(ctx context.Context, value any) error {
|
|
_ = ctx
|
|
_ = value
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSBlockingConn) ReadMessage(ctx context.Context) ([]byte, error) {
|
|
delay := c.readDelay
|
|
if delay <= 0 {
|
|
delay = 10 * time.Millisecond
|
|
}
|
|
timer := time.NewTimer(delay)
|
|
defer timer.Stop()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-timer.C:
|
|
return []byte(`{"type":"response.completed","response":{"id":"resp_blocking"}}`), nil
|
|
}
|
|
}
|
|
|
|
func (c *openAIWSBlockingConn) Ping(ctx context.Context) error {
|
|
_ = ctx
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSBlockingConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
type openAIWSWriteBlockingConn struct{}
|
|
|
|
func (c *openAIWSWriteBlockingConn) WriteJSON(ctx context.Context, _ any) error {
|
|
<-ctx.Done()
|
|
return ctx.Err()
|
|
}
|
|
|
|
func (c *openAIWSWriteBlockingConn) ReadMessage(context.Context) ([]byte, error) {
|
|
return []byte(`{"type":"response.completed","response":{"id":"resp_write_block"}}`), nil
|
|
}
|
|
|
|
func (c *openAIWSWriteBlockingConn) Ping(context.Context) error {
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSWriteBlockingConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
type openAIWSPingFailConn struct{}
|
|
|
|
func (c *openAIWSPingFailConn) WriteJSON(context.Context, any) error {
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSPingFailConn) ReadMessage(context.Context) ([]byte, error) {
|
|
return []byte(`{"type":"response.completed","response":{"id":"resp_ping_fail"}}`), nil
|
|
}
|
|
|
|
func (c *openAIWSPingFailConn) Ping(context.Context) error {
|
|
return errors.New("ping failed")
|
|
}
|
|
|
|
func (c *openAIWSPingFailConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
type openAIWSContextProbeConn struct {
|
|
lastWriteCtx context.Context
|
|
}
|
|
|
|
func (c *openAIWSContextProbeConn) WriteJSON(ctx context.Context, _ any) error {
|
|
c.lastWriteCtx = ctx
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSContextProbeConn) ReadMessage(context.Context) ([]byte, error) {
|
|
return []byte(`{"type":"response.completed","response":{"id":"resp_ctx_probe"}}`), nil
|
|
}
|
|
|
|
func (c *openAIWSContextProbeConn) Ping(context.Context) error {
|
|
return nil
|
|
}
|
|
|
|
func (c *openAIWSContextProbeConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
type openAIWSNilConnDialer struct{}
|
|
|
|
func (d *openAIWSNilConnDialer) Dial(
|
|
ctx context.Context,
|
|
wsURL string,
|
|
headers http.Header,
|
|
proxyURL string,
|
|
) (openAIWSClientConn, int, http.Header, error) {
|
|
_ = ctx
|
|
_ = wsURL
|
|
_ = headers
|
|
_ = proxyURL
|
|
return nil, 200, nil, nil
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_DialConnNilConnection(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
|
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
|
|
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
pool.setClientDialerForTest(&openAIWSNilConnDialer{})
|
|
account := &Account{ID: 91, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
|
|
|
_, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
|
Account: account,
|
|
WSURL: "wss://example.com/v1/responses",
|
|
})
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "nil connection")
|
|
}
|
|
|
|
func TestOpenAIWSConnPool_SnapshotTransportMetrics(t *testing.T) {
|
|
cfg := &config.Config{}
|
|
pool := newOpenAIWSConnPool(cfg)
|
|
|
|
dialer, ok := pool.clientDialer.(*coderOpenAIWSClientDialer)
|
|
require.True(t, ok)
|
|
|
|
_, err := dialer.proxyHTTPClient("http://127.0.0.1:28080")
|
|
require.NoError(t, err)
|
|
_, err = dialer.proxyHTTPClient("http://127.0.0.1:28080")
|
|
require.NoError(t, err)
|
|
_, err = dialer.proxyHTTPClient("http://127.0.0.1:28081")
|
|
require.NoError(t, err)
|
|
|
|
snapshot := pool.SnapshotTransportMetrics()
|
|
require.Equal(t, int64(1), snapshot.ProxyClientCacheHits)
|
|
require.Equal(t, int64(2), snapshot.ProxyClientCacheMisses)
|
|
require.InDelta(t, 1.0/3.0, snapshot.TransportReuseRatio, 0.0001)
|
|
}
|