fix(wsrelay): reserve pending requests before reuse

Author: DeliciousBuding

proxy/wsrelay/executor.go

@@ -89,26 +89,22 @@ func (e *Executor) ExecuteRequestViaWebsocket(
 	headers := e.prepareWebsocketHeaders(accessToken, accountIDStr, sessionID, apiKey, deviceCfg, ginHeaders)
 
 	// 获取或创建连接
-	wc, err := e.manager.AcquireConnection(ctx, account, wsURL, sessionID, headers, proxyOverride)
+	wc, pr, err := e.manager.AcquireConnection(ctx, account, wsURL, sessionID, headers, proxyOverride)
 	if err != nil {
 		return nil, err
 	}
 
-	// 创建请求
-	pr := wc.session.AddPendingRequest(sessionID)
-
 	// 发送请求
 	if err := e.sendRequest(wc, wsBody, pr.RequestID); err != nil {
 		// 发送失败，尝试重连一次
 		wc.session.RemovePendingRequest(pr.RequestID)
 		e.manager.RemoveConnection(account.ID(), wsURL, sessionID, proxyOverride)
 
-		wc, err = e.manager.AcquireConnection(ctx, account, wsURL, sessionID, headers, proxyOverride)
+		wc, pr, err = e.manager.AcquireConnection(ctx, account, wsURL, sessionID, headers, proxyOverride)
 		if err != nil {
 			return nil, err
 		}
 
-		pr = wc.session.AddPendingRequest(sessionID)
 		if err := e.sendRequest(wc, wsBody, pr.RequestID); err != nil {
 			wc.session.RemovePendingRequest(pr.RequestID)
 			e.manager.ReleaseConnection(wc)

proxy/wsrelay/manager.go

@@ -175,6 +175,9 @@ type Manager struct {
 
 	// 读写锁保护回调设置
 	mu sync.RWMutex
+
+	// pool key 级别串行化，避免同一逻辑 session 在 acquire 阶段竞争同一条连接
+	keyLocks sync.Map
 }
 
 // NewManager 创建连接池管理器
@@ -283,6 +286,17 @@ func (m *Manager) getOnConnected() func(accountID int64, session *Session) {
 	return m.onConnected
 }
 
+func (m *Manager) keyLock(key string) *sync.Mutex {
+	if v, ok := m.keyLocks.Load(key); ok {
+		return v.(*sync.Mutex)
+	}
+	mu := &sync.Mutex{}
+	if actual, loaded := m.keyLocks.LoadOrStore(key, mu); loaded {
+		return actual.(*sync.Mutex)
+	}
+	return mu
+}
+
 // AcquireConnection 获取或创建连接
 // 仅在同一逻辑 session 且连接空闲时复用，避免不同会话共用一条已握手连接。
 func (m *Manager) AcquireConnection(
@@ -292,34 +306,52 @@ func (m *Manager) AcquireConnection(
 	sessionKey string,
 	headers http.Header,
 	proxyOverride string,
-) (*WsConnection, error) {
+) (*WsConnection, *PendingRequest, error) {
 	key := m.poolKey(account.ID(), wsURL, sessionKey, effectiveProxyURL(account, proxyOverride))
+	lock := m.keyLock(key)
+	wait := 10 * time.Millisecond
 
-	if v, ok := m.connections.Load(key); ok {
-		wc := v.(*WsConnection)
-		if canReuseConnection(wc) {
-			wc.Touch()
-			return wc, nil
+	for {
+		lock.Lock()
+		if v, ok := m.connections.Load(key); ok {
+			wc := v.(*WsConnection)
+			if canReuseConnection(wc) {
+				pr := wc.session.AddPendingRequest(sessionKey)
+				wc.Touch()
+				lock.Unlock()
+				return wc, pr, nil
+			}
+			if wc.IsConnected() && !wc.IsExpired() && wc.session != nil {
+				lock.Unlock()
+				select {
+				case <-ctx.Done():
+					return nil, nil, ctx.Err()
+				case <-time.After(wait):
+				}
+				continue
+			}
+			m.connections.Delete(key)
+			m.sessions.Delete(key)
+			wc.Close()
 		}
-		m.connections.Delete(key)
-		wc.Close()
-	}
 
-	// 始终创建新连接，避免多个请求复用同一个 websocket.Conn 导致并发读取
-	wc, err := m.createConnection(ctx, account, wsURL, sessionKey, headers, proxyOverride)
-	if err != nil {
-		return nil, err
-	}
+		wc, err := m.createConnection(ctx, account, wsURL, sessionKey, headers, proxyOverride)
+		if err != nil {
+			lock.Unlock()
+			return nil, nil, err
+		}
 
-	// 存储新连接
-	m.connections.Store(key, wc)
+		// 存储新连接并立即占位 pending request，避免返回后才记账产生竞态
+		m.connections.Store(key, wc)
+		pr := wc.session.AddPendingRequest(sessionKey)
+		lock.Unlock()
 
-	// 调用连接回调
-	if fn := m.getOnConnected(); fn != nil {
-		fn(account.ID(), wc.session)
-	}
+		if fn := m.getOnConnected(); fn != nil {
+			fn(account.ID(), wc.session)
+		}
 
-	return wc, nil
+		return wc, pr, nil
+	}
 }
 
 func canReuseConnection(wc *WsConnection) bool {
@@ -459,7 +491,7 @@ func (m *Manager) ReplaceConnection(
 	sessionKey string,
 	headers http.Header,
 	proxyOverride string,
-) (*WsConnection, error) {
+) (*WsConnection, *PendingRequest, error) {
 	// 先移除旧连接
 	m.RemoveConnection(account.ID(), wsURL, sessionKey, effectiveProxyURL(account, proxyOverride))
 

proxy/wsrelay/manager_test.go

@@ -30,13 +30,20 @@ func TestAcquireConnectionReusesIdleConnectedConnection(t *testing.T) {
 	manager.connections.Store(key, conn)
 	manager.sessions.Store(key, session)
 
-	got, err := manager.AcquireConnection(context.Background(), account, wsURL, "session-1", http.Header{}, "")
+	got, pr, err := manager.AcquireConnection(context.Background(), account, wsURL, "session-1", http.Header{}, "")
 	if err != nil {
 		t.Fatalf("AcquireConnection() error = %v", err)
 	}
 	if got != conn {
 		t.Fatal("expected existing connection to be reused")
 	}
+	if pr == nil {
+		t.Fatal("expected pending request reservation")
+	}
+	if session.PendingCount() != 1 {
+		t.Fatalf("PendingCount = %d, want %d", session.PendingCount(), 1)
+	}
+	session.RemovePendingRequest(pr.RequestID)
 }
 
 func TestPoolKeyIncludesSessionKey(t *testing.T) {
@@ -115,3 +122,36 @@ func TestCanReuseConnection(t *testing.T) {
 		}
 	})
 }
+
+func TestAcquireConnectionWaitsWhileSessionHasPendingRequest(t *testing.T) {
+	manager := NewManager()
+	t.Cleanup(manager.Stop)
+
+	account := &auth.Account{DBID: 42}
+	wsURL := "wss://example.test/responses"
+	key := manager.poolKey(account.ID(), wsURL, "session-1", "")
+
+	session := NewSession(account.ID(), manager)
+	session.SetConnected(true)
+	blocking := session.AddPendingRequest("session-1")
+	t.Cleanup(func() { session.RemovePendingRequest(blocking.RequestID) })
+
+	conn := &WsConnection{
+		session:  session,
+		URL:      wsURL,
+		PoolKey:  key,
+		httpResp: &http.Response{StatusCode: http.StatusSwitchingProtocols},
+	}
+	conn.SetState(StateConnected)
+	conn.Touch()
+	manager.connections.Store(key, conn)
+	manager.sessions.Store(key, session)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+
+	_, _, err := manager.AcquireConnection(ctx, account, wsURL, "session-1", http.Header{}, "")
+	if err == nil {
+		t.Fatal("expected acquire to stop when session stays busy until context timeout")
+	}
+}