diff --git a/crates/lanspread-peer/src/services/server.rs b/crates/lanspread-peer/src/services/server.rs index 39cde7a..77e22f1 100644 --- a/crates/lanspread-peer/src/services/server.rs +++ b/crates/lanspread-peer/src/services/server.rs @@ -1,6 +1,13 @@ //! QUIC server accept loop. -use std::{future::Future, net::SocketAddr, panic::AssertUnwindSafe, sync::Arc, time::Duration}; +use std::{ + collections::HashMap, + future::Future, + net::{IpAddr, SocketAddr}, + panic::AssertUnwindSafe, + sync::{Arc, Mutex}, + time::Duration, +}; use futures::FutureExt as _; use s2n_quic::{ @@ -42,10 +49,260 @@ const MAX_GLOBAL_BULK_TRANSFER_TASKS: usize = 48; /// Native archive extraction is substantially more expensive than ordinary /// catalog egress, so keep it behind a separate small global budget. const MAX_GLOBAL_STREAM_INSTALL_TASKS: usize = 2; +/// One observed host may retain only this many established connection scopes. +const MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN: usize = 8; +/// Bounds decoded and undecoded control work across every connection from one IP. +const MAX_CONTROL_STREAM_TASKS_PER_ORIGIN: usize = 16; +/// Preserves bulk capacity for other LAN origins while allowing two full +/// four-stream ordinary download windows from one host. +const MAX_BULK_TRANSFER_TASKS_PER_ORIGIN: usize = 8; +/// A single origin cannot occupy both native Stream Install provider slots. +const MAX_STREAM_INSTALL_TASKS_PER_ORIGIN: usize = 1; /// Application idle bound for an established connection with no active /// request streams. This is independent of transport keepalive traffic. const CONNECTION_NO_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(10); +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(super) struct ObservedOrigin(IpAddr); + +impl ObservedOrigin { + fn from_remote_addr(remote_addr: Option) -> Option { + remote_addr.map(|addr| Self(canonical_origin_ip(addr.ip()))) + } + + pub(super) const fn ip(self) -> IpAddr { + self.0 + } +} + +fn canonical_origin_ip(ip: IpAddr) -> IpAddr { + match ip { + IpAddr::V4(ip) => IpAddr::V4(ip), + IpAddr::V6(ip) => ip.to_ipv4_mapped().map_or(IpAddr::V6(ip), IpAddr::V4), + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum OriginAdmissionClass { + Connection, + Control, + Bulk, + StreamInstall, +} + +impl OriginAdmissionClass { + const fn limit(self) -> usize { + match self { + Self::Connection => MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN, + Self::Control => MAX_CONTROL_STREAM_TASKS_PER_ORIGIN, + Self::Bulk => MAX_BULK_TRANSFER_TASKS_PER_ORIGIN, + Self::StreamInstall => MAX_STREAM_INSTALL_TASKS_PER_ORIGIN, + } + } +} + +#[derive(Default)] +struct OriginUsage { + connections: usize, + controls: usize, + bulk: usize, + stream_installs: usize, +} + +impl OriginUsage { + const fn count(&self, class: OriginAdmissionClass) -> usize { + match class { + OriginAdmissionClass::Connection => self.connections, + OriginAdmissionClass::Control => self.controls, + OriginAdmissionClass::Bulk => self.bulk, + OriginAdmissionClass::StreamInstall => self.stream_installs, + } + } + + fn increment(&mut self, class: OriginAdmissionClass) { + let count = self.count_mut(class); + *count += 1; + } + + fn decrement(&mut self, class: OriginAdmissionClass) -> bool { + let count = self.count_mut(class); + if *count == 0 { + return false; + } + *count -= 1; + true + } + + const fn is_empty(&self) -> bool { + self.connections == 0 && self.controls == 0 && self.bulk == 0 && self.stream_installs == 0 + } + + fn count_mut(&mut self, class: OriginAdmissionClass) -> &mut usize { + match class { + OriginAdmissionClass::Connection => &mut self.connections, + OriginAdmissionClass::Control => &mut self.controls, + OriginAdmissionClass::Bulk => &mut self.bulk, + OriginAdmissionClass::StreamInstall => &mut self.stream_installs, + } + } +} + +#[derive(Default)] +struct OriginAdmission { + active: Mutex>, +} + +impl OriginAdmission { + fn try_acquire( + self: &Arc, + origin: ObservedOrigin, + class: OriginAdmissionClass, + ) -> Option { + let mut active = self + .active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let usage = active.entry(origin).or_default(); + if usage.count(class) >= class.limit() { + return None; + } + usage.increment(class); + drop(active); + + Some(OriginLease { + admission: Arc::clone(self), + origin, + class, + }) + } + + fn release(&self, origin: ObservedOrigin, class: OriginAdmissionClass) { + let mut active = self + .active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(usage) = active.get_mut(&origin) else { + log::error!( + "Origin admission lease for {:?} had no active entry", + origin.ip() + ); + return; + }; + if !usage.decrement(class) { + log::error!( + "Origin admission lease for {:?} had no active {class:?} count", + origin.ip() + ); + return; + } + if usage.is_empty() { + active.remove(&origin); + } + } + + #[cfg(test)] + fn active_origin_count(&self) -> usize { + self.active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len() + } +} + +struct OriginLease { + admission: Arc, + origin: ObservedOrigin, + class: OriginAdmissionClass, +} + +impl Drop for OriginLease { + fn drop(&mut self) { + self.admission.release(self.origin, self.class); + } +} + +pub(super) struct ControlAdmission { + _origin: OriginLease, + _global: OwnedSemaphorePermit, +} + +pub(super) struct BulkAdmission { + _origin: OriginLease, + _global: OwnedSemaphorePermit, +} + +pub(super) struct StreamInstallAdmission { + _origin: OriginLease, + global: Option, +} + +impl StreamInstallAdmission { + pub(super) fn take_global_permit(&mut self) -> OwnedSemaphorePermit { + self.global + .take() + .expect("Stream Install admission must own its global permit") + } +} + +pub(super) struct ServerAdmission { + origins: Arc, + control: Arc, + bulk: Arc, + stream_install: Arc, +} + +impl ServerAdmission { + fn new() -> Self { + Self { + origins: Arc::new(OriginAdmission::default()), + control: Arc::new(Semaphore::new(MAX_GLOBAL_CONTROL_STREAM_TASKS)), + bulk: Arc::new(Semaphore::new(MAX_GLOBAL_BULK_TRANSFER_TASKS)), + stream_install: Arc::new(Semaphore::new(MAX_GLOBAL_STREAM_INSTALL_TASKS)), + } + } + + fn try_acquire_connection(&self, origin: ObservedOrigin) -> Option { + self.origins + .try_acquire(origin, OriginAdmissionClass::Connection) + } + + pub(super) fn try_acquire_control(&self, origin: ObservedOrigin) -> Option { + let origin = self + .origins + .try_acquire(origin, OriginAdmissionClass::Control)?; + let global = Arc::clone(&self.control).try_acquire_owned().ok()?; + Some(ControlAdmission { + _origin: origin, + _global: global, + }) + } + + pub(super) fn try_acquire_bulk(&self, origin: ObservedOrigin) -> Option { + let origin = self + .origins + .try_acquire(origin, OriginAdmissionClass::Bulk)?; + let global = Arc::clone(&self.bulk).try_acquire_owned().ok()?; + Some(BulkAdmission { + _origin: origin, + _global: global, + }) + } + + pub(super) fn try_acquire_stream_install( + &self, + origin: ObservedOrigin, + ) -> Option { + let origin = self + .origins + .try_acquire(origin, OriginAdmissionClass::StreamInstall)?; + let global = Arc::clone(&self.stream_install).try_acquire_owned().ok()?; + Some(StreamInstallAdmission { + _origin: origin, + global: Some(global), + }) + } +} + struct BoundedEndpointLimits { inner: endpoint_limits::Default, } @@ -126,9 +383,7 @@ async fn run_server_body( Ok(()) }); let mut connection_tasks = JoinSet::new(); - let control_stream_permits = Arc::new(Semaphore::new(MAX_GLOBAL_CONTROL_STREAM_TASKS)); - let bulk_transfer_permits = Arc::new(Semaphore::new(MAX_GLOBAL_BULK_TRANSFER_TASKS)); - let stream_install_permits = Arc::new(Semaphore::new(MAX_GLOBAL_STREAM_INSTALL_TASKS)); + let admission = Arc::new(ServerAdmission::new()); let ready_addr = (*ctx.local_peer_addr.read().await).unwrap_or_else(|| direct_connect_addr(server_addr)); @@ -159,26 +414,44 @@ async fn run_server_body( break Err(eyre::eyre!("QUIC server accept loop ended unexpectedly")); }; + let remote_addr = match connection.remote_addr() { + Ok(remote_addr) => remote_addr, + Err(error) => { + log::warn!("Closing peer connection without an observed remote address: {error}"); + connection.close(application::Error::UNKNOWN); + continue; + } + }; + let Some(origin) = ObservedOrigin::from_remote_addr(Some(remote_addr)) else { + log::warn!("Closing peer connection without an observed origin"); + connection.close(application::Error::UNKNOWN); + continue; + }; + if !has_child_capacity(connection_tasks.len(), MAX_ESTABLISHED_CONNECTIONS) { log::warn!( - "Closing excess peer connection from {} at application limit {}", - connection.remote_addr().map_or_else( - |_| "unknown".to_owned(), - |addr| addr.to_string(), - ), + "Closing excess peer connection from {remote_addr} at application limit {}", MAX_ESTABLISHED_CONNECTIONS, ); connection.close(application::Error::UNKNOWN); continue; } + let Some(connection_admission) = admission.try_acquire_connection(origin) else { + log::warn!( + "Closing excess peer connection from {remote_addr} at per-origin limit {MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN}" + ); + connection.close(application::Error::UNKNOWN); + continue; + }; connection_tasks.spawn(handle_peer_connection( connection, + remote_addr, + origin, + connection_admission, ctx.clone(), server_children_shutdown.clone(), - Arc::clone(&control_stream_permits), - Arc::clone(&bulk_transfer_permits), - Arc::clone(&stream_install_permits), + Arc::clone(&admission), )); } } @@ -266,13 +539,13 @@ fn direct_connect_addr(server_addr: SocketAddr) -> SocketAddr { async fn handle_peer_connection( mut connection: Connection, + remote_addr: SocketAddr, + origin: ObservedOrigin, + _connection_admission: OriginLease, ctx: PeerCtx, server_shutdown: CancellationToken, - control_stream_permits: Arc, - bulk_transfer_permits: Arc, - stream_install_permits: Arc, + admission: Arc, ) -> eyre::Result<()> { - let remote_addr = connection.remote_addr()?; log::info!("{remote_addr} peer connected"); let connection_shutdown = server_shutdown.child_token(); @@ -303,9 +576,7 @@ async fn handle_peer_connection( let _ = stream.reset(application::Error::UNKNOWN); continue; } - let Ok(control_permit) = Arc::clone(&control_stream_permits) - .try_acquire_owned() - else { + let Some(control_admission) = admission.try_acquire_control(origin) else { let _ = stream.stop_sending(application::Error::UNKNOWN); let _ = stream.reset(application::Error::UNKNOWN); continue; @@ -315,11 +586,11 @@ async fn handle_peer_connection( stream_tasks.spawn(handle_admitted_peer_stream( stream, stream_ctx, - Some(remote_addr), + remote_addr, + origin, stream_shutdown, - control_permit, - Arc::clone(&bulk_transfer_permits), - Arc::clone(&stream_install_permits), + control_admission, + Arc::clone(&admission), )); } Ok(None) => break Ok(()), @@ -356,20 +627,20 @@ async fn handle_peer_connection( async fn handle_admitted_peer_stream( stream: BidirectionalStream, ctx: PeerCtx, - remote_addr: Option, + remote_addr: SocketAddr, + origin: ObservedOrigin, stream_shutdown: CancellationToken, - control_permit: OwnedSemaphorePermit, - bulk_transfer_permits: Arc, - stream_install_permits: Arc, + control_admission: ControlAdmission, + admission: Arc, ) -> eyre::Result<()> { handle_peer_stream( stream, ctx, remote_addr, + origin, stream_shutdown, - control_permit, - bulk_transfer_permits, - stream_install_permits, + control_admission, + admission, ) .await } @@ -400,6 +671,8 @@ async fn drain_joined_child_tasks(children: &mut JoinSet>, labe #[cfg(test)] mod tests { use std::{ + net::SocketAddr, + panic::{AssertUnwindSafe, catch_unwind}, sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -415,12 +688,20 @@ mod tests { use super::{ CONNECTION_NO_STREAM_IDLE_TIMEOUT, + MAX_BULK_TRANSFER_TASKS_PER_ORIGIN, MAX_CONTROL_STREAM_TASKS, + MAX_CONTROL_STREAM_TASKS_PER_ORIGIN, MAX_ESTABLISHED_CONNECTIONS, + MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN, MAX_GLOBAL_BULK_TRANSFER_TASKS, MAX_GLOBAL_CONTROL_STREAM_TASKS, MAX_GLOBAL_STREAM_INSTALL_TASKS, MAX_INFLIGHT_HANDSHAKES, + MAX_STREAM_INSTALL_TASKS_PER_ORIGIN, + ObservedOrigin, + OriginAdmission, + OriginAdmissionClass, + ServerAdmission, bounded_endpoint_limits, drain_joined_child_tasks, endpoint_connection_capacity_available, @@ -436,6 +717,10 @@ mod tests { assert_eq!(MAX_GLOBAL_CONTROL_STREAM_TASKS, 16); assert_eq!(MAX_GLOBAL_BULK_TRANSFER_TASKS, 48); assert_eq!(MAX_GLOBAL_STREAM_INSTALL_TASKS, 2); + assert_eq!(MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN, 8); + assert_eq!(MAX_CONTROL_STREAM_TASKS_PER_ORIGIN, 16); + assert_eq!(MAX_BULK_TRANSFER_TASKS_PER_ORIGIN, 8); + assert_eq!(MAX_STREAM_INSTALL_TASKS_PER_ORIGIN, 1); assert_eq!( u64::try_from(MAX_CONTROL_STREAM_TASKS).expect("stream task bound fits u64"), crate::quic_runtime::MAX_OPEN_BIDIRECTIONAL_STREAMS, @@ -491,6 +776,154 @@ mod tests { drop(held_bulk); } + fn origin(ip: [u8; 4], port: u16) -> ObservedOrigin { + ObservedOrigin::from_remote_addr(Some(SocketAddr::from((ip, port)))) + .expect("test address should identify an origin") + } + + #[test] + fn missing_origin_is_rejected_and_port_rotation_cannot_escape_connection_limit() { + assert!(ObservedOrigin::from_remote_addr(None).is_none()); + + let admission = Arc::new(OriginAdmission::default()); + let first_ip = [192, 168, 1, 20]; + let held = (0..MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN) + .map(|index| { + admission + .try_acquire( + origin( + first_ip, + 12_000 + u16::try_from(index).expect("test index should fit u16"), + ), + OriginAdmissionClass::Connection, + ) + .expect("connection below the per-origin limit should be admitted") + }) + .collect::>(); + assert!( + admission + .try_acquire(origin(first_ip, 13_000), OriginAdmissionClass::Connection,) + .is_none(), + "a new source port must not create a new origin budget" + ); + + let distinct = admission + .try_acquire( + origin([192, 168, 1, 21], 12_000), + OriginAdmissionClass::Connection, + ) + .expect("a distinct observed IP should keep its own capacity"); + assert_eq!(admission.active_origin_count(), 2); + + drop(held); + assert_eq!(admission.active_origin_count(), 1); + drop(distinct); + assert_eq!(admission.active_origin_count(), 0); + } + + #[test] + fn origin_leases_release_on_normal_drop_and_panic_unwind() { + let admission = Arc::new(OriginAdmission::default()); + let test_origin = origin([10, 0, 0, 7], 12_000); + { + let _lease = admission + .try_acquire(test_origin, OriginAdmissionClass::StreamInstall) + .expect("first provider should be admitted"); + assert_eq!(admission.active_origin_count(), 1); + } + assert_eq!(admission.active_origin_count(), 0); + + let unwind_admission = Arc::clone(&admission); + let unwind = catch_unwind(AssertUnwindSafe(move || { + let _lease = unwind_admission + .try_acquire(test_origin, OriginAdmissionClass::StreamInstall) + .expect("provider should be admitted before the injected panic"); + panic!("injected origin lease panic"); + })); + assert!(unwind.is_err()); + assert_eq!(admission.active_origin_count(), 0); + assert!( + admission + .try_acquire(test_origin, OriginAdmissionClass::StreamInstall) + .is_some(), + "panic unwinding must restore the origin capacity" + ); + } + + #[test] + fn origin_resource_classes_have_independent_active_limits() { + let admission = Arc::new(OriginAdmission::default()); + let test_origin = origin([10, 0, 0, 8], 12_000); + let connections = (0..MAX_ESTABLISHED_CONNECTIONS_PER_ORIGIN) + .map(|_| { + admission + .try_acquire(test_origin, OriginAdmissionClass::Connection) + .expect("connection class should retain independent capacity") + }) + .collect::>(); + let controls = (0..MAX_CONTROL_STREAM_TASKS_PER_ORIGIN) + .map(|_| { + admission + .try_acquire(test_origin, OriginAdmissionClass::Control) + .expect("control class should retain independent capacity") + }) + .collect::>(); + let bulk = (0..MAX_BULK_TRANSFER_TASKS_PER_ORIGIN) + .map(|_| { + admission + .try_acquire(test_origin, OriginAdmissionClass::Bulk) + .expect("bulk class should retain independent capacity") + }) + .collect::>(); + let stream_install = admission + .try_acquire(test_origin, OriginAdmissionClass::StreamInstall) + .expect("Stream Install class should retain independent capacity"); + + for class in [ + OriginAdmissionClass::Connection, + OriginAdmissionClass::Control, + OriginAdmissionClass::Bulk, + OriginAdmissionClass::StreamInstall, + ] { + assert!( + admission.try_acquire(test_origin, class).is_none(), + "{class:?} should close at its own active limit" + ); + } + + drop(connections); + drop(controls); + drop(bulk); + drop(stream_install); + assert_eq!(admission.active_origin_count(), 0); + } + + #[test] + fn saturated_stream_install_gate_preserves_ordinary_bulk_capacity() { + let admission = ServerAdmission::new(); + let first = admission + .try_acquire_stream_install(origin([10, 0, 0, 10], 12_000)) + .expect("first global provider slot should be available"); + let second = admission + .try_acquire_stream_install(origin([10, 0, 0, 11], 12_000)) + .expect("second global provider slot should be available"); + assert!( + admission + .try_acquire_stream_install(origin([10, 0, 0, 12], 12_000)) + .is_none(), + "the existing two-provider global gate must remain authoritative" + ); + assert!( + admission + .try_acquire_bulk(origin([10, 0, 0, 12], 12_001)) + .is_some(), + "provider saturation must not consume ordinary bulk capacity" + ); + + drop(first); + drop(second); + } + #[tokio::test] async fn draining_waits_for_every_childs_natural_cleanup() { let owner_closed = CancellationToken::new(); diff --git a/crates/lanspread-peer/src/services/stream.rs b/crates/lanspread-peer/src/services/stream.rs index c495e74..56d1dc1 100644 --- a/crates/lanspread-peer/src/services/stream.rs +++ b/crates/lanspread-peer/src/services/stream.rs @@ -15,7 +15,6 @@ use s2n_quic::{ application, stream::{BidirectionalStream, SendStream}, }; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio_util::{ codec::{FramedRead, FramedWrite, LengthDelimitedCodec}, sync::CancellationToken, @@ -25,6 +24,7 @@ use crate::{ context::PeerCtx, services::{ remote_state, + server::{ControlAdmission, ObservedOrigin, ServerAdmission}, state_sync::StateDomain, transfer::{ChunkDispatch, handle_file_chunk_request, handle_stream_install_request}, }, @@ -57,21 +57,21 @@ fn request_codec() -> LengthDelimitedCodec { pub(super) async fn handle_peer_stream( stream: BidirectionalStream, ctx: PeerCtx, - remote_addr: Option, + remote_addr: SocketAddr, + origin: ObservedOrigin, stream_shutdown: CancellationToken, - control_permit: OwnedSemaphorePermit, - bulk_transfer_permits: Arc, - stream_install_permits: Arc, + control_admission: ControlAdmission, + admission: Arc, ) -> eyre::Result<()> { let (rx, tx) = stream.split(); let mut framed_rx = FramedRead::new(rx, request_codec()); let mut framed_tx = FramedWrite::new(tx, control_codec()); - log::trace!("{remote_addr:?} peer stream opened"); + log::trace!("{remote_addr} peer stream opened"); - let source_ip = remote_addr.map(|addr| addr.ip()); + let source_ip = Some(origin.ip()); let first_frame = read_expected_frame(&mut framed_rx, &stream_shutdown).await; - let mut control_permit = Some(control_permit); - let mut _bulk_permit = None; + let mut control_admission = Some(control_admission); + let mut _bulk_admission = None; let mut response_reset = false; match first_frame { FrameRead::Frame(data) => { @@ -79,23 +79,23 @@ pub(super) async fn handle_peer_stream( if trailing == TrailingRead::Eof { match Request::decode(data.freeze()) { Ok(request) => { - log::debug!("{remote_addr:?} msg: {request:?}"); + log::debug!("{remote_addr} msg: {request:?}"); if request_is_bulk(&request) { - let bulk_permit = - Arc::clone(&bulk_transfer_permits).try_acquire_owned(); + let bulk_admission = admission.try_acquire_bulk(origin); // Once the single bounded request is decoded, bulk // work moves to its smaller pool so it cannot hold // every control-plane permit during long egress. - drop(control_permit.take()); - if let Ok(permit) = bulk_permit { - _bulk_permit = Some(permit); + drop(control_admission.take()); + if let Some(bulk_admission) = bulk_admission { + _bulk_admission = Some(bulk_admission); let dispatched = dispatch_request( &ctx, request, source_ip, framed_tx, &stream_shutdown, - &stream_install_permits, + &admission, + origin, ) .await; framed_tx = dispatched.writer; @@ -113,7 +113,8 @@ pub(super) async fn handle_peer_stream( source_ip, framed_tx, &stream_shutdown, - &stream_install_permits, + &admission, + origin, ) .await; framed_tx = dispatched.writer; @@ -121,9 +122,7 @@ pub(super) async fn handle_peer_stream( } } Err(error) => { - log::warn!( - "Rejecting invalid control request from {remote_addr:?}: {error}" - ); + log::warn!("Rejecting invalid control request from {remote_addr}: {error}"); framed_tx = send_response( framed_tx, Response::Error(ControlErrorCode::InvalidRequest), @@ -134,7 +133,7 @@ pub(super) async fn handle_peer_stream( } } } else if trailing != TrailingRead::Cancelled { - log::warn!("Rejecting non-singular control request from {remote_addr:?}"); + log::warn!("Rejecting non-singular control request from {remote_addr}"); framed_tx = send_response( framed_tx, Response::Error(ControlErrorCode::InvalidRequest), @@ -145,7 +144,7 @@ pub(super) async fn handle_peer_stream( } } FrameRead::Invalid(error) => { - log::warn!("Rejecting malformed control frame from {remote_addr:?}: {error}"); + log::warn!("Rejecting malformed control frame from {remote_addr}: {error}"); framed_tx = send_response( framed_tx, Response::Error(ControlErrorCode::InvalidRequest), @@ -154,7 +153,7 @@ pub(super) async fn handle_peer_stream( ) .await; } - FrameRead::Eof => log::trace!("{remote_addr:?} peer stream closed without a request"), + FrameRead::Eof => log::trace!("{remote_addr} peer stream closed without a request"), FrameRead::Cancelled => {} } @@ -232,7 +231,8 @@ async fn dispatch_request( source_ip: Option, framed_tx: ResponseWriter, stream_shutdown: &CancellationToken, - stream_install_permits: &Arc, + admission: &ServerAdmission, + origin: ObservedOrigin, ) -> DispatchResult { match request { Request::Ping => { @@ -327,23 +327,24 @@ async fn dispatch_request( game_id, content_id, } => { - let Ok(stream_install_permit) = Arc::clone(stream_install_permits).try_acquire_owned() + let Some(mut stream_install_admission) = admission.try_acquire_stream_install(origin) else { let mut tx = framed_tx.into_inner(); let _ = tx.reset(application::Error::UNKNOWN); return DispatchResult::reset(FramedWrite::new(tx, control_codec())); }; - DispatchResult::close( - handle_stream_install_request( - ctx, - game_id, - content_id, - framed_tx, - stream_shutdown, - stream_install_permit, - ) - .await, + let stream_install_permit = stream_install_admission.take_global_permit(); + let writer = handle_stream_install_request( + ctx, + game_id, + content_id, + framed_tx, + stream_shutdown, + stream_install_permit, ) + .await; + drop(stream_install_admission); + DispatchResult::close(writer) } } } @@ -409,7 +410,7 @@ async fn send_response( async fn close_or_reset_stream( framed_rx: FramedRead, mut framed_tx: ResponseWriter, - remote_addr: Option, + remote_addr: SocketAddr, cancellation: &CancellationToken, response_reset: bool, ) { @@ -437,7 +438,7 @@ async fn close_or_reset_stream( return; } if let Some(Err(error)) = close_result { - log::debug!("{remote_addr:?} failed to close peer response stream: {error}"); + log::debug!("{remote_addr} failed to close peer response stream: {error}"); } }