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

fix(openai): remove websocket pool slot caps

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-08-25 11:22:07 +08:00
co-authored by Codex
parent b93b8859f9
commit 75f9d2a290
9 changed files with 94 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/config.rs
+1 -12
View File
@@ -155,8 +155,6 @@ impl CodexApiConfig {
pub struct CodexWebSocketPoolSettings {
pub enabled: bool,
pub max_age_ms: u64,
pub max_per_account: usize,
pub max_total: usize,
pub max_connecting: usize,
pub initial_event_timeout_ms: u64,
}
@@ -166,8 +164,6 @@ impl Default for CodexWebSocketPoolSettings {
Self {
enabled: true,
max_age_ms: 55 * 60 * 1000,
max_per_account: 8,
max_total: 64,
max_connecting: 8,
initial_event_timeout_ms: 20_000,
}
@@ -176,12 +172,7 @@ impl Default for CodexWebSocketPoolSettings {
impl CodexWebSocketPoolSettings {
fn validate(&self) -> Result<(), OpenAiConfigError> {
if self.max_age_ms == 0
|| self.max_per_account == 0
|| self.max_total == 0
|| self.max_connecting == 0
|| self.max_connecting > self.max_total
{
if self.max_age_ms == 0 || self.max_connecting == 0 {
return Err(OpenAiConfigError::InvalidField("openai.ws_pool"));
}
Ok(())
@@ -191,8 +182,6 @@ impl CodexWebSocketPoolSettings {
CodexWebSocketPoolConfig {
enabled: self.enabled,
max_age: Duration::from_millis(self.max_age_ms),
max_per_account: self.max_per_account,
max_total: self.max_total,
max_connecting: self.max_connecting,
initial_event_timeout: (self.initial_event_timeout_ms != 0)
.then(|| Duration::from_millis(self.initial_event_timeout_ms)),
backend/crates/providers/openai/src/transport/websocket/pool/mod.rs
+20 -48
View File
@@ -10,13 +10,13 @@ use std::{
time::Duration,
};
use tokio::{sync::Semaphore, time::Instant};
use tokio::sync::Semaphore;
use tokio_util::{sync::CancellationToken, task::TaskTracker};
use uuid::Uuid;
use self::state::{
WebSocketPoolConnecting, WebSocketPoolSlot, WebSocketPoolState, account_slot_count,
close_pooled_connection, close_pooled_connections, take_lru_idle_connection,
WebSocketPoolConnecting, WebSocketPoolSlot, WebSocketPoolState, close_pooled_connection,
close_pooled_connections,
};
use super::pump::PumpKeepalive;
@@ -31,8 +31,6 @@ pub(crate) use self::{
},
};
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAX_TOTAL: usize = 64;
const DEFAULT_MAX_CONNECTING: usize = 16;
const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
@@ -64,10 +62,6 @@ pub struct CodexWebSocketPoolConfig {
pub enabled: bool,
/// 单个 socket 的最大生命周期。
pub max_age: Duration,
/// 单个账号允许占用的最大池 slot 数。
pub max_per_account: usize,
/// 所有账号合计允许占用的最大池 slot 数。
pub max_total: usize,
/// 所有账号合计允许并发执行的 opening 数。
pub max_connecting: usize,
/// 后台维护间隔；`None` 表示不启动后台任务。
@@ -87,8 +81,6 @@ impl Default for CodexWebSocketPoolConfig {
Self {
enabled: true,
max_age: DEFAULT_MAX_AGE,
max_per_account: DEFAULT_MAX_PER_ACCOUNT,
max_total: DEFAULT_MAX_TOTAL,
max_connecting: DEFAULT_MAX_CONNECTING,
maintenance_interval: Some(DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(DEFAULT_PING_INTERVAL),
@@ -113,10 +105,9 @@ impl CodexWebSocketPoolConfig {
}
impl CodexWebSocketPool {
/// 构造连接池策略和状态。
pub fn new(max_per_account: usize, max_age: Duration) -> Self {
/// 构造不限制累计 slot 数的连接池策略和状态。
pub fn new(max_age: Duration) -> Self {
Self::with_config(CodexWebSocketPoolConfig {
max_per_account: max_per_account.max(1),
max_age,
maintenance_interval: None,
ping_interval: None,
@@ -218,43 +209,25 @@ impl CodexWebSocketPool {
let acquire = if required_response_id.is_some() {
WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::ContinuationNotFound)
} else if self.config.max_per_account == 0
|| account_slot_count(&state.slots, key.account_id()) >= self.config.max_per_account
{
WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap)
} else {
let connect_permit = self.connect_semaphore.clone().try_acquire_owned().ok();
match connect_permit {
Some(connect_permit) => {
let has_total_capacity = if self.config.max_total == 0 {
false
} else if state.slots.len() < self.config.max_total {
true
} else if let Some(connection) = take_lru_idle_connection(&mut state) {
connections_to_close.push(connection);
true
} else {
false
};
if has_total_capacity {
let lease = WebSocketPoolConnectLease::reserve(
self.clone(),
key.clone(),
connect_permit,
);
state.slots.insert(
key,
WebSocketPoolSlot::Connecting(WebSocketPoolConnecting {
id: lease.id,
started_at: lease.started_at,
outcome: lease.outcome.clone(),
cancellation: lease.cancellation.clone(),
}),
);
WebSocketPoolAcquire::Connect(lease)
} else {
WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap)
}
let lease = WebSocketPoolConnectLease::reserve(
self.clone(),
key.clone(),
connect_permit,
);
state.slots.insert(
key,
WebSocketPoolSlot::Connecting(WebSocketPoolConnecting {
id: lease.id,
started_at: lease.started_at,
outcome: lease.outcome.clone(),
cancellation: lease.cancellation.clone(),
}),
);
WebSocketPoolAcquire::Connect(lease)
}
None => WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap),
}
@@ -292,7 +265,6 @@ impl CodexWebSocketPool {
key.clone(),
WebSocketPoolSlot::Idle {
connection: Box::new(connection),
last_used_at: Instant::now(),
},
);
}
backend/crates/providers/openai/src/transport/websocket/pool/state.rs
-29
View File
@@ -108,7 +108,6 @@ impl WebSocketContinuationState {
pub(super) enum WebSocketPoolSlot {
Idle {
connection: Box<PooledWebSocketConnection>,
last_used_at: Instant,
},
Busy(WebSocketPoolReservation),
Connecting(WebSocketPoolConnecting),
@@ -138,34 +137,6 @@ pub(super) struct WebSocketPoolReservation {
pub(super) latest_response_id: Option<String>,
}
pub(super) fn account_slot_count(
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
account_id: &str,
) -> usize {
slots
.keys()
.filter(|key| key.account_id() == account_id)
.count()
}
pub(super) fn take_lru_idle_connection(
state: &mut WebSocketPoolState,
) -> Option<PooledWebSocketConnection> {
let key = state
.slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle { last_used_at, .. } => Some((key, *last_used_at)),
WebSocketPoolSlot::Busy(_) | WebSocketPoolSlot::Connecting(_) => None,
})
.min_by_key(|(_, last_used_at)| *last_used_at)
.map(|(key, _)| key.clone())?;
let Some(WebSocketPoolSlot::Idle { connection, .. }) = state.slots.remove(&key) else {
return None;
};
Some(*connection)
}
pub(super) async fn close_pooled_connection(connection: PooledWebSocketConnection) {
connection.websocket.close().await;
}
backend/crates/providers/openai/tests/config.rs
-4
View File
@@ -101,8 +101,6 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
config.api.base_url.as_str(),
config.ws_pool.enabled,
config.ws_pool.max_age_ms,
config.ws_pool.max_per_account,
config.ws_pool.max_total,
config.ws_pool.max_connecting,
config.ws_pool.initial_event_timeout_ms,
config.quota.refresh_interval_minutes,
@@ -115,8 +113,6 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
true,
3_300_000,
8,
64,
8,
20_000,
15,
true,
backend/crates/providers/openai/tests/transport/headers.rs
+3 -3
View File
@@ -265,7 +265,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
format!("http://{address}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))));
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
let response = backend
.create_response(
@@ -480,7 +480,7 @@ async fn backend_websocket_should_drop_only_unrepresentable_opaque_header_values
format!("http://{address}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))));
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
client
.create_response(
@@ -932,7 +932,7 @@ async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_la
});
let profile = test_wire_profile();
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let client = CodexBackendClient::new(
reqwest::Client::builder()
.no_proxy()
backend/crates/providers/openai/tests/transport/latency.rs
+5 -5
View File
@@ -300,7 +300,7 @@ async fn account_eviction_should_cancel_a_background_websocket_opening() {
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -354,7 +354,7 @@ async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() {
.unwrap();
assert_eq!(read, 0);
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -733,7 +733,7 @@ async fn concurrent_multi_conversation_continuations_should_select_each_exact_so
tokio::join!(first_continuation, second_continuation);
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -1190,7 +1190,7 @@ async fn cancelled_half_open_opening_should_allow_another_probe() {
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -1325,7 +1325,7 @@ async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() {
write_completed_sse_response(&mut second_http).await;
}
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
backend/crates/providers/openai/tests/transport/mod.rs
-2
View File
@@ -495,8 +495,6 @@ fn websocket_pool_config_for_tests(
CodexWebSocketPoolConfig {
enabled: true,
max_age: Duration::from_mins(1),
max_per_account: 8,
max_total: 64,
max_connecting: 16,
maintenance_interval,
ping_interval,
backend/crates/providers/openai/tests/transport/websocket.rs
+3 -3
View File
@@ -126,7 +126,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_
websocket.send(Message::Text(payload.into())).await.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -932,7 +932,7 @@ async fn codex_backend_client_should_not_reuse_websocket_rate_limit_headers() {
.unwrap();
websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -1914,7 +1914,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_
let mut request = codex_request("gpt-5.5", "be brief", Vec::new());
request.set_previous_response_id(Some("resp_previous".to_string()));
request.previous_response_scope = Some(PreviousResponseScope::Persisted);
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
backend/crates/providers/openai/tests/transport/websocket_pool.rs
+62 -156
View File
@@ -35,7 +35,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and
}
websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -97,7 +97,7 @@ async fn exact_continuation_reuses_owning_socket_after_connection_profile_update
format!("http://{addr}"),
profile.clone(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(1, Duration::from_mins(1))));
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
let first_request = pooled_websocket_request("wire-profile-continuation");
let first = backend
.create_response(
@@ -401,7 +401,7 @@ async fn codex_backend_client_stream_should_keep_reused_socket_after_structural_
}
#[tokio::test]
async fn codex_backend_client_stream_should_use_http_when_pool_is_at_cap() {
async fn codex_backend_client_stream_should_use_http_when_connecting_limit_is_exhausted() {
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));
@@ -414,7 +414,7 @@ async fn codex_backend_client_stream_should_use_http_when_pool_is_at_cap() {
write_completed_sse_response(&mut first_stream).await;
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
max_per_account: 0,
max_connecting: 0,
initial_event_timeout: Some(Duration::from_millis(30)),
..websocket_pool_config_for_tests(None, None, None)
}));
@@ -432,7 +432,7 @@ async fn codex_backend_client_stream_should_use_http_when_pool_is_at_cap() {
request_context("req_structural_bypass", Some("chatgpt-account")),
)
.await
.expect("pool cap should select HTTP before sending payload");
.expect("exhausted connection-opening limit should select HTTP before sending payload");
assert_eq!(response.transport, CodexBackendTransport::HttpSse);
assert_eq!(
response.transport_metrics.decision,
@@ -472,7 +472,7 @@ async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_acc
websocket.close(None).await.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -841,140 +841,32 @@ async fn websocket_pool_shutdown_should_cancel_and_join_active_stream() {
}
#[tokio::test]
async fn websocket_pool_should_use_http_for_new_keys_after_account_cap() {
async fn websocket_pool_should_keep_sixty_fifth_conversation_on_same_account_on_websocket() {
const REQUEST_COUNT: usize = 65;
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 {
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_cap_first", 2, 1).into(),
))
.await
.unwrap();
for _ in 0..2 {
let (mut stream, _) = listener.accept().await.unwrap();
let mut websockets = Vec::with_capacity(REQUEST_COUNT);
for index in 0..REQUEST_COUNT {
let (stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let request = read_http_request(&mut stream).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut stream).await;
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
completed_websocket_response(&format!("resp_same_account_{index}"), 2, 1)
.into(),
))
.await
.unwrap();
websockets.push(websocket);
}
drop(first_websocket);
});
let pool = Arc::new(CodexWebSocketPool::new(1, Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
test_wire_profile(),
)
.with_websocket_pool(pool);
let first_request = pooled_websocket_request("conversation-cap-one");
let second_request = pooled_websocket_request("conversation-cap-two");
let first = backend
.create_response(
&first_request,
request_context("req_cap_first", Some("chatgpt-account")),
)
.await
.expect("first capped websocket response should succeed");
let second = backend
.create_response(
&second_request,
request_context("req_cap_second", Some("chatgpt-account")),
)
.await
.expect("new key over account cap should use HTTP");
let third = backend
.create_response(
&second_request,
request_context("req_cap_third", Some("chatgpt-account")),
)
.await
.expect("capped key should keep using HTTP instead of opening one-shot sockets");
server.await.unwrap();
assert!(first.body.contains("resp_cap_first"));
assert!(second.body.contains("response.completed"));
assert!(third.body.contains("response.completed"));
assert_eq!(first.websocket_pool_decision.unwrap().kind(), "new");
assert_eq!(
second.transport_metrics.decision,
Some(CodexTransportDecision::Http2PoolUnavailable)
);
assert_eq!(
third.transport_metrics.decision,
Some(CodexTransportDecision::Http2PoolUnavailable)
);
assert_eq!(accepted_connections.load(Ordering::SeqCst), 3);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn websocket_pool_should_bound_concurrent_openings_across_accounts() {
let (websockets, http) = concurrent_pool_transport_counts(8, 2).await;
assert_eq!((websockets, http), (2, 6));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn websocket_pool_should_bound_total_slots_across_accounts() {
let (websockets, http) = concurrent_pool_transport_counts(3, 8).await;
assert_eq!((websockets, http), (3, 5));
}
#[tokio::test]
async fn websocket_pool_should_evict_lru_idle_connection_at_total_cap() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first = accept_codex_test_websocket(first_stream).await;
let _ = first.next().await.unwrap().unwrap();
first
.send(Message::Text(
completed_websocket_response("resp_global_lru_first", 2, 1).into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second = accept_codex_test_websocket(second_stream).await;
let _ = second.next().await.unwrap().unwrap();
second
.send(Message::Text(
completed_websocket_response("resp_global_lru_second", 2, 1).into(),
))
.await
.unwrap();
let close = timeout(Duration::from_secs(2), first.next())
.await
.expect("global LRU eviction should close the oldest idle socket")
.expect("oldest idle socket should emit a close frame")
.expect("oldest idle socket close should be valid");
std::assert_matches!(close, Message::Close(_));
let (third_stream, _) = listener.accept().await.unwrap();
let mut third = accept_codex_test_websocket(third_stream).await;
let _ = third.next().await.unwrap().unwrap();
third
.send(Message::Text(
completed_websocket_response("resp_global_lru_third", 2, 1).into(),
))
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
max_per_account: 8,
max_total: 2,
max_connecting: 2,
max_connecting: 1,
..websocket_pool_config_for_tests(None, None, None)
}));
let backend = CodexBackendClient::new(
@@ -983,36 +875,52 @@ async fn websocket_pool_should_evict_lru_idle_connection_at_total_cap() {
test_wire_profile(),
)
.with_websocket_pool(Arc::clone(&pool));
let mut responses = Vec::new();
for index in 0..3 {
let request = pooled_websocket_request(&format!("conversation-global-lru-{index}"));
let account_id = format!("account-global-lru-{index}");
let request_id = format!("req_global_lru_{index}");
let mut responses = Vec::with_capacity(REQUEST_COUNT);
for index in 0..REQUEST_COUNT {
let request = pooled_websocket_request(&format!("conversation-same-account-{index}"));
responses.push(
backend
.create_response(&request, request_context(&request_id, Some(&account_id)))
.create_response(
&request,
request_context(
&format!("req_same_account_{index}"),
Some("chatgpt-account"),
),
)
.await
.expect("global LRU request should keep using websocket"),
.expect("same-account conversation should use websocket"),
);
}
server.await.unwrap();
assert!(responses[0].body.contains("resp_global_lru_first"));
assert!(responses[1].body.contains("resp_global_lru_second"));
assert!(responses[2].body.contains("resp_global_lru_third"));
assert!(
responses
.iter()
.all(|response| response.transport == CodexBackendTransport::WebSocket)
assert_eq!(
(
responses
.iter()
.filter(|response| response.transport == CodexBackendTransport::WebSocket)
.count(),
accepted_connections.load(Ordering::SeqCst),
),
(REQUEST_COUNT, REQUEST_COUNT)
);
pool.shutdown().await;
}
async fn concurrent_pool_transport_counts(
max_total: usize,
max_connecting: usize,
) -> (usize, usize) {
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn websocket_pool_should_bound_concurrent_openings_across_accounts() {
let (websockets, http) = concurrent_pool_transport_counts(2).await;
assert_eq!((websockets, http), (2, 6));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn websocket_pool_should_allow_all_openings_when_connect_permits_are_available() {
let (websockets, http) = concurrent_pool_transport_counts(8).await;
assert_eq!((websockets, http), (8, 0));
}
async fn concurrent_pool_transport_counts(max_connecting: usize) -> (usize, usize) {
const REQUEST_COUNT: usize = 8;
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -1051,7 +959,7 @@ async fn concurrent_pool_transport_counts(
websocket
.send(Message::Text(
completed_websocket_response(
&format!("resp_global_cap_{max_total}_{max_connecting}_{index}"),
&format!("resp_connecting_cap_{max_connecting}_{index}"),
2,
1,
)
@@ -1063,8 +971,6 @@ async fn concurrent_pool_transport_counts(
(websocket_count, http_count)
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
max_per_account: REQUEST_COUNT,
max_total,
max_connecting,
..websocket_pool_config_for_tests(None, None, None)
}));
@@ -1096,7 +1002,7 @@ async fn concurrent_pool_transport_counts(
let response = request
.await
.unwrap()
.expect("globally capped request should complete");
.expect("concurrent request should complete");
if response.transport == CodexBackendTransport::WebSocket {
websocket_responses += 1;
}
@@ -1910,7 +1816,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_error_termin
.unwrap();
second_websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
@@ -1987,7 +1893,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_unknown_resp
.unwrap();
second_websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),