From 8b2b88f30ca0971e889c2589adb247a24e3dca6b Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:44:17 +0300 Subject: [PATCH] Session Residence + Ownership Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/web/http/websocket.rs | 2 + src/web/http/websocket/driver.rs | 87 +++++++++++++--- src/web/http/websocket/driver/io.rs | 22 +++-- src/web/manager.rs | 19 ++++ src/web/manager/budget.rs | 2 - src/web/manager/websocket.rs | 100 ++++++++++++++----- src/web/session.rs | 12 +++ src/web/session/downlink.rs | 67 ++++++++++--- src/web/session/lanes.rs | 49 +++++++--- src/web/session/resident.rs | 147 ++++++++++++++++++++++++++++ 10 files changed, 436 insertions(+), 71 deletions(-) create mode 100644 src/web/session/resident.rs diff --git a/src/web/http/websocket.rs b/src/web/http/websocket.rs index 5776f35..df90792 100644 --- a/src/web/http/websocket.rs +++ b/src/web/http/websocket.rs @@ -284,11 +284,13 @@ pub(super) async fn handle( .admit_websocket( session.profile_key(), session.trace_session_id(), + session.token_hash(), effective_ip, kind, BASE_BUDGET_BYTES, Duration::from_secs(timeouts.long_poll_secs), Duration::from_secs(timeouts.websocket_eviction_secs), + session.carrier_cancellation(), ) .await { diff --git a/src/web/http/websocket/driver.rs b/src/web/http/websocket/driver.rs index 9fa931c..989ceab 100644 --- a/src/web/http/websocket/driver.rs +++ b/src/web/http/websocket/driver.rs @@ -30,7 +30,16 @@ pub(super) async fn run_upgraded( trace: Option, acknowledge_commit: bool, ) { - let Ok(upgraded) = on_upgrade.await else { + let cancellation = connection.cancellation(); + let timeouts = runtime.active_generation().config().web.timeouts.clone(); + let upgraded = tokio::select! { + _ = cancellation.cancelled() => return, + result = tokio::time::timeout( + Duration::from_secs(timeouts.websocket_upgrade_secs), + on_upgrade, + ) => result, + }; + let Ok(Ok(upgraded)) = upgraded else { return; }; let Ok(parts) = upgraded.downcast::>() else { @@ -51,7 +60,6 @@ pub(super) async fn run_upgraded( .max_frame_size(Some(limits.carrier_batch_bytes)); let mut socket = WebSocketStream::from_raw_socket(io, Role::Server, Some(config)).await; connection.mark_opened(); - let cancellation = connection.cancellation(); if let Some(reservation) = lane_reservation.as_mut() { let _ = run_lane( &mut socket, @@ -84,7 +92,9 @@ pub(super) async fn run_upgraded( .timeouts .websocket_eviction_secs, ); - let _ = tokio::time::timeout(eviction, socket.close(None)).await; + if !cancellation.is_cancelled() { + let _ = tokio::time::timeout(eviction, socket.close(None)).await; + } if let Some(reservation) = lane_reservation { session.close_websocket_lane(reservation.lane_id()); drop(reservation); @@ -111,11 +121,22 @@ async fn run_multiplex( let mut read_budget = None; let liveness_interval = connection.liveness_interval(); let mut next_ping = Instant::now() + liveness_interval; + let open_deadline = Instant::now() + + Duration::from_secs( + runtime + .active_generation() + .config() + .web + .timeouts + .websocket_open_secs, + ); + let mut active = false; loop { let down = session.poll_down(cursor); tokio::pin!(down); let event = tokio::select! { _ = cancellation.cancelled() => return Err(()), + _ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()), _ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness, incoming = read_message( socket, @@ -145,7 +166,13 @@ async fn run_multiplex( result?; if acknowledge_commit && sequence == 1 && session.is_carrier_committed() { let started = Instant::now(); - send(socket, runtime, Message::Binary(Bytes::new())).await?; + send( + socket, + runtime, + Message::Binary(Bytes::new()), + &cancellation, + ) + .await?; record_message( runtime, trace, @@ -155,6 +182,10 @@ async fn run_multiplex( started, ); } + if !active { + connection.mark_active(); + active = true; + } sequence = sequence.checked_add(1).ok_or(())?; connection.mark_peer_activity(); next_ping = Instant::now() + liveness_interval; @@ -173,7 +204,7 @@ async fn run_multiplex( } Message::Ping(payload) => { let started = Instant::now(); - flush(socket, runtime).await?; + flush(socket, runtime, &cancellation).await?; record_message( runtime, trace, @@ -220,7 +251,7 @@ async fn run_multiplex( DriverEvent::Down(result) => { if result.body.is_empty() { let started = Instant::now(); - send(socket, runtime, Message::Ping(Bytes::new())).await?; + send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?; record_message( runtime, trace, @@ -241,7 +272,7 @@ async fn run_multiplex( let body = result.body; let started = Instant::now(); if trace.is_some() { - send(socket, runtime, Message::Binary(body.clone())).await?; + send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?; record_message( runtime, trace, @@ -251,7 +282,7 @@ async fn run_multiplex( started, ); } else { - send(socket, runtime, Message::Binary(body)).await?; + send(socket, runtime, Message::Binary(body), &cancellation).await?; } connection.mark_progress(); } @@ -259,7 +290,7 @@ async fn run_multiplex( } DriverEvent::Liveness => { let started = Instant::now(); - send(socket, runtime, Message::Ping(Bytes::new())).await?; + send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?; record_message( runtime, trace, @@ -291,11 +322,22 @@ async fn run_lane( let mut read_budget = None; let liveness_interval = connection.liveness_interval(); let mut next_ping = Instant::now() + liveness_interval; + let open_deadline = Instant::now() + + Duration::from_secs( + runtime + .active_generation() + .config() + .web + .timeouts + .websocket_open_secs, + ); + let mut active = false; loop { let down = session.poll_down_lane(reservation.lane_id(), cursor); tokio::pin!(down); let event = tokio::select! { _ = cancellation.cancelled() => return Err(()), + _ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()), _ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness, incoming = read_message( socket, @@ -332,7 +374,18 @@ async fn run_lane( result?; if acknowledge_commit && sequence == 1 && session.is_carrier_committed() { let started = Instant::now(); - send(socket, runtime, Message::Binary(Bytes::new())).await?; + if send( + socket, + runtime, + Message::Binary(Bytes::new()), + &cancellation, + ) + .await + .is_err() + { + session.close(); + return Err(()); + } record_message( runtime, trace, @@ -342,6 +395,10 @@ async fn run_lane( started, ); } + if !active { + connection.mark_active(); + active = true; + } sequence = sequence.checked_add(1).ok_or(())?; connection.mark_peer_activity(); next_ping = Instant::now() + liveness_interval; @@ -360,7 +417,7 @@ async fn run_lane( } Message::Ping(payload) => { let started = Instant::now(); - flush(socket, runtime).await?; + flush(socket, runtime, &cancellation).await?; record_message( runtime, trace, @@ -410,7 +467,7 @@ async fn run_lane( } if result.body.is_empty() { let started = Instant::now(); - send(socket, runtime, Message::Ping(Bytes::new())).await?; + send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?; record_message( runtime, trace, @@ -431,7 +488,7 @@ async fn run_lane( let body = result.body; let started = Instant::now(); if trace.is_some() { - send(socket, runtime, Message::Binary(body.clone())).await?; + send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?; record_message( runtime, trace, @@ -441,7 +498,7 @@ async fn run_lane( started, ); } else { - send(socket, runtime, Message::Binary(body)).await?; + send(socket, runtime, Message::Binary(body), &cancellation).await?; } connection.mark_progress(); } @@ -449,7 +506,7 @@ async fn run_lane( } DriverEvent::Liveness => { let started = Instant::now(); - send(socket, runtime, Message::Ping(Bytes::new())).await?; + send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?; record_message( runtime, trace, diff --git a/src/web/http/websocket/driver/io.rs b/src/web/http/websocket/driver/io.rs index 4ec87ce..ddc4941 100644 --- a/src/web/http/websocket/driver/io.rs +++ b/src/web/http/websocket/driver/io.rs @@ -160,6 +160,7 @@ pub(super) async fn send( socket: &mut CarrierSocket, runtime: &WebProcessRuntime, message: Message, + cancellation: &CancellationToken, ) -> Result<(), ()> { let timeout = Duration::from_secs( runtime @@ -169,15 +170,18 @@ pub(super) async fn send( .timeouts .websocket_write_secs, ); - tokio::time::timeout(timeout, socket.send(message)) - .await - .map_err(|_| ())? - .map_err(|_| ()) + tokio::select! { + _ = cancellation.cancelled() => Err(()), + result = tokio::time::timeout(timeout, socket.send(message)) => { + result.map_err(|_| ())?.map_err(|_| ()) + } + } } pub(super) async fn flush( socket: &mut CarrierSocket, runtime: &WebProcessRuntime, + cancellation: &CancellationToken, ) -> Result<(), ()> { let timeout = Duration::from_secs( runtime @@ -187,10 +191,12 @@ pub(super) async fn flush( .timeouts .websocket_write_secs, ); - tokio::time::timeout(timeout, socket.flush()) - .await - .map_err(|_| ())? - .map_err(|_| ()) + tokio::select! { + _ = cancellation.cancelled() => Err(()), + result = tokio::time::timeout(timeout, socket.flush()) => { + result.map_err(|_| ())?.map_err(|_| ()) + } + } } pub(super) fn record_message( diff --git a/src/web/manager.rs b/src/web/manager.rs index 1efce7a..e8ea123 100644 --- a/src/web/manager.rs +++ b/src/web/manager.rs @@ -290,6 +290,21 @@ impl WebProcessRuntime { Some((reader, body)) } + /// Reserves transient bytes while one downlink batch replaces queued frames. + pub(crate) fn try_downlink_staging_budget( + &self, + bytes: usize, + ) -> Option { + let bytes = u32::try_from(bytes).ok()?; + let permit = Arc::clone(&self.body_bytes) + .try_acquire_many_owned(bytes) + .ok(); + if permit.is_none() { + self.record_limit_hit(); + } + permit + } + /// Reserves bounded process-wide queue capacity for data or control traffic. pub(crate) fn try_reserve_pending( &self, @@ -351,21 +366,25 @@ impl WebProcessRuntime { self: &Arc, owner: ProfileKey, session_id: u64, + session_hash: TokenHash, client_ip: IpAddr, kind: WebSocketKind, base_bytes: usize, liveness_interval: Duration, eviction_timeout: Duration, + parent_cancellation: CancellationToken, ) -> Result { websocket::admit( self, owner, session_id, + session_hash, client_ip, kind, base_bytes, liveness_interval, eviction_timeout, + parent_cancellation, ) .await } diff --git a/src/web/manager/budget.rs b/src/web/manager/budget.rs index b1fbeba..af1712f 100644 --- a/src/web/manager/budget.rs +++ b/src/web/manager/budget.rs @@ -163,7 +163,6 @@ impl WebDataBudget { } remove_owner(&mut state.owner_bytes, owner, bytes); drop(state); - self.pressured.store(false, Ordering::Release); self.notify.notify_waiters(); } @@ -259,7 +258,6 @@ impl WebDataBudget { state.websocket_bytes = state.websocket_bytes.saturating_sub(bytes); remove_owner(&mut state.owner_bytes, owner, bytes); drop(state); - self.pressured.store(false, Ordering::Release); self.notify.notify_waiters(); } } diff --git a/src/web/manager/websocket.rs b/src/web/manager/websocket.rs index c0af4c8..e8350d8 100644 --- a/src/web/manager/websocket.rs +++ b/src/web/manager/websocket.rs @@ -275,20 +275,10 @@ impl WebProcessRuntime { pub(super) fn cleanup_websockets(&self) { let now = self.websocket_tick(); - let mut victims = self - .websockets - .lock() - .entries - .values() - .filter(|entry| { - now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire)) - >= dead_after(entry) - }) - .cloned() - .collect::>(); + let mut victims = claim_stale_victims(self, now); if victims.is_empty() && self.data_budget.take_pressure() - && let Some(victim) = select_pressure_victim(self, now) + && let Some(victim) = select_pressure_victim(self, now, true) { victims.push(victim); } @@ -315,16 +305,22 @@ fn select_victim( session_id: u64, client_ip: IpAddr, excluded_id: Option, + claim: bool, ) -> Option> { let fair_share = runtime.data_budget.fair_share(Some(owner)); let requester_usage = runtime.data_budget.owner_usage(owner); let now = runtime.websocket_tick(); - runtime - .websockets - .lock() + let mut registry = runtime.websockets.lock(); + if claim + && registry.evictions_in_flight >= runtime.limits.max_websocket_evictions_in_flight + { + return None; + } + let selected = registry .entries .values() .filter(|entry| Some(entry.id) != excluded_id) + .filter(|entry| !entry.closing.load(Ordering::Acquire)) .filter_map(|entry| { let owner_rank = if entry.session_id == session_id { 0 @@ -353,15 +349,28 @@ fn select_victim( )) }) .min_by_key(|(key, _)| *key) - .map(|(_, entry)| entry) + .map(|(_, entry)| entry)?; + if claim && !claim_entry(&mut registry, &selected, runtime) { + return None; + } + Some(selected) } -fn select_pressure_victim(runtime: &WebProcessRuntime, now: u64) -> Option> { - runtime - .websockets - .lock() +fn select_pressure_victim( + runtime: &WebProcessRuntime, + now: u64, + claim: bool, +) -> Option> { + let mut registry = runtime.websockets.lock(); + if claim + && registry.evictions_in_flight >= runtime.limits.max_websocket_evictions_in_flight + { + return None; + } + let selected = registry .entries .values() + .filter(|entry| !entry.closing.load(Ordering::Acquire)) .map(|entry| { ( ( @@ -374,11 +383,58 @@ fn select_pressure_victim(runtime: &WebProcessRuntime, now: u64) -> Option Vec> { + let mut registry = runtime.websockets.lock(); + let available = runtime + .limits + .max_websocket_evictions_in_flight + .saturating_sub(registry.evictions_in_flight); + let candidates = registry + .entries + .values() + .filter(|entry| !entry.closing.load(Ordering::Acquire)) + .filter(|entry| { + now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire)) + >= dead_after(entry) + }) + .take(available) + .cloned() + .collect::>(); + candidates + .into_iter() + .filter(|entry| claim_entry(&mut registry, entry, runtime)) + .collect() +} + +fn claim_entry( + registry: &mut WebSocketRegistry, + entry: &Arc, + runtime: &WebProcessRuntime, +) -> bool { + if registry.evictions_in_flight >= runtime.limits.max_websocket_evictions_in_flight + || entry + .closing + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return false; + } + entry + .phase + .store(WebSocketPhase::Closing as u8, Ordering::Release); + registry.evictions_in_flight += 1; + true } fn entry_priority(entry: &WebSocketEntry, now: u64) -> u8 { - if !entry.opened.load(Ordering::Acquire) + if entry.phase.load(Ordering::Acquire) < WebSocketPhase::Active as u8 || now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire)) >= dead_after(entry) { 0 diff --git a/src/web/session.rs b/src/web/session.rs index 2d385a4..0a8d83b 100644 --- a/src/web/session.rs +++ b/src/web/session.rs @@ -21,6 +21,8 @@ use crate::web::manager::CarrierLearningContext; mod backend; // Downlink queues own cursor replay, flow control, and memory reservations. mod downlink; +// Response ownership keeps detached batches charged until the last body clone drops. +mod resident; // Lane carrier state isolates request sequencing and downlink replay per logical stream. mod lanes; // WebSocket carrier state owns pre-OPEN lane reservations and failure isolation. @@ -71,6 +73,7 @@ struct QueuedFrame { struct DownBatch { body: Bytes, + lease: Arc, base_cursor: u64, next_cursor: u64, data_bytes: usize, @@ -83,6 +86,7 @@ struct CarrierLane { instance: u64, pending_bytes: usize, pending_items: usize, + resident: Arc, pending_frames: VecDeque, pending_windows: HashMap, unacked: Option, @@ -100,6 +104,7 @@ impl CarrierLane { instance, pending_bytes: 0, pending_items: 0, + resident: Arc::new(resident::ResidentCounters::default()), pending_frames: VecDeque::new(), pending_windows: HashMap::new(), unacked: None, @@ -169,6 +174,7 @@ pub(crate) struct WebSession { cancel: CancellationToken, tasks_live: AtomicUsize, tasks_done: Arc, + resident: Arc, finished: AtomicBool, up_active: AtomicBool, } @@ -251,6 +257,7 @@ impl WebSession { cancel: CancellationToken::new(), tasks_live: AtomicUsize::new(0), tasks_done: Arc::new(Notify::new()), + resident: Arc::new(resident::ResidentCounters::default()), finished: AtomicBool::new(false), up_active: AtomicBool::new(false), }) @@ -281,6 +288,11 @@ impl WebSession { self.trace_session_id } + /// Creates a child cancellation boundary for one owned carrier task. + pub(crate) fn carrier_cancellation(&self) -> CancellationToken { + self.cancel.child_token() + } + /// Returns whether accepted carrier progress made this attempt immutable. pub(crate) fn is_carrier_committed(&self) -> bool { self.state.lock().negotiation_phase == SessionNegotiationPhase::Committed diff --git a/src/web/session/downlink.rs b/src/web/session/downlink.rs index 33d5f22..76a1207 100644 --- a/src/web/session/downlink.rs +++ b/src/web/session/downlink.rs @@ -1,7 +1,9 @@ +use std::sync::Arc; use std::time::{Duration, Instant}; use bytes::{BufMut, Bytes, BytesMut}; +use super::resident::{OwnedBatchBody, PendingCounts, PendingResponseLease}; use super::{ DownBatch, PendingClass, PollResult, QUEUE_ITEM_COST, QueuedFrame, SessionState, WebSession, }; @@ -60,6 +62,9 @@ impl WebSession { if !state.pending_frames.is_empty() { let batch = match self.take_down_batch_locked(&mut state, cursor) { Ok(batch) => batch, + Err(ManagerError::Backpressure) => { + return Err(ManagerError::Backpressure); + } Err(error) => { drop(state); self.close(); @@ -121,6 +126,15 @@ impl WebSession { .limits .pending_items_per_session .saturating_sub(item_reserve); + let resident = self.resident.snapshot(); + let pending_bytes = state.pending_bytes.saturating_add(resident.bytes()); + let pending_items = state.pending_items.saturating_add(resident.items()); + let pending_control_bytes = state + .pending_control_bytes + .saturating_add(resident.control_bytes); + let pending_control_items = state + .pending_control_items + .saturating_add(resident.control_items); if state.closed { return false; } @@ -128,20 +142,16 @@ impl WebSession { let fits = if control { bytes <= self.limits.control_bytes_per_session && items <= item_reserve - && state.pending_bytes + && pending_bytes <= self.limits.pending_bytes_per_session.saturating_sub(bytes) - && state.pending_items + && pending_items <= self.limits.pending_items_per_session.saturating_sub(items) - && state.pending_control_bytes + && pending_control_bytes <= self.limits.control_bytes_per_session.saturating_sub(bytes) - && state.pending_control_items <= item_reserve.saturating_sub(items) + && pending_control_items <= item_reserve.saturating_sub(items) } else { - let data_bytes = state - .pending_bytes - .saturating_sub(state.pending_control_bytes); - let data_items = state - .pending_items - .saturating_sub(state.pending_control_items); + let data_bytes = pending_bytes.saturating_sub(pending_control_bytes); + let data_items = pending_items.saturating_sub(pending_control_items); let (byte_limit, item_limit) = if class == PendingClass::Downlink { let uplink_bytes = self.limits.max_body_bytes.saturating_add( self.limits @@ -203,6 +213,21 @@ impl WebSession { } } + pub(super) fn release_local_locked( + &self, + state: &mut SessionState, + bytes: usize, + items: usize, + control: bool, + ) { + state.pending_bytes = state.pending_bytes.saturating_sub(bytes); + state.pending_items = state.pending_items.saturating_sub(items); + if control { + state.pending_control_bytes = state.pending_control_bytes.saturating_sub(bytes); + state.pending_control_items = state.pending_control_items.saturating_sub(items); + } + } + /// Coalesces one flow-control update into the bounded control queue. pub(super) fn queue_window_locked( &self, @@ -351,6 +376,12 @@ impl WebSession { body_len += queued.encoded.len(); count += 1; } + let Some(manager) = self.manager.upgrade() else { + return Err(ManagerError::Closed); + }; + let Some(_staging) = manager.try_downlink_staging_budget(body_len) else { + return Err(ManagerError::Backpressure); + }; let mut body = BytesMut::with_capacity(body_len); let mut data_bytes = 0usize; let mut data_items = 0usize; @@ -383,8 +414,17 @@ impl WebSession { *index = index.saturating_sub(count); } state.down_cursor = next_cursor; + let counts = PendingCounts { + data_bytes, + data_items, + control_bytes, + control_items, + }; + let lease = PendingResponseLease::new(self, counts, None); + let body = Bytes::from_owner(OwnedBatchBody::new(body.freeze(), Arc::clone(&lease))); Ok(DownBatch { - body: body.freeze(), + body, + lease, base_cursor: cursor, next_cursor, data_bytes, @@ -398,8 +438,9 @@ impl WebSession { let Some(batch) = state.unacked.take() else { return; }; - self.release_locked(state, batch.data_bytes, batch.data_items, false); - self.release_locked(state, batch.control_bytes, batch.control_items, true); + batch.lease.detach(); + self.release_local_locked(state, batch.data_bytes, batch.data_items, false); + self.release_local_locked(state, batch.control_bytes, batch.control_items, true); for stream in state.streams.values_mut() { if let Some(waker) = stream.write_waker.take() { waker.wake(); diff --git a/src/web/session/lanes.rs b/src/web/session/lanes.rs index 2620ab1..ccf988a 100644 --- a/src/web/session/lanes.rs +++ b/src/web/session/lanes.rs @@ -6,6 +6,7 @@ use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; use super::uplink::{inbound_reservation, validate_batch}; +use super::resident::{OwnedBatchBody, PendingCounts, PendingResponseLease}; use super::{ CarrierLane, DownBatch, PendingClass, PollResult, QUEUE_ITEM_COST, QueuedFrame, SessionState, WebSession, insert_carrier_lane, remember_closed, @@ -241,8 +242,14 @@ impl WebSession { .pending_items .saturating_sub(batch.data_items.saturating_add(batch.control_items)); } - self.release_locked(&mut state, batch.data_bytes, batch.data_items, false); - self.release_locked(&mut state, batch.control_bytes, batch.control_items, true); + batch.lease.detach(); + self.release_local_locked(&mut state, batch.data_bytes, batch.data_items, false); + self.release_local_locked( + &mut state, + batch.control_bytes, + batch.control_items, + true, + ); if let Some(stream) = state.streams.get_mut(&lane_id) && let Some(waker) = stream.write_waker.take() { @@ -282,8 +289,11 @@ impl WebSession { }); } if !lane.pending_frames.is_empty() { - let batch = match take_lane_down_batch(&self.limits, lane, cursor) { + let batch = match take_lane_down_batch(self, &self.limits, lane, cursor) { Ok(batch) => batch, + Err(ManagerError::Backpressure) => { + return Err(ManagerError::Backpressure); + } Err(error) => { drop(state); self.close(); @@ -457,7 +467,8 @@ impl WebSession { }); if can_coalesce { if state.carrier_lanes.get(&stream_id).is_none_or(|lane| { - lane.pending_bytes + let resident = lane.resident.snapshot(); + lane.pending_bytes.saturating_add(resident.bytes()) > self .limits .pending_bytes_per_lane @@ -491,9 +502,10 @@ impl WebSession { PendingClass::Downlink }; if state.carrier_lanes.get(&stream_id).is_none_or(|lane| { - lane.pending_bytes + let resident = lane.resident.snapshot(); + lane.pending_bytes.saturating_add(resident.bytes()) > self.limits.pending_bytes_per_lane.saturating_sub(cost) - || lane.pending_items + || lane.pending_items.saturating_add(resident.items()) >= self.limits.pending_items_per_lane }) { return false; @@ -561,10 +573,9 @@ impl WebSession { } } if let Some(batch) = lane.unacked.take() { - data_bytes = data_bytes.saturating_add(batch.data_bytes); - data_items = data_items.saturating_add(batch.data_items); - control_bytes = control_bytes.saturating_add(batch.control_bytes); - control_items = control_items.saturating_add(batch.control_items); + batch.lease.detach(); + self.release_local_locked(state, batch.data_bytes, batch.data_items, false); + self.release_local_locked(state, batch.control_bytes, batch.control_items, true); } self.release_locked(state, data_bytes, data_items, false); self.release_locked(state, control_bytes, control_items, true); @@ -593,6 +604,7 @@ fn only_late_frames(frames: &[Frame<'_>]) -> bool { } fn take_lane_down_batch( + session: &WebSession, limits: &WebLimitsConfig, lane: &mut CarrierLane, cursor: u64, @@ -613,6 +625,12 @@ fn take_lane_down_batch( body_len += queued.encoded.len(); count += 1; } + let Some(manager) = session.manager.upgrade() else { + return Err(ManagerError::Closed); + }; + let Some(_staging) = manager.try_downlink_staging_budget(body_len) else { + return Err(ManagerError::Backpressure); + }; let mut body = BytesMut::with_capacity(body_len); let mut data_bytes = 0usize; let mut data_items = 0usize; @@ -645,8 +663,17 @@ fn take_lane_down_batch( *index = index.saturating_sub(count); } lane.down_cursor = next_cursor; + let counts = PendingCounts { + data_bytes, + data_items, + control_bytes, + control_items, + }; + let lease = PendingResponseLease::new(session, counts, Some(Arc::clone(&lane.resident))); + let body = Bytes::from_owner(OwnedBatchBody::new(body.freeze(), Arc::clone(&lease))); Ok(DownBatch { - body: body.freeze(), + body, + lease, base_cursor: cursor, next_cursor, data_bytes, diff --git a/src/web/session/resident.rs b/src/web/session/resident.rs new file mode 100644 index 0000000..e6e59f9 --- /dev/null +++ b/src/web/session/resident.rs @@ -0,0 +1,147 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use bytes::Bytes; + +use super::WebSession; +use crate::web::manager::{ProfileKey, WebProcessRuntime}; + +#[derive(Clone, Copy, Default)] +pub(super) struct PendingCounts { + pub(super) data_bytes: usize, + pub(super) data_items: usize, + pub(super) control_bytes: usize, + pub(super) control_items: usize, +} + +impl PendingCounts { + pub(super) fn bytes(self) -> usize { + self.data_bytes.saturating_add(self.control_bytes) + } + + pub(super) fn items(self) -> usize { + self.data_items.saturating_add(self.control_items) + } +} + +#[derive(Default)] +pub(super) struct ResidentCounters { + data_bytes: AtomicUsize, + data_items: AtomicUsize, + control_bytes: AtomicUsize, + control_items: AtomicUsize, +} + +impl ResidentCounters { + pub(super) fn snapshot(&self) -> PendingCounts { + PendingCounts { + data_bytes: self.data_bytes.load(Ordering::Acquire), + data_items: self.data_items.load(Ordering::Acquire), + control_bytes: self.control_bytes.load(Ordering::Acquire), + control_items: self.control_items.load(Ordering::Acquire), + } + } + + fn add(&self, counts: PendingCounts) { + self.data_bytes.fetch_add(counts.data_bytes, Ordering::AcqRel); + self.data_items.fetch_add(counts.data_items, Ordering::AcqRel); + self.control_bytes + .fetch_add(counts.control_bytes, Ordering::AcqRel); + self.control_items + .fetch_add(counts.control_items, Ordering::AcqRel); + } + + fn remove(&self, counts: PendingCounts) { + self.data_bytes.fetch_sub(counts.data_bytes, Ordering::AcqRel); + self.data_items.fetch_sub(counts.data_items, Ordering::AcqRel); + self.control_bytes + .fetch_sub(counts.control_bytes, Ordering::AcqRel); + self.control_items + .fetch_sub(counts.control_items, Ordering::AcqRel); + } +} + +pub(super) struct PendingResponseLease { + manager: std::sync::Weak, + owner: ProfileKey, + counts: PendingCounts, + session: Arc, + lane: Option>, + detached: AtomicBool, +} + +impl PendingResponseLease { + pub(super) fn new( + session: &WebSession, + counts: PendingCounts, + lane: Option>, + ) -> Arc { + Arc::new(Self { + manager: session.manager.clone(), + owner: session.profile_key, + counts, + session: Arc::clone(&session.resident), + lane, + detached: AtomicBool::new(false), + }) + } + + pub(super) fn detach(&self) { + if self + .detached + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + self.session.add(self.counts); + if let Some(lane) = &self.lane { + lane.add(self.counts); + } + } +} + +impl Drop for PendingResponseLease { + fn drop(&mut self) { + if self.detached.load(Ordering::Acquire) { + self.session.remove(self.counts); + if let Some(lane) = &self.lane { + lane.remove(self.counts); + } + } + if let Some(manager) = self.manager.upgrade() { + manager.release_pending( + self.owner, + self.counts.data_bytes, + self.counts.data_items, + false, + ); + manager.release_pending( + self.owner, + self.counts.control_bytes, + self.counts.control_items, + true, + ); + } + } +} + +pub(super) struct OwnedBatchBody { + bytes: Bytes, + _lease: Arc, +} + +impl OwnedBatchBody { + pub(super) fn new(bytes: Bytes, lease: Arc) -> Self { + Self { + bytes, + _lease: lease, + } + } +} + +impl AsRef<[u8]> for OwnedBatchBody { + fn as_ref(&self) -> &[u8] { + &self.bytes + } +}