diff --git a/src/maestro/me_startup.rs b/src/maestro/me_startup.rs index 4fe5273..3ad0c85 100644 --- a/src/maestro/me_startup.rs +++ b/src/maestro/me_startup.rs @@ -22,6 +22,9 @@ use crate::transport::middle_proxy::MePool; use super::generation::RuntimeTaskScope; use super::helpers::load_startup_proxy_config_snapshot; +// Cleanup capacity follows the configured client admission semantics. +mod capacity; + // Restarting supervisors for long-lived ME maintenance tasks. mod supervisor; use supervisor::spawn_me_supervisors; @@ -299,7 +302,7 @@ pub(crate) async fn initialize_me_pool( config.general.me_route_blocking_send_timeout_ms, config.general.me_route_inline_recovery_attempts, config.general.me_route_inline_recovery_wait_ms, - (config.server.max_connections as usize).saturating_add(128), + capacity::me_connection_cleanup_capacity(config.server.max_connections), ); startup_tracker .complete_component( diff --git a/src/maestro/me_startup/capacity.rs b/src/maestro/me_startup/capacity.rs new file mode 100644 index 0000000..7da0c6d --- /dev/null +++ b/src/maestro/me_startup/capacity.rs @@ -0,0 +1,57 @@ +/// Returns the cleanup reservation budget for simultaneously admitted ME clients. +pub(super) fn me_connection_cleanup_capacity(max_connections: u32) -> usize { + if max_connections == 0 { + // Match frontend admission; adding headroom here would exceed Tokio's limit. + tokio::sync::Semaphore::MAX_PERMITS + } else { + (max_connections as usize) + .saturating_add(128) + .min(tokio::sync::Semaphore::MAX_PERMITS) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::time::Duration; + use tokio::sync::Semaphore; + use crate::transport::middle_proxy::{ConnRegistry, MeResponse}; + + #[test] + fn me_cleanup_capacity_preserves_zero_and_finite_limits() { + for (configured, expected) in [ + (0, Semaphore::MAX_PERMITS), + (1, 129), + (128, 256), + (10_000, 10_128), + (u32::MAX, (u32::MAX as usize).saturating_add(128).min(Semaphore::MAX_PERMITS)), + ] { + let capacity = me_connection_cleanup_capacity(configured); + assert_eq!(capacity, expected); + assert!(capacity > 0 && capacity <= Semaphore::MAX_PERMITS); + } + } + + #[tokio::test] + async fn me_cleanup_capacity_zero_admits_512_simultaneous_leased_routes() { + let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity( + 8, me_connection_cleanup_capacity(0), + )); + let mut clients = Vec::new(); + for index in 0..512 { + let client = tokio::time::timeout(Duration::from_secs(1), registry.register_leased()) + .await.unwrap_or_else(|_| panic!("registration {} blocked", index + 1)) + .unwrap(); + clients.push(client); + } + for (lease, mut receiver) in clients { + registry.route_nowait(lease.conn_id(), MeResponse::Ack(1)).await; + assert!(matches!(receiver.try_recv(), Ok(MeResponse::Ack(1)))); + lease.unregister().await; + } + // All reservation capacity must remain reusable after explicit cleanup. + let (lease, _) = registry.register_leased().await.unwrap(); + lease.unregister().await; + } +} diff --git a/src/transport/middle_proxy/pool_lifecycle.rs b/src/transport/middle_proxy/pool_lifecycle.rs index 7734bff..c38a433 100644 --- a/src/transport/middle_proxy/pool_lifecycle.rs +++ b/src/transport/middle_proxy/pool_lifecycle.rs @@ -342,3 +342,7 @@ mod tests { assert!(lifecycle.spawn_producer(async {}).is_err()); } } + +// Cleanup reservations must remain cancellation-safe under capacity pressure. +#[cfg(test)] +mod connection_tests; diff --git a/src/transport/middle_proxy/pool_lifecycle/connection_tests.rs b/src/transport/middle_proxy/pool_lifecycle/connection_tests.rs new file mode 100644 index 0000000..7cdc996 --- /dev/null +++ b/src/transport/middle_proxy/pool_lifecycle/connection_tests.rs @@ -0,0 +1,20 @@ +use super::*; +use std::future::{Future, poll_fn}; +use std::task::Poll; +use super::super::pool_writer_security_tests::make_pool; +use super::super::ConnRegistry; + +#[tokio::test] +async fn shutdown_cancels_registration_waiting_for_cleanup_capacity() { + let mut pool = make_pool().await; + Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().routing).unwrap().registry = + Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1)); + let (lease, _receiver) = pool.register_connection().await.unwrap(); + let mut waiting = Box::pin(pool.register_connection()); + assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending)); + pool.begin_shutdown(); + assert!(tokio::time::timeout(Duration::from_secs(1), waiting).await.unwrap().is_none()); + lease.unregister().await; + assert!(pool.shutdown_until(Duration::from_secs(1)).await); + assert!(pool.register_connection().await.is_none()); +} diff --git a/src/transport/middle_proxy/registry.rs b/src/transport/middle_proxy/registry.rs index 1655c4b..24493d3 100644 --- a/src/transport/middle_proxy/registry.rs +++ b/src/transport/middle_proxy/registry.rs @@ -331,3 +331,7 @@ impl ConnRegistry { #[cfg(test)] mod tests; + +// Cleanup reservations must remain cancellation-safe under capacity pressure. +#[cfg(test)] +mod lease_capacity_tests; diff --git a/src/transport/middle_proxy/registry/lease_capacity_tests.rs b/src/transport/middle_proxy/registry/lease_capacity_tests.rs new file mode 100644 index 0000000..10c7f6b --- /dev/null +++ b/src/transport/middle_proxy/registry/lease_capacity_tests.rs @@ -0,0 +1,54 @@ +use super::*; +use std::future::{Future, poll_fn}; +use std::task::Poll; +use std::time::Duration; + +#[tokio::test] +async fn finite_cleanup_capacity_waits_until_dropped_route_is_received() { + let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 128)); + let mut cleanup = registry.take_cleanup_receiver().unwrap(); + let mut clients = Vec::new(); + for _ in 0..128 { + clients.push(registry.register_leased().await.unwrap()); + } + assert_eq!(registry.routing.map.len(), 128); + assert_eq!(registry.cleanup_tx.capacity(), 0); + assert!(matches!(cleanup.try_recv(), Err(mpsc::error::TryRecvError::Empty))); + let mut waiting = Box::pin(registry.register_leased()); + assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending)); + let retired = clients.last().unwrap().0.conn_id(); + drop(clients.pop().unwrap()); + assert!(!registry.routing.map.contains_key(&retired)); + assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending)); + assert_eq!(cleanup.recv().await, Some(retired)); + registry.unregister(retired).await; + let (lease, _) = tokio::time::timeout(Duration::from_secs(1), waiting).await.unwrap().unwrap(); + lease.unregister().await; + for (lease, _) in clients { + lease.unregister().await; + } + assert!(registry.routing.map.is_empty()); + assert!(registry.routing.byte_budget.is_empty()); + assert_eq!(registry.cleanup_tx.capacity(), 128); +} + +#[tokio::test] +async fn cancelled_registration_leaks_neither_route_nor_cleanup_permit() { + let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1)); + let _cleanup = registry.take_cleanup_receiver().unwrap(); + let (lease, _receiver) = registry.register_leased().await.unwrap(); + for _ in 0..64 { + let mut waiting = Box::pin(registry.register_leased()); + assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending)); + drop(waiting); + assert_eq!(registry.routing.map.len(), 1); + assert_eq!(registry.cleanup_tx.capacity(), 0); + } + lease.unregister().await; + assert!(registry.routing.map.is_empty()); + assert_eq!(registry.cleanup_tx.capacity(), 1); + let (lease, _) = tokio::time::timeout(Duration::from_secs(1), registry.register_leased()) + .await.unwrap().unwrap(); + lease.unregister().await; + assert_eq!(registry.cleanup_tx.capacity(), 1); +}