fix(openai-ws): recycle stale idle pool connections

This commit is contained in:
Ray
2026-08-31 21:04:29 +08:00
parent 52374af940
commit 6d5f027848
2 changed files with 42 additions and 2 deletions
+16 -2
View File
@@ -18,8 +18,11 @@ import (
)
const (
openAIWSConnMaxAge = 60 * time.Minute
openAIWSConnHealthCheckIdle = 90 * time.Second
openAIWSConnMaxAge = 60 * time.Minute
openAIWSConnHealthCheckIdle = 90 * time.Second
// coder/websocket cannot consume pong frames without a reader. Recycle
// unsupported idle sockets before the upstream keepalive window expires.
openAIWSConnIdleRecycleAfter = 90 * time.Second
openAIWSConnHealthCheckTO = 2 * time.Second
openAIWSConnPrewarmExtraDelay = 2 * time.Second
openAIWSAcquireCleanupInterval = 3 * time.Second
@@ -1353,6 +1356,17 @@ func (p *openAIWSConnPool) cleanupAccountLocked(ap *openAIWSAccountPool, now tim
if p.isConnPinnedLocked(ap, id) {
continue
}
if !conn.isLeased() && conn.waiters.Load() == 0 &&
!conn.supportsIdlePingWithoutReader() &&
conn.idleDuration(now) >= openAIWSConnIdleRecycleAfter {
delete(ap.conns, id)
if len(ap.pinnedConns) > 0 {
delete(ap.pinnedConns, id)
}
evicted = append(evicted, conn)
p.metrics.scaleDownTotal.Add(1)
continue
}
if maxAge > 0 && !conn.isLeased() && conn.age(now) > maxAge {
delete(ap.conns, id)
if len(ap.pinnedConns) > 0 {
@@ -1307,6 +1307,32 @@ func TestOpenAIWSConnPool_BackgroundCleanupSweep_WithoutAcquire(t *testing.T) {
require.False(t, exists, "后台清理应在无新 acquire 时也回收过期连接")
}
func TestOpenAIWSConnPool_RecyclesUnsupportedIdlePingConnection(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
pool := &openAIWSConnPool{cfg: cfg}
accountID := int64(303)
ap := &openAIWSAccountPool{conns: make(map[string]*openAIWSConn)}
conn := newOpenAIWSConn("stale_unsupported_idle_ping", accountID, &openAIWSIdlePingUnsupportedConn{}, nil)
conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnIdleRecycleAfter - time.Second).UnixNano())
ap.conns[conn.id] = conn
pool.accounts.Store(accountID, ap)
pool.runBackgroundCleanupSweep(time.Now())
ap.mu.Lock()
_, exists := ap.conns[conn.id]
ap.mu.Unlock()
require.False(t, exists, "不支持无 reader idle ping 的陈旧连接应被主动回收")
select {
case <-conn.closedCh:
default:
t.Fatal("被回收的陈旧连接应已关闭")
}
}
func TestOpenAIWSConnPool_BackgroundWorkerGuardBranches(t *testing.T) {
var nilPool *openAIWSConnPool
require.NotPanics(t, func() {