seewe/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-14 12:02:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #7049 from xlplbo/perf/openai-ws-pool-queue-wake

Browse Source 
perf(openai): WS 连接池排队等待者按账号池变更重选连接
This commit is contained in:
Wesley Liddick
2026-09-12 19:26:56 +08:00
committed by GitHub
parent 2a4f3d1a71 7cf90c79d6
commit 3873215908
2 changed files with 258 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/openai_ws_pool.go
+84 -19
View File
@@ -40,6 +40,7 @@ var (
errOpenAIWSConnClosed = errors.New("openai ws connection closed")
errOpenAIWSConnQueueFull = errors.New("openai ws connection queue full")
errOpenAIWSPreferredConnUnavailable = errors.New("openai ws preferred connection unavailable")
errOpenAIWSPoolChanged = errors.New("openai ws account pool changed")
)
type openAIWSDialError struct {
@@ -497,6 +498,32 @@ func (c *openAIWSConn) acquire(ctx context.Context) error {
}
}
// acquireOrPoolChanged 与 acquire 相同，但同时监听账号池的变更信号：
// 别的连接释放、连接被剔除或新拨号完成都会触发它，此时返回 errOpenAIWSPoolChanged，
// 调用方应放弃只等这一条连接，回到选择逻辑重新挑选。
func (c *openAIWSConn) acquireOrPoolChanged(ctx context.Context, poolChanged <-chan struct{}) error {
if c == nil {
return errOpenAIWSConnClosed
}
select {
case <-ctx.Done():
return ctx.Err()
case <-c.closedCh:
return errOpenAIWSConnClosed
case <-poolChanged:
return errOpenAIWSPoolChanged
case <-c.leaseCh:
if err := ctx.Err(); err != nil {
c.release()
return err
}
if !c.leaseTokenUsable() {
return errOpenAIWSConnClosed
}
return nil
}
}
func (c *openAIWSConn) release() {
if c == nil {
return
@@ -1071,11 +1098,32 @@ func (p *openAIWSConnPool) runBackgroundCleanupSweep(now time.Time) {
}
}
// openAIWSAcquireQueueWait 跨广播重选与递归重试累计一次获取的排队耗时，
// 由 Acquire 在统一出口写入租约与指标，任何成功路径都不会漏记。
type openAIWSAcquireQueueWait struct {
queued bool
rewoken bool
total time.Duration
}
func (p *openAIWSConnPool) Acquire(ctx context.Context, req openAIWSAcquireRequest) (*openAIWSConnLease, error) {
if p != nil {
p.metrics.acquireTotal.Add(1)
}
lease, err := p.acquire(ctx, cloneOpenAIWSAcquireRequest(req), 0)
queueWait := &openAIWSAcquireQueueWait{}
lease, err := p.acquire(ctx, cloneOpenAIWSAcquireRequest(req), 0, queueWait)
if lease != nil && queueWait.rewoken {
// 广播重选经 tryAcquire 拿令牌，不像排队分支那样在取得令牌后检查取消，
// 这里补上复查：上下文已取消就归还令牌并按取消返回。
if ctxErr := ctx.Err(); ctxErr != nil {
lease.Release()
return nil, ctxErr
}
}
if lease != nil && queueWait.total > 0 {
lease.queueWait = queueWait.total
p.metrics.acquireQueueWaitMs.Add(queueWait.total.Milliseconds())
}
if lease != nil && lease.conn != nil {
now := time.Now()
lease.idleBefore = lease.conn.idleDuration(now)
@@ -1084,13 +1132,16 @@ func (p *openAIWSConnPool) Acquire(ctx context.Context, req openAIWSAcquireReque
return lease, err
}
func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireRequest, retry int) (*openAIWSConnLease, error) {
func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireRequest, retry int, queueWait *openAIWSAcquireQueueWait) (*openAIWSConnLease, error) {
if p == nil || req.Account == nil || req.Account.ID <= 0 {
return nil, errors.New("invalid ws acquire request")
}
if stringsTrim(req.WSURL) == "" {
return nil, errors.New("ws url is empty")
}
if queueWait == nil {
queueWait = &openAIWSAcquireQueueWait{}
}
retryAcquire:
accountID := req.Account.ID
@@ -1139,7 +1190,7 @@ retryAcquire:
preferredConn.close()
p.evictConn(accountID, preferredConn.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, err
}
@@ -1181,7 +1232,7 @@ retryAcquire:
if errors.Is(err, errOpenAIWSConnClosed) {
p.evictConn(accountID, preferredConn.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
}
return nil, err
@@ -1192,7 +1243,7 @@ retryAcquire:
preferredConn.close()
p.evictConn(accountID, preferredConn.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, err
}
@@ -1225,7 +1276,7 @@ retryAcquire:
conn.close()
p.evictConn(accountID, conn.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, err
}
@@ -1254,7 +1305,7 @@ retryAcquire:
best.close()
p.evictConn(accountID, best.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, err
}
@@ -1282,7 +1333,7 @@ retryAcquire:
conn.close()
p.evictConn(accountID, conn.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, err
}
@@ -1362,7 +1413,7 @@ retryAcquire:
conn.close()
}
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, errOpenAIWSConnClosed
}
@@ -1419,20 +1470,36 @@ acquireAtCapacity:
return nil, errOpenAIWSConnQueueFull
}
target.waiters.Add(1)
// 排队时不只等这一条连接的令牌：账号池任何容量变化（别的连接释放、被剔除、
// 新拨号完成）都会唤醒等待者回到 retryAcquire 重新选择。变更通道必须在锁内取，
// 否则会漏掉解锁到开始等待之间的信号。
changedCh := ap.changeChannelLocked()
ap.mu.Unlock()
closeOpenAIWSConns(evicted)
defer target.waiters.Add(-1)
waitStart := time.Now()
p.metrics.acquireQueueWaitTotal.Add(1)
if !queueWait.queued {
queueWait.queued = true
p.metrics.acquireQueueWaitTotal.Add(1)
}
if err := target.acquire(ctx); err != nil {
if errors.Is(err, errOpenAIWSConnClosed) {
waitErr := target.acquireOrPoolChanged(ctx, changedCh)
target.waiters.Add(-1)
queueWait.total += time.Since(waitStart)
if waitErr != nil {
if errors.Is(waitErr, errOpenAIWSPoolChanged) {
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
queueWait.rewoken = true
goto retryAcquire
}
if errors.Is(waitErr, errOpenAIWSConnClosed) {
p.evictConn(accountID, target.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
}
return nil, err
return nil, waitErr
}
if p.shouldHealthCheckConn(target) {
if err := target.pingWithTimeout(openAIWSConnHealthCheckTO); err != nil {
@@ -1440,15 +1507,13 @@ acquireAtCapacity:
target.close()
p.evictConn(accountID, target.id)
if retry < 1 {
return p.acquire(ctx, req, retry+1)
return p.acquire(ctx, req, retry+1, queueWait)
}
return nil, err
}
}
queueWait := time.Since(waitStart)
p.metrics.acquireQueueWaitMs.Add(queueWait.Milliseconds())
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: target, queueWait: queueWait, connPick: connPick, reused: true}
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: target, connPick: connPick, reused: true}
p.metrics.acquireReuseTotal.Add(1)
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
p.ensureTargetIdleAsync(accountID)
backend/internal/service/openai_ws_pool_test.go
+174
View File
@@ -322,6 +322,180 @@ func TestOpenAIWSConnPool_AcquireQueueWaitMetrics(t *testing.T) {
require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1))
}
func TestOpenAIWSConnPool_AcquireAtCapacityWakesWhenAnotherConnReleases(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
pool := newOpenAIWSConnPool(cfg)
accountID := int64(993)
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
target := newOpenAIWSConn("target", accountID, &openAIWSFakeConn{}, nil)
other := newOpenAIWSConn("other", accountID, &openAIWSFakeConn{}, nil)
require.True(t, target.tryAcquire())
require.True(t, other.tryAcquire())
// other 上已有一个等待者，新来的等待者会挂到 target 上。
other.waiters.Add(1)
ap := pool.ensureAccountPoolLocked(accountID)
ap.mu.Lock()
ap.conns[target.id] = target
ap.conns[other.id] = other
ap.lastAcquire = &req
ap.mu.Unlock()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
type result struct {
lease *openAIWSConnLease
err error
}
resultCh := make(chan result, 1)
go func() {
lease, err := pool.Acquire(ctx, req)
resultCh <- result{lease: lease, err: err}
}()
require.Eventually(t, func() bool { return target.waiters.Load() == 1 }, time.Second, 5*time.Millisecond)
time.Sleep(40 * time.Millisecond)
otherLease := &openAIWSConnLease{pool: pool, accountID: accountID, conn: other}
otherLease.Release()
select {
case got := <-resultCh:
require.NoError(t, got.err)
require.NotNil(t, got.lease)
require.Equal(t, other.id, got.lease.ConnID())
require.True(t, got.lease.Reused())
require.GreaterOrEqual(t, got.lease.QueueWaitDuration(), 30*time.Millisecond, "queue wait accumulated before the wake-up must be carried into the lease")
got.lease.Release()
case <-time.After(500 * time.Millisecond):
t.Fatal("waiter queued on a busy connection must be woken when another connection is released")
}
require.Equal(t, int32(0), target.waiters.Load())
metrics := pool.SnapshotMetrics()
require.Equal(t, int64(1), metrics.AcquireQueueWaitTotal)
require.GreaterOrEqual(t, metrics.AcquireQueueWaitMsTotal, int64(30))
}
func TestOpenAIWSConnPool_AcquireAtCapacityWakesWhenCapacityFreedByEviction(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
pool := newOpenAIWSConnPool(cfg)
dialer := &openAIWSCountingDialer{}
pool.setClientDialerForTest(dialer)
accountID := int64(994)
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
target := newOpenAIWSConn("target", accountID, &openAIWSFakeConn{}, nil)
other := newOpenAIWSConn("other", accountID, &openAIWSFakeConn{}, nil)
require.True(t, target.tryAcquire())
require.True(t, other.tryAcquire())
other.waiters.Add(1)
ap := pool.ensureAccountPoolLocked(accountID)
ap.mu.Lock()
ap.conns[target.id] = target
ap.conns[other.id] = other
ap.lastAcquire = &req
ap.mu.Unlock()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
type result struct {
lease *openAIWSConnLease
err error
}
resultCh := make(chan result, 1)
go func() {
lease, err := pool.Acquire(ctx, req)
resultCh <- result{lease: lease, err: err}
}()
require.Eventually(t, func() bool { return target.waiters.Load() == 1 }, time.Second, 5*time.Millisecond)
// 剔除另一条连接腾出名额，等待者应重新选择并新拨号，而不是继续等 target。
time.Sleep(40 * time.Millisecond)
pool.evictConn(accountID, other.id)
select {
case got := <-resultCh:
require.NoError(t, got.err)
require.NotNil(t, got.lease)
require.False(t, got.lease.Reused())
require.NotEqual(t, target.id, got.lease.ConnID())
require.GreaterOrEqual(t, got.lease.QueueWaitDuration(), 30*time.Millisecond, "queue wait must be carried into a lease obtained by dialing after the wake-up")
got.lease.Release()
case <-time.After(500 * time.Millisecond):
t.Fatal("waiter queued on a busy connection must be woken when pool capacity is freed")
}
require.Equal(t, 1, dialer.DialCount())
require.Equal(t, int32(0), target.waiters.Load())
metrics := pool.SnapshotMetrics()
require.Equal(t, int64(1), metrics.AcquireQueueWaitTotal)
require.GreaterOrEqual(t, metrics.AcquireQueueWaitMsTotal, int64(30))
}
func TestOpenAIWSConnPool_AcquireAtCapacityCanceledWaiterDoesNotTakeReleasedConn(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
accountID := int64(995)
account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
pool := newOpenAIWSConnPool(cfg)
target := newOpenAIWSConn("target", accountID, &openAIWSFakeConn{}, nil)
other := newOpenAIWSConn("other", accountID, &openAIWSFakeConn{}, nil)
require.True(t, target.tryAcquire())
require.True(t, other.tryAcquire())
other.waiters.Add(1)
ap := pool.ensureAccountPoolLocked(accountID)
ap.mu.Lock()
ap.conns[target.id] = target
ap.conns[other.id] = other
ap.lastAcquire = &req
ap.mu.Unlock()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
type result struct {
lease *openAIWSConnLease
err error
}
resultCh := make(chan result, 1)
go func() {
lease, err := pool.Acquire(ctx, req)
resultCh <- result{lease: lease, err: err}
}()
require.Eventually(t, func() bool { return target.waiters.Load() == 1 }, time.Second, 5*time.Millisecond)
// 持锁广播：等待者被唤醒后卡在重新取锁上，此时取消请求并归还 other 的令牌，
// 解锁后重选会看到一条空闲连接，但请求已经取消，不能带着租约返回。
ap.mu.Lock()
ap.signalChangedLocked()
cancel()
other.release()
ap.mu.Unlock()
select {
case got := <-resultCh:
require.ErrorIs(t, got.err, context.Canceled)
require.Nil(t, got.lease, "a canceled waiter must not come back with a lease after a pool change wake-up")
case <-time.After(500 * time.Millisecond):
t.Fatal("canceled waiter must return promptly")
}
require.True(t, other.tryAcquire(), "the released token must stay available to other acquirers")
other.release()
require.Equal(t, int32(0), target.waiters.Load())
}
func TestOpenAIWSConnPool_DialSuccessWakesTopologyWaiterAndCanceledWaiterDoesNotLoseLease(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1