mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix(transport): preserve websocket openings after fallback
Continue pool-owned WebSocket handshakes after the foreground budget expires, retain late successes for reuse, and cancel them on account eviction or shutdown. Keep breaker accounting and transport metrics aligned with the foreground decision. Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -309,7 +309,7 @@ pub enum CodexTransportDecision {
|
||||
ConnectedWebSocket,
|
||||
ExactWebSocket,
|
||||
RequiredWebSocket,
|
||||
Http2WebSocketSlow,
|
||||
Http2WebSocketBudgetExhausted,
|
||||
Http2BreakerOpen,
|
||||
Http2PoolUnavailable,
|
||||
Http2PreSendFailure,
|
||||
@@ -323,7 +323,7 @@ impl CodexTransportDecision {
|
||||
Self::ConnectedWebSocket => "ws_connected_fast",
|
||||
Self::ExactWebSocket => "ws_exact_required",
|
||||
Self::RequiredWebSocket => "ws_required",
|
||||
Self::Http2WebSocketSlow => "http2_ws_slow",
|
||||
Self::Http2WebSocketBudgetExhausted => "http2_ws_budget_exhausted",
|
||||
Self::Http2BreakerOpen => "http2_breaker_open",
|
||||
Self::Http2PoolUnavailable => "http2_pool_unavailable",
|
||||
Self::Http2PreSendFailure => "http2_ws_pre_send_failure",
|
||||
@@ -627,7 +627,7 @@ fn websocket_success_decision(
|
||||
fn http_fallback_decision(error: &CodexWebSocketExchangeError) -> CodexTransportDecision {
|
||||
match error {
|
||||
CodexWebSocketExchangeError::FastPathTimeout { .. } => {
|
||||
CodexTransportDecision::Http2WebSocketSlow
|
||||
CodexTransportDecision::Http2WebSocketBudgetExhausted
|
||||
}
|
||||
CodexWebSocketExchangeError::OriginCircuitOpen
|
||||
| CodexWebSocketExchangeError::OriginHalfOpenBusy => {
|
||||
|
||||
@@ -314,11 +314,7 @@ impl CodexBackendClient {
|
||||
route: PreparedResponseRoute::Http,
|
||||
metrics: CodexTransportMetrics {
|
||||
decision: Some(decision),
|
||||
ws_connect_ms: matches!(
|
||||
error,
|
||||
CodexWebSocketExchangeError::FastPathTimeout { .. }
|
||||
)
|
||||
.then_some(wait_ms),
|
||||
ws_connect_ms: None,
|
||||
transport_decision_wait_ms: Some(wait_ms),
|
||||
..CodexTransportMetrics::default()
|
||||
},
|
||||
|
||||
@@ -12,7 +12,7 @@ use bytes::Bytes;
|
||||
use chrono::Utc;
|
||||
use futures::Stream;
|
||||
use thiserror::Error;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio::{sync::Mutex, time::timeout};
|
||||
use tokio_tungstenite::{Connector, connect_async_tls_with_config};
|
||||
use tungstenite::{
|
||||
@@ -39,12 +39,13 @@ use crate::upstream::openai::protocol::websocket::{
|
||||
|
||||
use super::websocket_breaker::{
|
||||
WebSocketOriginBreaker, WebSocketOriginBreakerDecision, WebSocketOriginBreakerPermit,
|
||||
WebSocketOriginFastPathReporter,
|
||||
};
|
||||
use super::websocket_pool::{
|
||||
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey,
|
||||
PooledWebSocketConnection, WebSocketContinuationState, WebSocketPoolAcquire,
|
||||
WebSocketPoolBypassReason, WebSocketPoolConnectOutcome, WebSocketPoolDecision,
|
||||
WebSocketPoolLease,
|
||||
WebSocketPoolBypassReason, WebSocketPoolConnectLease, WebSocketPoolConnectOutcome,
|
||||
WebSocketPoolDecision, WebSocketPoolLease,
|
||||
};
|
||||
use super::websocket_pump::{PumpExitReason, PumpKeepalive, PumpedWebSocket, RawWsStream};
|
||||
use super::{client::CodexResponseMetadata, diagnostics::CodexUpstreamDiagnostics, response_meta};
|
||||
@@ -576,7 +577,7 @@ async fn prepare_unpooled_websocket(
|
||||
async fn prepare_pooled_websocket(
|
||||
request: &CodexWebSocketRequest,
|
||||
pool: &CodexWebSocketPool,
|
||||
connect_lease: super::websocket_pool::WebSocketPoolConnectLease,
|
||||
connect_lease: WebSocketPoolConnectLease,
|
||||
breaker: &WebSocketOriginBreaker,
|
||||
origin_key: &str,
|
||||
fast_path_budget: Option<Duration>,
|
||||
@@ -589,27 +590,21 @@ async fn prepare_pooled_websocket(
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let connected =
|
||||
connect_with_budget(request.connection(), pool.keepalive(), fast_path_budget).await;
|
||||
let (connection, connect_elapsed) = match finish_breaker_attempt(permit, connected) {
|
||||
Ok(connected) => connected,
|
||||
Err(error) => {
|
||||
connect_lease.failed().await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let connected = match connect_lease.connected(connection).await {
|
||||
Ok(connected) => connected,
|
||||
Err(connection) => {
|
||||
connection.websocket.close().await;
|
||||
return Err(continuation_unavailable(
|
||||
PreviousResponseUnavailableReason::PoolUnavailable,
|
||||
));
|
||||
}
|
||||
let mut waiter = start_pooled_websocket_connect(
|
||||
request.connection().clone(),
|
||||
pool.clone(),
|
||||
connect_lease,
|
||||
permit,
|
||||
);
|
||||
let connect_elapsed = waiter.wait(fast_path_budget).await?;
|
||||
let Some((connection, lease)) = pool.take_idle(waiter.key()).await else {
|
||||
return Err(continuation_unavailable(
|
||||
PreviousResponseUnavailableReason::PoolUnavailable,
|
||||
));
|
||||
};
|
||||
Ok(PreparedWebSocket {
|
||||
connection: *connected.connection,
|
||||
lease: Some(connected.lease),
|
||||
connection: *connection,
|
||||
lease: Some(lease),
|
||||
reused: false,
|
||||
pool_decision: Some(WebSocketPoolDecision::new()),
|
||||
connect_elapsed: Some(connect_elapsed),
|
||||
@@ -618,6 +613,131 @@ async fn prepare_pooled_websocket(
|
||||
})
|
||||
}
|
||||
|
||||
struct PooledWebSocketConnectWaiter {
|
||||
key: CodexWebSocketPoolKey,
|
||||
started_at: tokio::time::Instant,
|
||||
receiver: oneshot::Receiver<Result<Duration, CodexWebSocketExchangeError>>,
|
||||
fast_path_reporter: WebSocketOriginFastPathReporter,
|
||||
}
|
||||
|
||||
impl PooledWebSocketConnectWaiter {
|
||||
fn key(&self) -> &CodexWebSocketPoolKey {
|
||||
&self.key
|
||||
}
|
||||
|
||||
async fn wait(
|
||||
&mut self,
|
||||
fast_path_budget: Option<Duration>,
|
||||
) -> Result<Duration, CodexWebSocketExchangeError> {
|
||||
let received = match fast_path_budget {
|
||||
Some(budget) => {
|
||||
let remaining = budget.saturating_sub(self.started_at.elapsed());
|
||||
match timeout(remaining, &mut self.receiver).await {
|
||||
Ok(received) => received,
|
||||
Err(_) => {
|
||||
self.fast_path_reporter.missed();
|
||||
return Err(CodexWebSocketExchangeError::FastPathTimeout {
|
||||
timeout: budget,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
None => (&mut self.receiver).await,
|
||||
};
|
||||
received.map_err(|_| CodexWebSocketExchangeError::SharedConnectFailed)?
|
||||
}
|
||||
}
|
||||
|
||||
fn start_pooled_websocket_connect(
|
||||
connection: CodexWebSocketConnection,
|
||||
pool: CodexWebSocketPool,
|
||||
connect_lease: WebSocketPoolConnectLease,
|
||||
permit: WebSocketOriginBreakerPermit,
|
||||
) -> PooledWebSocketConnectWaiter {
|
||||
let key = connect_lease.key().clone();
|
||||
let started_at = connect_lease.started_at();
|
||||
let cancellation = connect_lease.cancellation_token();
|
||||
let fast_path_reporter = permit.fast_path_reporter();
|
||||
let keepalive = pool.keepalive();
|
||||
let task_key = key.clone();
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
pool.spawn_connect_task(async move {
|
||||
let connected = tokio::select! {
|
||||
biased;
|
||||
_ = cancellation.cancelled() => {
|
||||
permit.cancel();
|
||||
connect_lease.failed().await;
|
||||
let _ = sender.send(Err(CodexWebSocketExchangeError::SharedConnectFailed));
|
||||
tracing::info!(
|
||||
account_id = task_key.account_id(),
|
||||
conversation_id_hash = task_key.conversation_id_hash(),
|
||||
ws_preconnect_duration_ms = duration_millis_u64(started_at.elapsed()),
|
||||
ws_preconnect_outcome = "cancelled",
|
||||
"WebSocket pool connect finished"
|
||||
);
|
||||
return;
|
||||
}
|
||||
result = connect_with_budget(&connection, keepalive, None) => result,
|
||||
};
|
||||
match finish_breaker_attempt(permit, connected) {
|
||||
Ok((connection, connect_elapsed)) => {
|
||||
match connect_lease.connected_idle(connection).await {
|
||||
Ok(()) => {
|
||||
let foreground_waiting = sender.send(Ok(connect_elapsed)).is_ok();
|
||||
tracing::info!(
|
||||
account_id = task_key.account_id(),
|
||||
conversation_id_hash = task_key.conversation_id_hash(),
|
||||
ws_preconnect_duration_ms = duration_millis_u64(connect_elapsed),
|
||||
foreground_waiting,
|
||||
ws_preconnect_outcome = "ready",
|
||||
"WebSocket pool connect finished"
|
||||
);
|
||||
}
|
||||
Err(connection) => {
|
||||
connection.websocket.close().await;
|
||||
let _ = sender.send(Err(continuation_unavailable(
|
||||
PreviousResponseUnavailableReason::PoolUnavailable,
|
||||
)));
|
||||
tracing::info!(
|
||||
account_id = task_key.account_id(),
|
||||
conversation_id_hash = task_key.conversation_id_hash(),
|
||||
ws_preconnect_duration_ms = duration_millis_u64(connect_elapsed),
|
||||
ws_preconnect_outcome = "rejected",
|
||||
"WebSocket pool connect finished"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
let error_message = error.to_string();
|
||||
connect_lease.failed().await;
|
||||
let foreground_waiting = sender.send(Err(error)).is_ok();
|
||||
tracing::warn!(
|
||||
account_id = task_key.account_id(),
|
||||
conversation_id_hash = task_key.conversation_id_hash(),
|
||||
ws_preconnect_duration_ms = duration_millis_u64(started_at.elapsed()),
|
||||
foreground_waiting,
|
||||
error = %error_message,
|
||||
ws_preconnect_outcome = "failed",
|
||||
"WebSocket pool connect finished"
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
PooledWebSocketConnectWaiter {
|
||||
key,
|
||||
started_at,
|
||||
receiver,
|
||||
fast_path_reporter,
|
||||
}
|
||||
}
|
||||
|
||||
fn duration_millis_u64(duration: Duration) -> u64 {
|
||||
u64::try_from(duration.as_millis())
|
||||
.unwrap_or(u64::MAX)
|
||||
.max(1)
|
||||
}
|
||||
|
||||
async fn connect_with_budget(
|
||||
connection: &CodexWebSocketConnection,
|
||||
keepalive: PumpKeepalive,
|
||||
@@ -683,11 +803,11 @@ fn finish_breaker_attempt(
|
||||
async fn wait_for_shared_connect(
|
||||
waiter: super::websocket_pool::WebSocketPoolConnectWaiter,
|
||||
fast_path_budget: Option<Duration>,
|
||||
decision_started_at: Instant,
|
||||
_decision_started_at: Instant,
|
||||
) -> Result<WebSocketPoolConnectOutcome, CodexWebSocketExchangeError> {
|
||||
match fast_path_budget {
|
||||
Some(budget) => {
|
||||
let remaining = budget.saturating_sub(decision_started_at.elapsed());
|
||||
let remaining = waiter.remaining_budget(budget);
|
||||
timeout(remaining, waiter.wait())
|
||||
.await
|
||||
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })
|
||||
|
||||
@@ -2,7 +2,10 @@
|
||||
|
||||
use std::{
|
||||
collections::{HashMap, VecDeque},
|
||||
sync::{Arc, Mutex},
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
@@ -73,8 +76,8 @@ impl WebSocketOriginBreaker {
|
||||
.or_insert_with(CircuitState::closed);
|
||||
|
||||
match state {
|
||||
CircuitState::Closed { fast_timeouts } => {
|
||||
retain_window(fast_timeouts, now, self.config.failure_window);
|
||||
CircuitState::Closed { failures } => {
|
||||
retain_window(failures, now, self.config.failure_window);
|
||||
WebSocketOriginBreakerDecision::Allowed(WebSocketOriginBreakerPermit::new(
|
||||
self.clone(),
|
||||
origin_key.to_string(),
|
||||
@@ -105,7 +108,7 @@ impl WebSocketOriginBreaker {
|
||||
}
|
||||
}
|
||||
|
||||
fn record_fast_timeout(&self, origin_key: &str, probe_id: Option<Uuid>) {
|
||||
fn record_failure(&self, origin_key: &str, probe_id: Option<Uuid>) {
|
||||
let now = Instant::now();
|
||||
let mut circuits = self
|
||||
.inner
|
||||
@@ -127,35 +130,17 @@ impl WebSocketOriginBreaker {
|
||||
let state = circuits
|
||||
.entry(origin_key.to_string())
|
||||
.or_insert_with(CircuitState::closed);
|
||||
let CircuitState::Closed { fast_timeouts } = state else {
|
||||
let CircuitState::Closed { failures } = state else {
|
||||
return;
|
||||
};
|
||||
retain_window(fast_timeouts, now, self.config.failure_window);
|
||||
fast_timeouts.push_back(now);
|
||||
if fast_timeouts.len() >= self.config.failure_threshold {
|
||||
retain_window(failures, now, self.config.failure_window);
|
||||
failures.push_back(now);
|
||||
if failures.len() >= self.config.failure_threshold {
|
||||
*state = CircuitState::Open {
|
||||
until: now + self.config.open_duration,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
fn record_failure(&self, origin_key: &str, probe_id: Option<Uuid>) {
|
||||
let Some(probe_id) = probe_id else {
|
||||
return;
|
||||
};
|
||||
let mut circuits = self
|
||||
.inner
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if permit_still_owns_state(circuits.get(origin_key), Some(probe_id)) {
|
||||
circuits.insert(
|
||||
origin_key.to_string(),
|
||||
CircuitState::Open {
|
||||
until: Instant::now() + self.config.open_duration,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 一次新的 WebSocket opening 许可;消费式完成保证 half-open 探针不会泄漏。
|
||||
@@ -163,6 +148,7 @@ pub struct WebSocketOriginBreakerPermit {
|
||||
breaker: WebSocketOriginBreaker,
|
||||
origin_key: String,
|
||||
probe_id: Option<Uuid>,
|
||||
degradation_recorded: Arc<AtomicBool>,
|
||||
armed: bool,
|
||||
}
|
||||
|
||||
@@ -172,6 +158,7 @@ impl WebSocketOriginBreakerPermit {
|
||||
breaker,
|
||||
origin_key,
|
||||
probe_id,
|
||||
degradation_recorded: Arc::new(AtomicBool::new(false)),
|
||||
armed: true,
|
||||
}
|
||||
}
|
||||
@@ -180,33 +167,71 @@ impl WebSocketOriginBreakerPermit {
|
||||
self.probe_id.is_some()
|
||||
}
|
||||
|
||||
/// 创建可跨前台预算与后台 opening 生命周期共享的快路径观察器。
|
||||
pub(crate) fn fast_path_reporter(&self) -> WebSocketOriginFastPathReporter {
|
||||
WebSocketOriginFastPathReporter {
|
||||
breaker: self.breaker.clone(),
|
||||
origin_key: self.origin_key.clone(),
|
||||
probe_id: self.probe_id,
|
||||
degradation_recorded: Arc::clone(&self.degradation_recorded),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn succeed(mut self) {
|
||||
self.breaker.record_success(&self.origin_key, self.probe_id);
|
||||
if !self.degradation_recorded.load(Ordering::Acquire) {
|
||||
self.breaker.record_success(&self.origin_key, self.probe_id);
|
||||
}
|
||||
self.armed = false;
|
||||
}
|
||||
|
||||
pub fn fast_timeout(mut self) {
|
||||
self.breaker
|
||||
.record_fast_timeout(&self.origin_key, self.probe_id);
|
||||
self.fast_path_reporter().missed();
|
||||
self.armed = false;
|
||||
}
|
||||
|
||||
pub fn fail(mut self) {
|
||||
self.breaker.record_failure(&self.origin_key, self.probe_id);
|
||||
self.record_failure_once();
|
||||
self.armed = false;
|
||||
}
|
||||
|
||||
/// opening 因账号驱逐或服务关闭而取消,不污染 origin 可用性。
|
||||
pub(crate) fn cancel(mut self) {
|
||||
self.armed = false;
|
||||
}
|
||||
|
||||
fn record_failure_once(&self) {
|
||||
if !self.degradation_recorded.swap(true, Ordering::AcqRel) {
|
||||
self.breaker.record_failure(&self.origin_key, self.probe_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WebSocketOriginBreakerPermit {
|
||||
fn drop(&mut self) {
|
||||
if self.armed {
|
||||
self.record_failure_once();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 前台等待结束后仍可报告同一次后台 opening 错过快路径预算。
|
||||
pub(crate) struct WebSocketOriginFastPathReporter {
|
||||
breaker: WebSocketOriginBreaker,
|
||||
origin_key: String,
|
||||
probe_id: Option<Uuid>,
|
||||
degradation_recorded: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl WebSocketOriginFastPathReporter {
|
||||
pub(crate) fn missed(&self) {
|
||||
if !self.degradation_recorded.swap(true, Ordering::AcqRel) {
|
||||
self.breaker.record_failure(&self.origin_key, self.probe_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum CircuitState {
|
||||
Closed { fast_timeouts: VecDeque<Instant> },
|
||||
Closed { failures: VecDeque<Instant> },
|
||||
Open { until: Instant },
|
||||
HalfOpen { probe_id: Uuid },
|
||||
}
|
||||
@@ -214,7 +239,7 @@ enum CircuitState {
|
||||
impl CircuitState {
|
||||
fn closed() -> Self {
|
||||
Self::Closed {
|
||||
fast_timeouts: VecDeque::new(),
|
||||
failures: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! Codex WebSocket 连接池。
|
||||
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
use std::{collections::HashMap, future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use serde_json::{Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -8,6 +8,10 @@ use tokio::{
|
||||
sync::{Mutex, watch},
|
||||
time::Instant,
|
||||
};
|
||||
use tokio_util::{
|
||||
sync::CancellationToken,
|
||||
task::{TaskTracker, task_tracker::TaskTrackerToken},
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::websocket_pump::{PumpKeepalive, PumpedWebSocket};
|
||||
@@ -63,6 +67,8 @@ impl CodexWebSocketPoolKey {
|
||||
pub struct CodexWebSocketPool {
|
||||
inner: Arc<Mutex<WebSocketPoolState>>,
|
||||
config: CodexWebSocketPoolConfig,
|
||||
tasks: TaskTracker,
|
||||
shutdown: CancellationToken,
|
||||
}
|
||||
|
||||
impl Default for CodexWebSocketPool {
|
||||
@@ -138,6 +144,8 @@ impl CodexWebSocketPool {
|
||||
let pool = Self {
|
||||
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
|
||||
config,
|
||||
tasks: TaskTracker::new(),
|
||||
shutdown: CancellationToken::new(),
|
||||
};
|
||||
pool.spawn_maintenance_task();
|
||||
pool
|
||||
@@ -153,6 +161,11 @@ impl CodexWebSocketPool {
|
||||
self.config.initial_event_timeout
|
||||
}
|
||||
|
||||
/// 注册由连接池生命周期托管的 opening 任务。
|
||||
pub(crate) fn spawn_connect_task(&self, future: impl Future<Output = ()> + Send + 'static) {
|
||||
drop(self.tasks.spawn(future));
|
||||
}
|
||||
|
||||
pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
|
||||
let mut expired_connection = None;
|
||||
let acquire = {
|
||||
@@ -167,6 +180,7 @@ impl CodexWebSocketPool {
|
||||
Some(WebSocketPoolSlot::Connecting(connecting)) => {
|
||||
return WebSocketPoolAcquire::Wait(WebSocketPoolConnectWaiter {
|
||||
receiver: connecting.outcome.subscribe(),
|
||||
started_at: connecting.started_at,
|
||||
});
|
||||
}
|
||||
Some(WebSocketPoolSlot::Idle(_)) => {
|
||||
@@ -199,8 +213,9 @@ impl CodexWebSocketPool {
|
||||
key.clone(),
|
||||
WebSocketPoolSlot::Connecting(WebSocketPoolConnecting {
|
||||
id: lease.id,
|
||||
started_at: Instant::now(),
|
||||
started_at: lease.started_at,
|
||||
outcome: lease.outcome.clone(),
|
||||
cancellation: lease.cancellation.clone(),
|
||||
}),
|
||||
);
|
||||
WebSocketPoolAcquire::Connect(lease)
|
||||
@@ -216,6 +231,39 @@ impl CodexWebSocketPool {
|
||||
acquire
|
||||
}
|
||||
|
||||
/// 只租用已经处于 Idle 的连接,不创建新的 Connecting slot。
|
||||
pub(crate) async fn take_idle(
|
||||
&self,
|
||||
key: &CodexWebSocketPoolKey,
|
||||
) -> Option<(Box<PooledWebSocketConnection>, WebSocketPoolLease)> {
|
||||
let mut expired_connection = None;
|
||||
let acquired = {
|
||||
let mut state = self.inner.lock().await;
|
||||
if !self.config.enabled || state.shutting_down {
|
||||
return None;
|
||||
}
|
||||
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else {
|
||||
return None;
|
||||
};
|
||||
if connection.created_at.elapsed() < self.config.max_age
|
||||
&& !connection.websocket.is_closed()
|
||||
{
|
||||
let lease = WebSocketPoolLease::reserve(self.clone(), key.clone());
|
||||
state
|
||||
.slots
|
||||
.insert(key.clone(), WebSocketPoolSlot::Busy(lease.reservation));
|
||||
Some((connection, lease))
|
||||
} else {
|
||||
expired_connection = Some(*connection);
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some(connection) = expired_connection {
|
||||
close_pooled_connection(connection).await;
|
||||
}
|
||||
acquired
|
||||
}
|
||||
|
||||
async fn put_reserved(
|
||||
&self,
|
||||
key: &CodexWebSocketPoolKey,
|
||||
@@ -256,26 +304,22 @@ impl CodexWebSocketPool {
|
||||
}
|
||||
}
|
||||
|
||||
async fn finish_connect(
|
||||
async fn finish_connect_idle(
|
||||
&self,
|
||||
key: &CodexWebSocketPoolKey,
|
||||
connect_id: Uuid,
|
||||
connection: PooledWebSocketConnection,
|
||||
) -> Result<WebSocketPoolConnected, Box<PooledWebSocketConnection>> {
|
||||
) -> Result<(), Box<PooledWebSocketConnection>> {
|
||||
let mut state = self.inner.lock().await;
|
||||
let owns_connect = matches!(
|
||||
state.slots.get(key),
|
||||
Some(WebSocketPoolSlot::Connecting(connecting)) if connecting.id == connect_id
|
||||
);
|
||||
if owns_connect && !state.shutting_down && self.config.enabled {
|
||||
let lease = WebSocketPoolLease::reserve(self.clone(), key.clone());
|
||||
state
|
||||
.slots
|
||||
.insert(key.clone(), WebSocketPoolSlot::Busy(lease.reservation));
|
||||
Ok(WebSocketPoolConnected {
|
||||
connection: Box::new(connection),
|
||||
lease,
|
||||
})
|
||||
.insert(key.clone(), WebSocketPoolSlot::Idle(Box::new(connection)));
|
||||
Ok(())
|
||||
} else {
|
||||
if owns_connect {
|
||||
state.slots.remove(key);
|
||||
@@ -294,7 +338,7 @@ impl CodexWebSocketPool {
|
||||
}
|
||||
}
|
||||
|
||||
/// 驱逐指定账号的 idle 连接,并阻止已占用 slot 回收到池中。
|
||||
/// 驱逐指定账号的全部 slot,取消 opening,并阻止 busy 连接回收到池中。
|
||||
pub async fn evict_account(&self, account_id: &str) {
|
||||
let mut idle_connections = Vec::new();
|
||||
{
|
||||
@@ -311,6 +355,7 @@ impl CodexWebSocketPool {
|
||||
idle_connections.push(*connection);
|
||||
}
|
||||
Some(WebSocketPoolSlot::Connecting(connecting)) => {
|
||||
connecting.cancellation.cancel();
|
||||
connecting
|
||||
.outcome
|
||||
.send_replace(WebSocketPoolConnectOutcome::Failed);
|
||||
@@ -322,8 +367,10 @@ impl CodexWebSocketPool {
|
||||
close_pooled_connections(idle_connections).await;
|
||||
}
|
||||
|
||||
/// 关闭连接池,关闭所有 idle 连接,并让后续 acquire 直接绕过池。
|
||||
/// 关闭连接池,取消受管任务、关闭 idle 连接,并让后续 acquire 直接绕过池。
|
||||
pub async fn shutdown(&self) {
|
||||
self.tasks.close();
|
||||
self.shutdown.cancel();
|
||||
let idle_connections = {
|
||||
let mut state = self.inner.lock().await;
|
||||
state.shutting_down = true;
|
||||
@@ -333,6 +380,7 @@ impl CodexWebSocketPool {
|
||||
.filter_map(|(_, slot)| match slot {
|
||||
WebSocketPoolSlot::Idle(connection) => Some(*connection),
|
||||
WebSocketPoolSlot::Connecting(connecting) => {
|
||||
connecting.cancellation.cancel();
|
||||
connecting
|
||||
.outcome
|
||||
.send_replace(WebSocketPoolConnectOutcome::Failed);
|
||||
@@ -343,6 +391,7 @@ impl CodexWebSocketPool {
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
close_pooled_connections(idle_connections).await;
|
||||
self.tasks.wait().await;
|
||||
}
|
||||
|
||||
/// 维护池 slot:清扫已死亡或超龄的 idle 连接,以及异常残留的 Busy reservation。
|
||||
@@ -364,20 +413,30 @@ impl CodexWebSocketPool {
|
||||
};
|
||||
let inner = Arc::downgrade(&self.inner);
|
||||
let config = self.config;
|
||||
tokio::spawn(async move {
|
||||
let tasks = self.tasks.clone();
|
||||
let shutdown = self.shutdown.clone();
|
||||
drop(self.tasks.spawn(async move {
|
||||
let mut interval = tokio::time::interval(interval_duration);
|
||||
loop {
|
||||
interval.tick().await;
|
||||
tokio::select! {
|
||||
_ = shutdown.cancelled() => break,
|
||||
_ = interval.tick() => {}
|
||||
}
|
||||
let Some(inner) = inner.upgrade() else {
|
||||
break;
|
||||
};
|
||||
let pool = CodexWebSocketPool { inner, config };
|
||||
let pool = CodexWebSocketPool {
|
||||
inner,
|
||||
config,
|
||||
tasks: tasks.clone(),
|
||||
shutdown: shutdown.clone(),
|
||||
};
|
||||
if pool.is_shutdown().await {
|
||||
break;
|
||||
}
|
||||
pool.maintain_idle_connections().await;
|
||||
}
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
/// 返回连接池是否已进入关闭状态。
|
||||
@@ -428,6 +487,7 @@ impl CodexWebSocketPool {
|
||||
);
|
||||
}
|
||||
Some(WebSocketPoolSlot::Connecting(connecting)) => {
|
||||
connecting.cancellation.cancel();
|
||||
connecting
|
||||
.outcome
|
||||
.send_replace(WebSocketPoolConnectOutcome::Failed);
|
||||
@@ -503,6 +563,7 @@ struct WebSocketPoolConnecting {
|
||||
id: Uuid,
|
||||
started_at: Instant,
|
||||
outcome: watch::Sender<WebSocketPoolConnectOutcome>,
|
||||
cancellation: CancellationToken,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -514,9 +575,14 @@ pub(crate) enum WebSocketPoolConnectOutcome {
|
||||
|
||||
pub(crate) struct WebSocketPoolConnectWaiter {
|
||||
receiver: watch::Receiver<WebSocketPoolConnectOutcome>,
|
||||
started_at: Instant,
|
||||
}
|
||||
|
||||
impl WebSocketPoolConnectWaiter {
|
||||
pub(crate) fn remaining_budget(&self, budget: Duration) -> Duration {
|
||||
budget.saturating_sub(self.started_at.elapsed())
|
||||
}
|
||||
|
||||
pub(crate) async fn wait(mut self) -> WebSocketPoolConnectOutcome {
|
||||
loop {
|
||||
let outcome = *self.receiver.borrow_and_update();
|
||||
@@ -534,34 +600,50 @@ pub(crate) struct WebSocketPoolConnectLease {
|
||||
pool: CodexWebSocketPool,
|
||||
key: CodexWebSocketPoolKey,
|
||||
id: Uuid,
|
||||
started_at: Instant,
|
||||
outcome: watch::Sender<WebSocketPoolConnectOutcome>,
|
||||
cancellation: CancellationToken,
|
||||
// slot 分配时即注册,封闭 acquire 与后台 task spawn 之间的 shutdown 竞态。
|
||||
_task_registration: TaskTrackerToken,
|
||||
armed: bool,
|
||||
}
|
||||
|
||||
pub(crate) struct WebSocketPoolConnected {
|
||||
pub(crate) connection: Box<PooledWebSocketConnection>,
|
||||
pub(crate) lease: WebSocketPoolLease,
|
||||
}
|
||||
|
||||
impl WebSocketPoolConnectLease {
|
||||
fn reserve(pool: CodexWebSocketPool, key: CodexWebSocketPoolKey) -> Self {
|
||||
let (outcome, _) = watch::channel(WebSocketPoolConnectOutcome::Pending);
|
||||
let cancellation = pool.shutdown.child_token();
|
||||
let task_registration = pool.tasks.token();
|
||||
Self {
|
||||
pool,
|
||||
key,
|
||||
id: Uuid::new_v4(),
|
||||
started_at: Instant::now(),
|
||||
outcome,
|
||||
cancellation,
|
||||
_task_registration: task_registration,
|
||||
armed: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn connected(
|
||||
pub(crate) fn started_at(&self) -> Instant {
|
||||
self.started_at
|
||||
}
|
||||
|
||||
pub(crate) fn key(&self) -> &CodexWebSocketPoolKey {
|
||||
&self.key
|
||||
}
|
||||
|
||||
pub(crate) fn cancellation_token(&self) -> CancellationToken {
|
||||
self.cancellation.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn connected_idle(
|
||||
mut self,
|
||||
connection: PooledWebSocketConnection,
|
||||
) -> Result<WebSocketPoolConnected, Box<PooledWebSocketConnection>> {
|
||||
) -> Result<(), Box<PooledWebSocketConnection>> {
|
||||
let result = self
|
||||
.pool
|
||||
.finish_connect(&self.key, self.id, connection)
|
||||
.finish_connect_idle(&self.key, self.id, connection)
|
||||
.await;
|
||||
let outcome = if result.is_ok() {
|
||||
WebSocketPoolConnectOutcome::Ready
|
||||
@@ -574,6 +656,7 @@ impl WebSocketPoolConnectLease {
|
||||
}
|
||||
|
||||
pub(crate) async fn failed(mut self) {
|
||||
self.cancellation.cancel();
|
||||
self.pool.fail_connect(&self.key, self.id).await;
|
||||
self.outcome
|
||||
.send_replace(WebSocketPoolConnectOutcome::Failed);
|
||||
@@ -586,6 +669,7 @@ impl Drop for WebSocketPoolConnectLease {
|
||||
if !self.armed {
|
||||
return;
|
||||
}
|
||||
self.cancellation.cancel();
|
||||
self.outcome
|
||||
.send_replace(WebSocketPoolConnectOutcome::Failed);
|
||||
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
|
||||
|
||||
@@ -19,7 +19,7 @@ fn new_chain_request(conversation_id: &str) -> CodexResponsesRequest {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cold_websocket_should_fall_back_to_same_account_http_after_fast_budget() {
|
||||
async fn cold_websocket_should_fall_back_without_recording_a_successful_connect() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
@@ -32,12 +32,13 @@ async fn cold_websocket_should_fall_back_to_same_account_http_after_fast_budget(
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut http).await;
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::default());
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::new(CodexWebSocketPool::default()))
|
||||
.with_websocket_pool(Arc::clone(&pool))
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(30));
|
||||
|
||||
let response = backend
|
||||
@@ -52,10 +53,251 @@ async fn cold_websocket_should_fall_back_to_same_account_http_after_fast_budget(
|
||||
assert_eq!(response.transport, CodexBackendTransport::HttpSse);
|
||||
assert_eq!(
|
||||
response.transport_metrics.decision,
|
||||
Some(CodexTransportDecision::Http2WebSocketSlow)
|
||||
Some(CodexTransportDecision::Http2WebSocketBudgetExhausted)
|
||||
);
|
||||
assert_eq!(response.transport_metrics.ws_connect_ms, None);
|
||||
assert!(response.transport_metrics.first_event_ms.is_some());
|
||||
assert!(response.body.contains("response.completed"));
|
||||
pool.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_request() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let (websocket_ready_tx, websocket_ready_rx) = tokio::sync::oneshot::channel();
|
||||
let server = tokio::spawn(async move {
|
||||
let (websocket_stream, _) = listener.accept().await.unwrap();
|
||||
let websocket_server = tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_millis(60)).await;
|
||||
let mut websocket = accept_codex_test_websocket(websocket_stream).await;
|
||||
websocket_ready_tx.send(()).unwrap();
|
||||
let _payload = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_background_ready", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
let (mut http, _) = listener.accept().await.unwrap();
|
||||
let request = read_http_request(&mut http).await;
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut http).await;
|
||||
websocket_server.await.unwrap();
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::default());
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool))
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(30));
|
||||
let request = new_chain_request("conversation-background-ready");
|
||||
|
||||
let first = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_background_http", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("foreground request should use HTTP after its fast-path budget");
|
||||
websocket_ready_rx.await.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
let second = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_background_reuse", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("next request should reuse the background websocket");
|
||||
server.await.unwrap();
|
||||
|
||||
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
|
||||
assert_eq!(first.transport_metrics.ws_connect_ms, None);
|
||||
assert_eq!(second.transport, CodexBackendTransport::WebSocket);
|
||||
assert_eq!(
|
||||
second.transport_metrics.decision,
|
||||
Some(CodexTransportDecision::ReusedWebSocket)
|
||||
);
|
||||
assert!(second.body.contains("resp_background_ready"));
|
||||
pool.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_websocket_opening_should_keep_the_original_fast_path_deadline() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stalled_websocket, _) = listener.accept().await.unwrap();
|
||||
let opening = read_http_request(&mut stalled_websocket).await;
|
||||
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
|
||||
|
||||
for _ in 0..2 {
|
||||
let (mut http, _) = listener.accept().await.unwrap();
|
||||
let request = read_http_request(&mut http).await;
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut http).await;
|
||||
}
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::default());
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool))
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(200));
|
||||
let request = new_chain_request("conversation-original-deadline");
|
||||
|
||||
let first = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_original_deadline_first", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("first request should use HTTP after the opening budget");
|
||||
let second = timeout(
|
||||
Duration::from_millis(100),
|
||||
backend.create_response(
|
||||
&request,
|
||||
request_context("req_original_deadline_second", Some("chatgpt-account")),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("shared opening must not grant the second request another full budget")
|
||||
.expect("second request should use HTTP");
|
||||
server.await.unwrap();
|
||||
|
||||
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
|
||||
assert_eq!(second.transport, CodexBackendTransport::HttpSse);
|
||||
assert_eq!(
|
||||
second.transport_metrics.decision,
|
||||
Some(CodexTransportDecision::Http2WebSocketBudgetExhausted)
|
||||
);
|
||||
pool.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_eviction_should_cancel_a_background_websocket_opening() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stale_opening, _) = listener.accept().await.unwrap();
|
||||
let opening = read_http_request(&mut stale_opening).await;
|
||||
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
|
||||
|
||||
let (mut http, _) = listener.accept().await.unwrap();
|
||||
let request = read_http_request(&mut http).await;
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut http).await;
|
||||
|
||||
let mut byte = [0_u8; 1];
|
||||
let read = timeout(Duration::from_secs(1), stale_opening.read(&mut byte))
|
||||
.await
|
||||
.expect("account eviction should close the opening socket")
|
||||
.unwrap();
|
||||
assert_eq!(read, 0);
|
||||
|
||||
let (fresh_stream, _) = listener.accept().await.unwrap();
|
||||
let mut websocket = accept_codex_test_websocket(fresh_stream).await;
|
||||
let _payload = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_after_eviction", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool))
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(30));
|
||||
let request = new_chain_request("conversation-evict-opening");
|
||||
|
||||
let first = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_evict_opening_first", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("first request should use HTTP after the opening budget");
|
||||
pool.evict_account("chatgpt-account").await;
|
||||
let second = timeout(
|
||||
Duration::from_secs(2),
|
||||
backend
|
||||
.clone()
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(500))
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_evict_opening_second", Some("chatgpt-account")),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("fresh websocket request should finish")
|
||||
.expect("the next request should build a fresh websocket");
|
||||
timeout(Duration::from_secs(2), server)
|
||||
.await
|
||||
.expect("test server should finish")
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
|
||||
assert_eq!(second.transport, CodexBackendTransport::WebSocket);
|
||||
assert!(second.body.contains("resp_after_eviction"));
|
||||
timeout(Duration::from_secs(2), pool.shutdown())
|
||||
.await
|
||||
.expect("pool shutdown should finish");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stalled_websocket, _) = listener.accept().await.unwrap();
|
||||
let opening = read_http_request(&mut stalled_websocket).await;
|
||||
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
|
||||
|
||||
let (mut http, _) = listener.accept().await.unwrap();
|
||||
let request = read_http_request(&mut http).await;
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut http).await;
|
||||
|
||||
let mut byte = [0_u8; 1];
|
||||
let read = timeout(Duration::from_secs(1), stalled_websocket.read(&mut byte))
|
||||
.await
|
||||
.expect("pool shutdown should close the opening socket")
|
||||
.unwrap();
|
||||
assert_eq!(read, 0);
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool))
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(30));
|
||||
|
||||
backend
|
||||
.create_response(
|
||||
&new_chain_request("conversation-shutdown-opening"),
|
||||
request_context("req_shutdown_opening", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("foreground request should use HTTP after the opening budget");
|
||||
timeout(Duration::from_secs(1), pool.shutdown())
|
||||
.await
|
||||
.expect("shutdown should join the cancelled opening task");
|
||||
server.await.unwrap();
|
||||
|
||||
assert!(pool.is_shutdown().await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -399,6 +641,145 @@ async fn origin_breaker_should_open_then_allow_only_one_half_open_probe() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn origin_breaker_should_count_hard_opening_failures() {
|
||||
let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig {
|
||||
failure_threshold: 1,
|
||||
failure_window: Duration::from_secs(1),
|
||||
open_duration: Duration::from_secs(1),
|
||||
});
|
||||
let WebSocketOriginBreakerDecision::Allowed(permit) =
|
||||
breaker.try_acquire("https://example.test:443")
|
||||
else {
|
||||
panic!("closed breaker should allow an opening");
|
||||
};
|
||||
permit.fail();
|
||||
|
||||
assert!(matches!(
|
||||
breaker.try_acquire("https://example.test:443"),
|
||||
WebSocketOriginBreakerDecision::Open
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let (late_failure_tx, late_failure_rx) = tokio::sync::oneshot::channel();
|
||||
let (second_used_websocket_tx, second_used_websocket_rx) = tokio::sync::oneshot::channel();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut delayed_opening, _) = listener.accept().await.unwrap();
|
||||
let opening = read_http_request(&mut delayed_opening).await;
|
||||
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
|
||||
let delayed_failure = tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_millis(80)).await;
|
||||
delayed_opening
|
||||
.write_all(
|
||||
b"HTTP/1.1 503 Service Unavailable\r\ncontent-length: 0\r\nconnection: close\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
late_failure_tx.send(()).unwrap();
|
||||
});
|
||||
|
||||
let (mut first_http, _) = listener.accept().await.unwrap();
|
||||
let request = read_http_request(&mut first_http).await;
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut first_http).await;
|
||||
delayed_failure.await.unwrap();
|
||||
|
||||
let (second_stream, _) = listener.accept().await.unwrap();
|
||||
let second_is_websocket = timeout(Duration::from_secs(1), async {
|
||||
let mut prefix = [0_u8; 4];
|
||||
loop {
|
||||
let read = second_stream.peek(&mut prefix).await.unwrap();
|
||||
if read == 0 {
|
||||
return false;
|
||||
}
|
||||
if read == prefix.len() {
|
||||
return prefix == *b"GET ";
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("second request should reach the upstream");
|
||||
if second_is_websocket {
|
||||
second_used_websocket_tx.send(true).unwrap();
|
||||
let mut websocket = accept_codex_test_websocket(second_stream).await;
|
||||
let _payload = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_after_late_failure", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
second_used_websocket_tx.send(false).unwrap();
|
||||
let mut second_http = second_stream;
|
||||
let request = read_http_request(&mut second_http).await;
|
||||
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
|
||||
write_completed_sse_response(&mut second_http).await;
|
||||
}
|
||||
});
|
||||
let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig {
|
||||
failure_threshold: 2,
|
||||
failure_window: Duration::from_secs(1),
|
||||
open_duration: Duration::from_secs(1),
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool))
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(30))
|
||||
.with_websocket_origin_breaker(breaker);
|
||||
|
||||
let first = backend
|
||||
.create_response(
|
||||
&new_chain_request("conversation-late-failure-first"),
|
||||
request_context("req_late_failure_first", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("first request should use HTTP after the opening budget");
|
||||
timeout(Duration::from_secs(2), late_failure_rx)
|
||||
.await
|
||||
.expect("delayed opening should receive its 503 response")
|
||||
.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
let second = timeout(
|
||||
Duration::from_secs(2),
|
||||
backend
|
||||
.clone()
|
||||
.with_websocket_fast_path_budget(Duration::from_millis(500))
|
||||
.create_response(
|
||||
&new_chain_request("conversation-late-failure-second"),
|
||||
request_context("req_late_failure_second", Some("chatgpt-account")),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("second request should finish")
|
||||
.expect("one degraded opening must not open a threshold-two breaker");
|
||||
let second_used_websocket = timeout(Duration::from_secs(2), second_used_websocket_rx)
|
||||
.await
|
||||
.expect("test server should observe the second transport")
|
||||
.unwrap();
|
||||
timeout(Duration::from_secs(2), server)
|
||||
.await
|
||||
.expect("test server should finish")
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
|
||||
assert!(second_used_websocket);
|
||||
assert_eq!(second.transport, CodexBackendTransport::WebSocket);
|
||||
assert!(second.body.contains("resp_after_late_failure"));
|
||||
timeout(Duration::from_secs(2), pool.shutdown())
|
||||
.await
|
||||
.expect("pool shutdown should finish");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn half_open_upstream_response_should_close_origin_breaker() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
|
||||
@@ -162,7 +162,7 @@ async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocke
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_stream_should_keep_fresh_socket_after_structural_activity() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
@@ -550,7 +550,7 @@ async fn websocket_pool_should_use_http_while_exact_key_is_busy() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
#[tokio::test]
|
||||
async fn websocket_pool_should_release_slot_when_client_drops_stream() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
@@ -816,7 +816,7 @@ async fn websocket_pool_should_replace_idle_connection_after_pong_deadline() {
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
#[tokio::test]
|
||||
async fn websocket_pool_should_gc_expired_idle_connections() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
@@ -890,7 +890,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() {
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
#[tokio::test]
|
||||
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();
|
||||
|
||||
+18
-8
@@ -241,7 +241,7 @@ Controller 之间不直接调用。跨功能优先级只在 `ControllerSet` 声
|
||||
|
||||
## 上游双热传输
|
||||
|
||||
Responses transport 使用精确 WebSocket 与共享 HTTP/2 两条常热通道。WebSocket 保存 connection-local response 状态;HTTP/2 提供新链和可持久化续链的低延迟后备。
|
||||
Responses transport 使用会话级 WebSocket 池与 origin 级 HTTP/2 池两条常热通道。WebSocket 保存 connection-local response 状态;HTTP/2 提供新链和可持久化续链的低延迟后备。transport 只竞速连接准备,不并行发送业务 payload。
|
||||
|
||||
### TransportRequirement
|
||||
|
||||
@@ -262,11 +262,14 @@ Responses transport 使用精确 WebSocket 与共享 HTTP/2 两条常热通道
|
||||
|
||||
`PreparedResponseTransport` 只表示“账号、header 和连接已准备,但 payload 尚未发送”。WebSocket opening 与 `response.create` 发送是单向类型边界:
|
||||
|
||||
- 普通新链、persisted continuation 和 external unknown 的冷 WS opening 预算为 800ms。
|
||||
- 预算内连接成功或复用热 socket 时发送 WS payload。
|
||||
- 预算耗尽、DNS/TCP/TLS 失败、pool 不可用或 breaker 打开时,同账号使用 HTTP/2。
|
||||
- 热 WS 可用时立即租用并发送 WS payload。
|
||||
- 普通新链、persisted continuation 和 external unknown 没有热 WS 时创建一个 `Connecting` slot;前台从该 slot 的统一起点最多等待 800ms。
|
||||
- 800ms 内握手成功时,连接先原子进入 `Idle`,当前请求再租用为 `Busy` 并发送 WS payload。
|
||||
- 800ms 耗尽时,当前请求不发送 WS payload,立即在同账号使用 HTTP/2;未完成的握手由 pool 在后台继续,仍受 15 秒硬连接超时约束。
|
||||
- 后台握手成功后只进入 `Idle`,供后续请求复用;最终失败则删除 `Connecting` slot。相同 key 的后续请求沿用原 opening 的绝对 deadline,不会重新获得一段 800ms。
|
||||
- DNS/TCP/TLS 立即失败、pool 不可用或 breaker 打开时,同账号使用 HTTP/2。
|
||||
- opening 已返回明确上游响应时,事实直接进入 controller,不追加 HTTP 请求。
|
||||
- warmup 与 exact continuation 禁止 HTTP fallback。
|
||||
- warmup 与 exact continuation 必须等待所需 WS,不使用 HTTP fallback。
|
||||
- payload 发送后的连接断开、发送超时或首事件超时统一为 `PostSendAmbiguous`,不得换 transport 或账号。
|
||||
|
||||
AttemptRunner 将 prepared transport 与账号请求间隔并行,800ms、熔断阈值和连接状态只存在于 `upstream/openai/transport`。
|
||||
@@ -277,21 +280,26 @@ Pool key 为 `(base_url, account_id, local_conversation_id)`,slot 状态为:
|
||||
|
||||
- `Idle`:可租用的热 socket。
|
||||
- `Busy`:单 socket 上已有一个 in-flight response。
|
||||
- `Connecting`:同 key opening 单飞,其他请求不会再建第二条 socket。
|
||||
- `Connecting`:同 key opening 单飞,包含统一 started-at、结果通知和取消令牌。
|
||||
|
||||
每条连接记录最新 completed response ID。Exact continuation 只有 ID 完全匹配时才能租用;socket 缺失、busy、失活或 ID 不匹配时不新建连接碰运气,由 HistoryController 决定 full replay 或失败。
|
||||
|
||||
池化 socket 默认每 25 秒发送带唯一序列的 Ping,匹配 Pong 的 deadline 为 5 秒,最大寿命为 55 分钟。任意业务入站帧同样证明链路存活;下游背压导致 pump 无法读取 socket 时暂停主动探活。后台 pump 统一处理 ping/pong、EOF、close 和 liveness,复用前只读取原子连接状态,不增加网络往返。
|
||||
|
||||
`Idle`、`Busy` 和 `Connecting` 都计入 `ws_pool.max_per_account`,默认单账号最多 8 个 slot;后台续建不设置另一套队列或全局连接上限。账号驱逐会删除该账号全部 slot、关闭 idle socket 并取消 connecting;busy socket 在归还时因 reservation 已失效而关闭。进程关闭先取消所有 pool-owned opening,再等待维护任务和建连任务退出,迟到连接不能在驱逐或 shutdown 后重新入池。
|
||||
|
||||
### Origin breaker
|
||||
|
||||
Breaker key 为 origin 与 TLS profile,不按账号区分:
|
||||
|
||||
- 30 秒窗口内 3 次 WS 快路径超时后打开 30 秒。
|
||||
- 30 秒窗口内 3 次退化 opening 后打开 30 秒;退化包括 800ms 快路径预算耗尽、transport 失败和 5xx opening 响应。
|
||||
- 同一次 opening 的前台预算耗尽与后台最终失败只计一次;迟到成功也不抹去已经发生的快路径退化。
|
||||
- 预算内成功 opening 会清空该 origin 的历史退化记录。
|
||||
- 打开期间热 WS 仍可使用;普通冷 opening 直接走 HTTP/2。
|
||||
- 到期后只允许一个 half-open probe。
|
||||
- 账号或请求级 4xx opening 响应证明 origin 可达并关闭 half-open。
|
||||
- 5xx 或 transport probe 失败重新打开 breaker。
|
||||
- 账号驱逐和进程 shutdown 触发的主动取消不计为 origin 故障。
|
||||
|
||||
Breaker 不更新账号状态;账号失效只由 controller 的明确 typed fact 触发。
|
||||
|
||||
@@ -322,7 +330,9 @@ HTTP/2 连接按 origin/TLS profile 共享,账号 Authorization 只属于单
|
||||
- `httpVersion`
|
||||
- `websocketPool.kind`
|
||||
|
||||
稳定 decision 包括 `ws_reused`、`ws_connected_fast`、`ws_exact_required`、`ws_required`、`http2_ws_slow`、`http2_breaker_open`、`http2_pool_unavailable` 和 `http2_ws_pre_send_failure`。
|
||||
`wsConnectMs` 只表示已经成功的 WS 握手耗时;复用、预算耗尽和其他 HTTP fallback 不填写该字段。`transportDecisionWaitMs` 表示当前请求为 transport 选择实际等待的时间。后台 opening 另外记录 `ws_preconnect_duration_ms`、`ws_preconnect_outcome` 和 `foreground_waiting`,不把未成功握手的 800ms 伪装成连接耗时。
|
||||
|
||||
稳定 decision 包括 `ws_reused`、`ws_connected_fast`、`ws_exact_required`、`ws_required`、`http2_ws_budget_exhausted`、`http2_breaker_open`、`http2_pool_unavailable` 和 `http2_ws_pre_send_failure`。
|
||||
|
||||
## 账号调度与失效
|
||||
|
||||
|
||||
Reference in New Issue
Block a user