starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: move websocket keepalive into pump

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-05 19:38:15 +08:00
co-authored by Codex
parent 0c18b44246
commit 617bd676f0
5 changed files with 426 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/upstream/transport/mod.rs
+1
View File
@@ -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::{
backend/src/upstream/transport/websocket.rs
+35 -30
View File
@@ -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 = "<redacted>";
@@ -318,7 +319,7 @@ fn websocket_host_header(endpoint: &str) -> Option<String> {
async fn connect_websocket(
connection: &CodexWebSocketConnection,
) -> Result<(CodexWsStream, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
) -> Result<(RawWsStream, WsResponse<Option<Vec<u8>>>), 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<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
let (raw, response) = connect_websocket(connection).await?;
Ok((PumpedWebSocket::new(raw, keepalive), response))
}
fn websocket_handshake_request(
connection: &CodexWebSocketConnection,
) -> Result<WsRequest, tungstenite::http::Error> {
@@ -400,7 +413,8 @@ async fn execute_fresh_response_create_request(
request: &CodexWebSocketRequest,
started_at: Instant,
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
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<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
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<CodexWebSocketExchange, CodexWebSocketExchangeError> {
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<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
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<CodexWebSocketExchange, CodexWebSocketExchangeError> {
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<String, CodexWebSocketExchangeError> {
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<Instant>,
}
fn stream_websocket_response(
websocket: CodexWsStream,
websocket: PumpedWebSocket,
metadata: CodexWebSocketConnectionMetadata,
pool_return: Option<WebSocketStreamPoolReturn>,
prefetched_text: Option<String>,
@@ -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<WebSocketStreamPoolReturn>,
prefetched_text: Option<String>,
@@ -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<Option<Message>, 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<WebSocketStreamPoolReturn>,
) {
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<WebSocketStreamPoolReturn>,
) {
if let Some(pool_return) = pool_return {
pool_return.pool.discard(&pool_return.key).await;
}
let _ = websocket.close(None).await;
websocket.close().await;
}
// ---------------------------------------------------------------------------
backend/src/upstream/transport/websocket_pool.rs
+43 -153
View File
@@ -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::<Vec<_>>()
};
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<PooledWebSocketConnection> {
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::<Vec<_>>();
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::<Vec<_>>();
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<PooledWebSocketConnection>,
) {
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<MaybeTlsStream<TcpStream>>;
#[derive(Clone)]
pub(crate) struct CodexWebSocketConnectionMetadata {
pub(crate) turn_state: Option<String>,
@@ -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<Instant>,
}
enum WebSocketPoolSlot {
Idle(Box<PooledWebSocketConnection>),
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<PooledWebSocketConnection>) {
@@ -547,59 +485,11 @@ async fn close_pooled_connections(connections: Vec<PooledWebSocketConnection>) {
}
}
struct WebSocketPoolMaintenance {
close: Vec<PooledWebSocketConnection>,
probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>,
}
/// idle 连接是否应从池中摘除：被后台 pump 标记死亡，或已超过 `max_age`。
fn should_close_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
max_age: Duration,
liveness_timeout: Option<Duration>,
) -> 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<PooledWebSocketConnection> {
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
}
backend/src/upstream/transport/websocket_pump.rs
+236
View File
@@ -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<MaybeTlsStream<TcpStream>>;
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<Duration>,
/// 无入站活动多久后判定失活；`None` 表示只靠显式 close/传输错误发现死亡。
pub(crate) liveness_timeout: Option<Duration>,
}
impl PumpKeepalive {
/// 不做主动保活（用于即用即弃的非池化连接）。
pub(crate) fn disabled() -> Self {
Self {
ping_interval: None,
liveness_timeout: None,
}
}
}
enum PumpCommand {
Send {
message: Message,
ack: oneshot::Sender<Result<(), tungstenite::Error>>,
},
}
/// 由后台 pump 任务托管的 WebSocket 连接句柄。
pub(crate) struct PumpedWebSocket {
tx_command: mpsc::Sender<PumpCommand>,
rx_message: mpsc::UnboundedReceiver<Result<Message, tungstenite::Error>>,
closed: Arc<AtomicBool>,
pump: Option<JoinHandle<()>>,
}
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::<PumpCommand>(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<Result<Message, tungstenite::Error>> {
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<PumpCommand>,
tx_message: mpsc::UnboundedSender<Result<Message, tungstenite::Error>>,
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<tokio::time::Interval>) {
match ticker {
Some(ticker) => {
ticker.tick().await;
}
None => std::future::pending::<()>().await,
}
}
backend/tests/upstream/transport/websocket_pool.rs
+111 -128
View File
@@ -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]