From 7cf90c79d624eb4b9ccc8279e5facf628bc046e2 Mon Sep 17 00:00:00 2001 From: xlplbo Date: Sat, 12 Sep 2026 13:01:43 +0800 Subject: [PATCH] =?UTF-8?q?perf(openai):=20WS=20=E8=BF=9E=E6=8E=A5?= =?UTF-8?q?=E6=B1=A0=E6=8E=92=E9=98=9F=E7=AD=89=E5=BE=85=E8=80=85=E6=8C=89?= =?UTF-8?q?=E8=B4=A6=E5=8F=B7=E6=B1=A0=E5=8F=98=E6=9B=B4=E9=87=8D=E9=80=89?= =?UTF-8?q?=E8=BF=9E=E6=8E=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/internal/service/openai_ws_pool.go | 103 +++++++++-- .../internal/service/openai_ws_pool_test.go | 174 ++++++++++++++++++ 2 files changed, 258 insertions(+), 19 deletions(-) diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index ed86a36284..a500484232 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -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) diff --git a/backend/internal/service/openai_ws_pool_test.go b/backend/internal/service/openai_ws_pool_test.go index 1a4480dccb..e2639d5321 100644 --- a/backend/internal/service/openai_ws_pool_test.go +++ b/backend/internal/service/openai_ws_pool_test.go @@ -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