From 617bd676f0a341c156361be714919f67daaf6d70 Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 5 Jul 2026 19:38:15 +0800 Subject: [PATCH] refactor: move websocket keepalive into pump Co-authored-by: Codex --- backend/src/upstream/transport/mod.rs | 1 + backend/src/upstream/transport/websocket.rs | 65 ++--- .../src/upstream/transport/websocket_pool.rs | 196 ++++---------- .../src/upstream/transport/websocket_pump.rs | 236 +++++++++++++++++ .../upstream/transport/websocket_pool.rs | 239 ++++++++---------- 5 files changed, 426 insertions(+), 311 deletions(-) create mode 100644 backend/src/upstream/transport/websocket_pump.rs diff --git a/backend/src/upstream/transport/mod.rs b/backend/src/upstream/transport/mod.rs index 3b97e887a..c9d35b161 100644 --- a/backend/src/upstream/transport/mod.rs +++ b/backend/src/upstream/transport/mod.rs @@ -8,6 +8,7 @@ pub mod tls; pub mod usage; pub mod websocket; pub mod websocket_pool; +pub(crate) mod websocket_pump; pub use self::{ client::{ diff --git a/backend/src/upstream/transport/websocket.rs b/backend/src/upstream/transport/websocket.rs index 3d1338a1b..8732fee1e 100644 --- a/backend/src/upstream/transport/websocket.rs +++ b/backend/src/upstream/transport/websocket.rs @@ -10,7 +10,7 @@ use std::{ use bytes::Bytes; use chrono::Utc; -use futures::{SinkExt, Stream, StreamExt}; +use futures::Stream; use thiserror::Error; use tokio::sync::mpsc; use tokio::{sync::Mutex, time::timeout}; @@ -38,9 +38,10 @@ use crate::upstream::protocol::websocket::{ use super::response_meta; use super::websocket_pool::{ - CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, CodexWsStream, + CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision, }; +use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream}; use uuid::Uuid; const REDACTED_HEADER_VALUE: &str = ""; @@ -318,7 +319,7 @@ fn websocket_host_header(endpoint: &str) -> Option { async fn connect_websocket( connection: &CodexWebSocketConnection, -) -> Result<(CodexWsStream, WsResponse>>), CodexWebSocketExchangeError> { +) -> Result<(RawWsStream, WsResponse>>), CodexWebSocketExchangeError> { let request = websocket_handshake_request(connection)?; let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca() .map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))? @@ -330,6 +331,18 @@ async fn connect_websocket( } } +/// 建立连接并交给后台 pump 托管；`keepalive` 决定该连接是否做主动保活。 +/// +/// 非池化（即用即弃）连接用 [`PumpKeepalive::disabled`]；池化连接用连接池派生的保活策略， +/// 让空闲期在后台完成 ping/pong 与失活检测，从而复用前可零成本判活。 +async fn connect_pumped_websocket( + connection: &CodexWebSocketConnection, + keepalive: PumpKeepalive, +) -> Result<(PumpedWebSocket, WsResponse>>), CodexWebSocketExchangeError> { + let (raw, response) = connect_websocket(connection).await?; + Ok((PumpedWebSocket::new(raw, keepalive), response)) +} + fn websocket_handshake_request( connection: &CodexWebSocketConnection, ) -> Result { @@ -400,7 +413,8 @@ async fn execute_fresh_response_create_request( request: &CodexWebSocketRequest, started_at: Instant, ) -> Result { - let (mut websocket, response) = connect_websocket(request.connection()).await?; + let (websocket, response) = + connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; websocket .send(Message::Text(request.payload_text().to_string().into())) .await?; @@ -414,7 +428,8 @@ async fn execute_fresh_response_create_request( async fn execute_fresh_response_create_request_stream( request: &CodexWebSocketRequest, ) -> Result { - let (mut websocket, response) = connect_websocket(request.connection()).await?; + let (websocket, response) = + connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; websocket .send(Message::Text(request.payload_text().to_string().into())) .await?; @@ -543,7 +558,8 @@ async fn execute_fresh_pooled_response_create_request( key: CodexWebSocketPoolKey, started_at: Instant, ) -> Result { - let (mut websocket, response) = connect_websocket(request.connection()).await?; + let (websocket, response) = + connect_pumped_websocket(request.connection(), pool.keepalive()).await?; websocket .send(Message::Text(request.payload_text().to_string().into())) .await?; @@ -552,15 +568,12 @@ async fn execute_fresh_pooled_response_create_request( let metadata = websocket_connection_metadata(&response); let (exchange, websocket, metadata) = collect_websocket_response(websocket, metadata, false, started_at).await?; - let last_activity_at = Instant::now(); pool.put( key, PooledWebSocketConnection { websocket, metadata, created_at: now, - last_activity_at, - last_ping_at: None, }, ) .await; @@ -572,7 +585,8 @@ async fn execute_fresh_pooled_response_create_request_stream( pool: CodexWebSocketPool, key: CodexWebSocketPoolKey, ) -> Result { - let (mut websocket, response) = connect_websocket(request.connection()).await?; + let (websocket, response) = + connect_pumped_websocket(request.connection(), pool.keepalive()).await?; websocket .send(Message::Text(request.payload_text().to_string().into())) .await?; @@ -586,7 +600,6 @@ async fn execute_fresh_pooled_response_create_request_stream( pool, key, created_at: now, - last_ping_at: None, }), None, )) @@ -596,7 +609,7 @@ async fn execute_pooled_response_create_request( request: &CodexWebSocketRequest, pool: &CodexWebSocketPool, key: CodexWebSocketPoolKey, - mut connection: PooledWebSocketConnection, + connection: PooledWebSocketConnection, started_at: Instant, ) -> Result { if let Err(error) = connection @@ -613,20 +626,16 @@ async fn execute_pooled_response_create_request( } let created_at = connection.created_at; - let last_ping_at = connection.last_ping_at; match collect_websocket_response(connection.websocket, connection.metadata, true, started_at) .await { Ok((exchange, websocket, metadata)) => { - let last_activity_at = Instant::now(); pool.put( key, PooledWebSocketConnection { websocket, metadata, created_at, - last_activity_at, - last_ping_at, }, ) .await; @@ -667,14 +676,13 @@ async fn execute_pooled_response_create_request_stream( pool, key, created_at: connection.created_at, - last_ping_at: connection.last_ping_at, }), Some(prefetched_text), )) } async fn first_reused_stream_text( - websocket: &mut CodexWsStream, + websocket: &mut PumpedWebSocket, ) -> Result { loop { let message = match next_websocket_message(websocket, WEBSOCKET_RECEIVE_IDLE_TIMEOUT).await @@ -720,14 +728,14 @@ pub async fn execute_response_create_request( // --------------------------------------------------------------------------- async fn collect_websocket_response( - mut websocket: CodexWsStream, + mut websocket: PumpedWebSocket, mut metadata: CodexWebSocketConnectionMetadata, reused_connection: bool, started_at: Instant, ) -> Result< ( CodexWebSocketExchange, - CodexWsStream, + PumpedWebSocket, CodexWebSocketConnectionMetadata, ), CodexWebSocketExchangeError, @@ -849,11 +857,10 @@ struct WebSocketStreamPoolReturn { pool: CodexWebSocketPool, key: CodexWebSocketPoolKey, created_at: Instant, - last_ping_at: Option, } fn stream_websocket_response( - websocket: CodexWsStream, + websocket: PumpedWebSocket, metadata: CodexWebSocketConnectionMetadata, pool_return: Option, prefetched_text: Option, @@ -893,7 +900,7 @@ fn stream_websocket_response( } async fn forward_websocket_response_stream( - mut websocket: CodexWsStream, + mut websocket: PumpedWebSocket, mut metadata: CodexWebSocketConnectionMetadata, pool_return: Option, prefetched_text: Option, @@ -1024,7 +1031,7 @@ fn receive_idle_timeout(saw_response_frame: bool) -> Duration { } async fn next_websocket_message( - websocket: &mut CodexWsStream, + websocket: &mut PumpedWebSocket, receive_timeout: Duration, ) -> Result, CodexWebSocketExchangeError> { match timeout(receive_timeout, websocket.next()).await { @@ -1036,12 +1043,12 @@ async fn next_websocket_message( } async fn finish_stream_websocket( - mut websocket: CodexWsStream, + mut websocket: PumpedWebSocket, metadata: CodexWebSocketConnectionMetadata, pool_return: Option, ) { let Some(pool_return) = pool_return else { - let _ = websocket.close(None).await; + websocket.close().await; return; }; pool_return @@ -1052,21 +1059,19 @@ async fn finish_stream_websocket( websocket, metadata, created_at: pool_return.created_at, - last_activity_at: Instant::now(), - last_ping_at: pool_return.last_ping_at, }, ) .await; } async fn discard_stream_websocket( - mut websocket: CodexWsStream, + mut websocket: PumpedWebSocket, pool_return: Option, ) { if let Some(pool_return) = pool_return { pool_return.pool.discard(&pool_return.key).await; } - let _ = websocket.close(None).await; + websocket.close().await; } // --------------------------------------------------------------------------- diff --git a/backend/src/upstream/transport/websocket_pool.rs b/backend/src/upstream/transport/websocket_pool.rs index bec280af8..b55208781 100644 --- a/backend/src/upstream/transport/websocket_pool.rs +++ b/backend/src/upstream/transport/websocket_pool.rs @@ -6,11 +6,11 @@ use std::{ time::{Duration, Instant}, }; -use futures::{SinkExt, StreamExt}; use serde_json::{json, Value}; use sha2::{Digest, Sha256}; -use tokio::{net::TcpStream, sync::Mutex, time::timeout}; -use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream}; +use tokio::sync::Mutex; + +use super::websocket_pump::{PumpKeepalive, PumpedWebSocket}; const DEFAULT_MAX_PER_ACCOUNT: usize = 8; const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55); @@ -105,6 +105,16 @@ impl Default for CodexWebSocketPoolConfig { } } +impl CodexWebSocketPoolConfig { + /// pump 后台任务的保活策略：从连接池配置派生出 ping 间隔与 liveness 超时。 + pub(crate) fn keepalive(&self) -> PumpKeepalive { + PumpKeepalive { + ping_interval: self.ping_interval, + liveness_timeout: self.liveness_timeout, + } + } +} + impl CodexWebSocketPool { /// 构造连接池策略和状态。 pub fn new(max_per_account: usize, max_age: Duration) -> Self { @@ -128,6 +138,11 @@ impl CodexWebSocketPool { pool } + /// pump 后台任务的保活策略（供建连时传入）。 + pub(crate) fn keepalive(&self) -> PumpKeepalive { + self.config.keepalive() + } + pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire { let mut expired_connection = None; let acquire = { @@ -136,14 +151,18 @@ impl CodexWebSocketPool { return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Disabled); } match state.slots.get(key) { - Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) => { + Some(WebSocketPoolSlot::Busy) => { return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy); } Some(WebSocketPoolSlot::Idle(_)) => { let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else { return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy); }; - if connection.created_at.elapsed() < self.config.max_age { + // 零成本探活：后台 pump 已实时感知连接死亡（RST/Close/EOF/失活）， + // 复用前只需读取 is_closed 标志，避免复用到静默死连接后卡到超时。 + if connection.created_at.elapsed() < self.config.max_age + && !connection.websocket.is_closed() + { state.slots.insert(key.clone(), WebSocketPoolSlot::Busy); return WebSocketPoolAcquire::Reused(connection); } @@ -200,7 +219,7 @@ impl CodexWebSocketPool { let mut state = self.inner.lock().await; match state.slots.remove(key) { Some(WebSocketPoolSlot::Idle(connection)) => Some(*connection), - Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) | None => None, + Some(WebSocketPoolSlot::Busy) | None => None, } }; if let Some(connection) = idle_connection { @@ -238,21 +257,21 @@ impl CodexWebSocketPool { .drain() .filter_map(|(_, slot)| match slot { WebSocketPoolSlot::Idle(connection) => Some(*connection), - WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None, + WebSocketPoolSlot::Busy => None, }) .collect::>() }; close_pooled_connections(idle_connections).await; } - /// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。 + /// 维护 idle 连接：清扫已被后台 pump 判定死亡或已超龄的连接。 + /// + /// 保活（ping/pong）与失活检测已下沉到每条连接的 pump 任务内部，此处不再做 + /// 同步 ping 探活，只负责把「后台已标记 closed」或「超过 max_age」的 idle + /// 连接从池中摘除并关闭，避免它们占用 slot。 pub async fn maintain_idle_connections(&self) { - let maintenance = self.take_idle_connections_for_maintenance().await; - close_pooled_connections(maintenance.close).await; - for (key, connection) in maintenance.probe { - let connection = probe_idle_connection(connection, self.config.ping_timeout).await; - self.return_maintained_connection(key, connection).await; - } + let close = self.take_dead_idle_connections().await; + close_pooled_connections(close).await; } fn spawn_maintenance_task(&self) { @@ -285,107 +304,31 @@ impl CodexWebSocketPool { self.inner.lock().await.shutting_down } - async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance { + async fn take_dead_idle_connections(&self) -> Vec { let mut close = Vec::new(); - let mut probe = Vec::new(); let now = Instant::now(); - let Some(ping_interval) = self.config.ping_interval else { - let mut state = self.inner.lock().await; - if state.shutting_down || !self.config.enabled { - return WebSocketPoolMaintenance { close, probe }; - } - let keys = state - .slots - .iter() - .filter_map(|(key, slot)| match slot { - WebSocketPoolSlot::Idle(connection) - if should_close_idle_connection( - connection, - now, - self.config.max_age, - self.config.liveness_timeout, - ) => - { - Some(key.clone()) - } - WebSocketPoolSlot::Idle(_) - | WebSocketPoolSlot::Busy - | WebSocketPoolSlot::Checking => None, - }) - .collect::>(); - for key in keys { - if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { - close.push(*connection); - } - } - return WebSocketPoolMaintenance { close, probe }; - }; - let mut state = self.inner.lock().await; if state.shutting_down || !self.config.enabled { - return WebSocketPoolMaintenance { close, probe }; + return close; } let keys = state .slots .iter() .filter_map(|(key, slot)| match slot { WebSocketPoolSlot::Idle(connection) - if should_close_idle_connection( - connection, - now, - self.config.max_age, - self.config.liveness_timeout, - ) || should_probe_idle_connection(connection, now, ping_interval) => + if should_close_idle_connection(connection, now, self.config.max_age) => { Some(key.clone()) } - WebSocketPoolSlot::Idle(_) - | WebSocketPoolSlot::Busy - | WebSocketPoolSlot::Checking => None, + WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy => None, }) .collect::>(); for key in keys { - let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) else { - continue; - }; - if should_close_idle_connection( - &connection, - now, - self.config.max_age, - self.config.liveness_timeout, - ) { + if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { close.push(*connection); - } else { - state.slots.insert(key.clone(), WebSocketPoolSlot::Checking); - probe.push((key, *connection)); } } - WebSocketPoolMaintenance { close, probe } - } - - async fn return_maintained_connection( - &self, - key: CodexWebSocketPoolKey, - connection: Option, - ) { - let mut connection = connection; - { - let mut state = self.inner.lock().await; - if state.shutting_down || !self.config.enabled { - state.slots.remove(&key); - } else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Checking)) { - if let Some(connection) = connection.take() { - state - .slots - .insert(key, WebSocketPoolSlot::Idle(Box::new(connection))); - } else { - state.slots.remove(&key); - } - } - } - if let Some(connection) = connection { - close_pooled_connection(connection).await; - } + close } } @@ -395,8 +338,6 @@ struct WebSocketPoolState { shutting_down: bool, } -pub(crate) type CodexWsStream = WebSocketStream>; - #[derive(Clone)] pub(crate) struct CodexWebSocketConnectionMetadata { pub(crate) turn_state: Option, @@ -405,17 +346,14 @@ pub(crate) struct CodexWebSocketConnectionMetadata { } pub(crate) struct PooledWebSocketConnection { - pub(crate) websocket: CodexWsStream, + pub(crate) websocket: PumpedWebSocket, pub(crate) metadata: CodexWebSocketConnectionMetadata, pub(crate) created_at: Instant, - pub(crate) last_activity_at: Instant, - pub(crate) last_ping_at: Option, } enum WebSocketPoolSlot { Idle(Box), Busy, - Checking, } pub(crate) enum WebSocketPoolAcquire { @@ -538,7 +476,7 @@ fn account_slot_count( } async fn close_pooled_connection(mut connection: PooledWebSocketConnection) { - let _ = connection.websocket.send(Message::Close(None)).await; + connection.websocket.close().await; } async fn close_pooled_connections(connections: Vec) { @@ -547,59 +485,11 @@ async fn close_pooled_connections(connections: Vec) { } } -struct WebSocketPoolMaintenance { - close: Vec, - probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>, -} - +/// idle 连接是否应从池中摘除：被后台 pump 标记死亡，或已超过 `max_age`。 fn should_close_idle_connection( connection: &PooledWebSocketConnection, now: Instant, max_age: Duration, - liveness_timeout: Option, ) -> bool { - now.duration_since(connection.created_at) >= max_age - || liveness_timeout.is_some_and(|timeout| { - !timeout.is_zero() && now.duration_since(connection.last_activity_at) >= timeout - }) -} - -fn should_probe_idle_connection( - connection: &PooledWebSocketConnection, - now: Instant, - ping_interval: Duration, -) -> bool { - !ping_interval.is_zero() - && connection - .last_ping_at - .is_none_or(|last_ping_at| now.duration_since(last_ping_at) >= ping_interval) -} - -async fn probe_idle_connection( - mut connection: PooledWebSocketConnection, - ping_timeout: Duration, -) -> Option { - connection.last_ping_at = Some(Instant::now()); - if connection - .websocket - .send(Message::Ping(Vec::new().into())) - .await - .is_err() - { - close_pooled_connection(connection).await; - return None; - } - if ping_timeout.is_zero() { - return Some(connection); - } - match timeout(ping_timeout, connection.websocket.next()).await { - Ok(Some(Ok(Message::Close(_)) | Err(_)) | None) | Err(_) => { - close_pooled_connection(connection).await; - None - } - Ok(Some(Ok(_message))) => { - connection.last_activity_at = Instant::now(); - Some(connection) - } - } + connection.websocket.is_closed() || now.duration_since(connection.created_at) >= max_age } diff --git a/backend/src/upstream/transport/websocket_pump.rs b/backend/src/upstream/transport/websocket_pump.rs new file mode 100644 index 000000000..e61c2a71d --- /dev/null +++ b/backend/src/upstream/transport/websocket_pump.rs @@ -0,0 +1,236 @@ +//! 常驻后台的 WebSocket 泵（对齐官方 Codex CLI 的连接管理）。 +//! +//! 每条上游 WebSocket 连接都由一个后台 pump 任务独占： +//! - 持续读取 socket：一旦观察到 `Close` / EOF / 传输错误，立即把连接标记为 `closed`。 +//! - 自动回应上游 `Ping`（保活），并按 `ping_interval` 主动 `Ping`（穿透 NAT/中间盒空闲计时器）。 +//! - 可选 `liveness_timeout`：长时间无任何入站活动时判定连接失活并退出。 +//! +//! 因此空闲连接的“死没死”在后台被实时感知；复用方只需零成本读取 [`PumpedWebSocket::is_closed`]， +//! 无需在请求路径上做同步探活 ping（消除“复用到静默死连接才卡住超时”的长尾）。 +//! +//! 收发都通过 channel 与 pump 任务交互： +//! - 发送：`send` 走 command channel，等待 pump 回执。 +//! - 接收：`next` 从 message channel 取出 pump 转发的入站帧（`Ping`/`Pong` 已被 pump 吞掉）。 + +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; +use std::time::{Duration, Instant}; + +use futures::{SinkExt, StreamExt}; +use tokio::{ + net::TcpStream, + sync::{mpsc, oneshot}, + task::JoinHandle, + time::MissedTickBehavior, +}; +use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream}; + +/// 底层 tungstenite WebSocket 流。 +pub(crate) type RawWsStream = WebSocketStream>; + +const PUMP_COMMAND_BUFFER: usize = 32; + +/// pump 保活策略。 +#[derive(Debug, Clone, Copy)] +pub(crate) struct PumpKeepalive { + /// 主动 `Ping` 间隔；`None` 表示 pump 不主动 ping（仅被动读取 + 回应上游 ping）。 + pub(crate) ping_interval: Option, + /// 无入站活动多久后判定失活；`None` 表示只靠显式 close/传输错误发现死亡。 + pub(crate) liveness_timeout: Option, +} + +impl PumpKeepalive { + /// 不做主动保活（用于即用即弃的非池化连接）。 + pub(crate) fn disabled() -> Self { + Self { + ping_interval: None, + liveness_timeout: None, + } + } +} + +enum PumpCommand { + Send { + message: Message, + ack: oneshot::Sender>, + }, +} + +/// 由后台 pump 任务托管的 WebSocket 连接句柄。 +pub(crate) struct PumpedWebSocket { + tx_command: mpsc::Sender, + rx_message: mpsc::UnboundedReceiver>, + closed: Arc, + pump: Option>, +} + +impl std::fmt::Debug for PumpedWebSocket { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PumpedWebSocket") + .field("closed", &self.closed.load(Ordering::Acquire)) + .finish() + } +} + +impl PumpedWebSocket { + /// 用底层流启动一个 pump 任务并返回句柄。 + pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self { + let (tx_command, rx_command) = mpsc::channel::(PUMP_COMMAND_BUFFER); + let (tx_message, rx_message) = mpsc::unbounded_channel(); + let closed = Arc::new(AtomicBool::new(false)); + let closed_for_task = Arc::clone(&closed); + let pump = tokio::spawn(async move { + pump_loop(inner, rx_command, tx_message, keepalive).await; + closed_for_task.store(true, Ordering::Release); + }); + Self { + tx_command, + rx_message, + closed, + pump: Some(pump), + } + } + + /// 通过 pump 发送一帧，返回底层 `send` 的结果。 + pub(crate) async fn send(&self, message: Message) -> Result<(), tungstenite::Error> { + let (ack, rx_ack) = oneshot::channel(); + if self + .tx_command + .send(PumpCommand::Send { message, ack }) + .await + .is_err() + { + return Err(tungstenite::Error::ConnectionClosed); + } + rx_ack + .await + .unwrap_or(Err(tungstenite::Error::ConnectionClosed)) + } + + /// 取出下一帧入站消息（`Ping`/`Pong` 已被 pump 处理，不会到达这里）。 + /// + /// 返回 `None` 表示连接已结束（pump 已退出且缓冲已排空）。 + pub(crate) async fn next(&mut self) -> Option> { + self.rx_message.recv().await + } + + /// 连接是否已被后台 pump 判定关闭/失活。零成本，用于复用前探活。 + pub(crate) fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) || self.tx_command.is_closed() + } + + /// 主动关闭连接（best-effort 发送 Close 帧）。 + pub(crate) async fn close(&mut self) { + let _ = self.send(Message::Close(None)).await; + } +} + +impl Drop for PumpedWebSocket { + fn drop(&mut self) { + if let Some(pump) = self.pump.take() { + pump.abort(); + } + } +} + +async fn pump_loop( + mut inner: RawWsStream, + mut rx_command: mpsc::Receiver, + tx_message: mpsc::UnboundedSender>, + keepalive: PumpKeepalive, +) { + let mut last_activity = Instant::now(); + let ping_interval = keepalive.ping_interval.filter(|d| !d.is_zero()); + let liveness_timeout = keepalive.liveness_timeout.filter(|d| !d.is_zero()); + // ticker 由 ping 间隔或 liveness 检查间隔中较小者驱动：即使不主动 ping， + // 也要周期性醒来检查 liveness。两者都未配置时不启动 ticker。 + let tick_interval = match (ping_interval, liveness_timeout) { + (Some(ping), Some(live)) => Some(ping.min(live)), + (Some(ping), None) => Some(ping), + (None, Some(live)) => Some(live), + (None, None) => None, + }; + let mut ticker = tick_interval.map(|d| { + // 首个 tick 推迟一整个间隔：tokio::time::interval 默认会让首个 tick 立即就绪， + // 否则连接一建立就会立刻发一帧 Ping，与首个请求 Text 抢跑、打乱帧序。 + let mut ticker = tokio::time::interval_at(tokio::time::Instant::now() + d, d); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + ticker + }); + + loop { + tokio::select! { + command = rx_command.recv() => { + match command { + // 所有句柄已丢弃：连接不再被使用，退出 pump。 + None => break, + Some(PumpCommand::Send { message, ack }) => { + let result = inner.send(message).await; + let should_break = result.is_err(); + let _ = ack.send(result); + if should_break { + break; + } + } + } + } + message = inner.next() => { + match message { + None => break, + Some(Ok(Message::Ping(payload))) => { + last_activity = Instant::now(); + if inner.send(Message::Pong(payload)).await.is_err() { + break; + } + } + Some(Ok(Message::Pong(_))) => { + last_activity = Instant::now(); + } + Some(Ok(message)) => { + last_activity = Instant::now(); + let is_close = matches!(message, Message::Close(_)); + if tx_message.send(Ok(message)).is_err() { + break; + } + if is_close { + break; + } + } + Some(Err(err)) => { + let _ = tx_message.send(Err(err)); + break; + } + } + } + _ = tick(&mut ticker) => { + // 先检查 liveness：长时间无入站活动则判定失活退出。 + if let Some(timeout) = liveness_timeout { + if last_activity.elapsed() >= timeout { + break; + } + } + // 再按需主动 ping 保活（穿透 NAT/中间盒空闲计时器）。 + if ping_interval.is_some() + && inner.send(Message::Ping(Vec::new().into())).await.is_err() + { + break; + } + } + } + } + + // best-effort：尝试通知对端关闭，随后 drop 底层流释放 socket。 + let _ = inner.send(Message::Close(None)).await; +} + +/// 等待 ping ticker；`None` 时永远挂起，让 `select!` 分支实际禁用。 +async fn tick(ticker: &mut Option) { + match ticker { + Some(ticker) => { + ticker.tick().await; + } + None => std::future::pending::<()>().await, + } +} diff --git a/backend/tests/upstream/transport/websocket_pool.rs b/backend/tests/upstream/transport/websocket_pool.rs index 09dd4e919..77b3f85c0 100644 --- a/backend/tests/upstream/transport/websocket_pool.rs +++ b/backend/tests/upstream/transport/websocket_pool.rs @@ -74,6 +74,87 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and assert_eq!(accepted_connections.load(Ordering::SeqCst), 1); } +/// idle 连接被上游静默关闭后，后台 pump 会实时把它标记为 closed。 +/// 复用前的零成本 `is_closed` 检查应直接丢弃它并新建连接，不经过 +/// “发请求 → 等首帧超时 → stale-reuse 重试” 的长尾（无需任何 maintenance sweep）。 +#[tokio::test] +async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocket_died_silently() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + // 第一条连接：完成一次响应后由服务端主动关闭（模拟 idle 期间被上游/中间盒断开）。 + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_codex_test_websocket(first_stream).await; + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + completed_websocket_response("resp_silent_first", 3, 1).into(), + )) + .await + .unwrap(); + first_websocket.close(None).await.unwrap(); + + // 第二条连接：证明复用被跳过、直接新建。 + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_codex_test_websocket(second_stream).await; + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + completed_websocket_response("resp_silent_second", 4, 1).into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + // 无 maintenance、无主动 ping：完全依赖 pump 后台读取感知连接死亡。 + let pool = Arc::new(CodexWebSocketPool::with_config( + websocket_pool_config_for_tests(None, None, None), + )); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + crate::support::fingerprint::test_fingerprint(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.set_previous_response_id(Some("resp_previous".to_string())); + request.set_prompt_cache_key(Some("conversation-pool-silent".to_string())); + + let first = backend + .create_response( + &request, + request_context("req_pool_silent_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + // 给后台 pump 一点时间观察到服务端的 close 帧并标记连接死亡。 + tokio::time::sleep(Duration::from_millis(50)).await; + let second = backend + .create_response( + &request, + request_context("req_pool_silent_second", Some("chatgpt-account")), + ) + .await + .expect("second websocket response should open a fresh socket"); + server.await.unwrap(); + + assert!(first.body.contains("resp_silent_first")); + assert!(second.body.contains("resp_silent_second")); + // 死连接在 acquire 处被零成本识别 → 直接新建，而非 stale-reuse 重试。 + assert_eq!(first.websocket_pool_decision.unwrap().kind(), "new"); + assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + #[tokio::test] async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_accounts() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -403,112 +484,10 @@ async fn websocket_pool_should_bypass_new_keys_after_account_cap() { } #[tokio::test] -async fn codex_backend_client_should_ping_idle_pooled_websocket_during_maintenance() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let ping_count = Arc::new(AtomicUsize::new(0)); - let ping_count_for_server = Arc::clone(&ping_count); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_codex_test_websocket(stream).await; - let _first_message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_pool_keepalive_first", - "object": "response", - "output": [], - "usage": { - "input_tokens": 3, - "output_tokens": 1, - "total_tokens": 4 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - - let ping = timeout(Duration::from_secs(1), websocket.next()) - .await - .expect("pool maintenance should probe the idle websocket") - .expect("pool maintenance should send a websocket frame") - .expect("pool maintenance frame should be valid"); - let Message::Ping(payload) = ping else { - panic!("expected pool maintenance ping frame, got {ping:?}"); - }; - ping_count_for_server.fetch_add(1, Ordering::SeqCst); - websocket.send(Message::Pong(payload)).await.unwrap(); - - let _second_message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_pool_keepalive_second", - "object": "response", - "output": [], - "usage": { - "input_tokens": 3, - "output_tokens": 1, - "total_tokens": 4 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - let pool = Arc::new(CodexWebSocketPool::with_config( - websocket_pool_config_for_tests(None, Some(Duration::from_millis(1)), None), - )); - let backend = CodexBackendClient::new( - reqwest::Client::builder().no_proxy().build().unwrap(), - format!("http://{addr}"), - crate::support::fingerprint::test_fingerprint(), - ) - .with_websocket_pool(Arc::clone(&pool)); - let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "be brief", - Vec::new(), - ); - request.set_previous_response_id(Some("resp_previous".to_string())); - request.set_prompt_cache_key(Some("conversation-pool-keepalive".to_string())); - - let first = backend - .create_response( - &request, - request_context("req_pool_keepalive_first", Some("chatgpt-account")), - ) - .await - .expect("first pooled websocket response should succeed"); - pool.maintain_idle_connections().await; - let second = backend - .create_response( - &request, - request_context("req_pool_keepalive_second", Some("chatgpt-account")), - ) - .await - .expect("second pooled websocket response should reuse the probed socket"); - server.await.unwrap(); - - assert!(first.body.contains("resp_pool_keepalive_first")); - assert!(second.body.contains("resp_pool_keepalive_second")); - assert_eq!(ping_count.load(Ordering::SeqCst), 1); -} - -#[tokio::test] -async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() { +async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despite_pings() { + // 不变量：pump 自己发出的 keepalive ping 不算“入站活动”，因此一个只收 ping、 + // 从不回 pong / 从不发帧的静默连接，仍会被 liveness watchdog 判定失活并驱逐； + // 复用前 acquire 读到 is_closed 即开新连接。 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let accepted_connections = Arc::new(AtomicUsize::new(0)); @@ -524,6 +503,7 @@ async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() { )) .await .unwrap(); + // 故意不再读取该连接：既不回 pong，也不发任何帧，模拟静默失活。 let (second_stream, _) = listener.accept().await.unwrap(); accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); let mut second_websocket = accept_codex_test_websocket(second_stream).await; @@ -539,7 +519,7 @@ async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() { }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { ping_interval: Some(Duration::from_millis(1)), - ping_timeout: Duration::from_millis(20), + liveness_timeout: Some(Duration::from_millis(20)), maintenance_interval: None, ..websocket_pool_config_for_tests(None, None, None) })); @@ -558,7 +538,8 @@ async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() { ) .await .expect("first websocket response should succeed"); - pool.maintain_idle_connections().await; + // 等待 pump 的 liveness watchdog 判定静默连接失活。 + tokio::time::sleep(Duration::from_millis(60)).await; let second = backend .create_response( &request, @@ -648,7 +629,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() { } #[tokio::test] -async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_maintenance() { +async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_repeated_pings() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let ping_count = Arc::new(AtomicUsize::new(0)); @@ -678,18 +659,23 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_ .await .unwrap(); - let ping = timeout(Duration::from_secs(1), websocket.next()) - .await - .expect("background maintenance should probe the idle websocket") - .expect("background maintenance should send a websocket frame") - .expect("background maintenance frame should be valid"); - let Message::Ping(payload) = ping else { - panic!("expected background maintenance ping frame, got {ping:?}"); - }; - ping_count_for_server.fetch_add(1, Ordering::SeqCst); - websocket.send(Message::Pong(payload)).await.unwrap(); - - let _second_message = websocket.next().await.unwrap().unwrap(); + // pump 会在 idle 期间反复发送 keepalive ping；服务端计数并回 pong， + // 直到下一个业务请求（response.create）到达为止。 + loop { + let message = timeout(Duration::from_secs(1), websocket.next()) + .await + .expect("pump keepalive / second request should arrive") + .expect("frame should be present") + .expect("frame should be valid"); + match message { + Message::Ping(payload) => { + ping_count_for_server.fetch_add(1, Ordering::SeqCst); + websocket.send(Message::Pong(payload)).await.unwrap(); + } + Message::Text(_) => break, + other => panic!("unexpected frame while idle: {other:?}"), + } + } websocket .send(Message::Text( json!({ @@ -713,11 +699,7 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_ websocket.close(None).await.unwrap(); }); let pool = Arc::new(CodexWebSocketPool::with_config( - websocket_pool_config_for_tests( - Some(Duration::from_millis(20)), - Some(Duration::from_mins(1)), - None, - ), + websocket_pool_config_for_tests(None, Some(Duration::from_millis(10)), None), )); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), @@ -741,6 +723,7 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_ ) .await .expect("first pooled websocket response should succeed"); + // 让 pump 有时间发出多轮 keepalive ping。 tokio::time::sleep(Duration::from_millis(80)).await; let second = backend .create_response( @@ -748,13 +731,13 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_ request_context("req_pool_background_second", Some("chatgpt-account")), ) .await - .expect("second pooled websocket response should reuse the background-probed socket"); + .expect("second pooled websocket response should reuse the kept-alive socket"); server.await.unwrap(); pool.shutdown().await; assert!(first.body.contains("resp_pool_background_first")); assert!(second.body.contains("resp_pool_background_second")); - assert_eq!(ping_count.load(Ordering::SeqCst), 1); + assert!(ping_count.load(Ordering::SeqCst) >= 1); } #[tokio::test]