Session Residence + Ownership

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-08-26 18:44:17 +03:00
parent 016ad247a4
commit 8b2b88f30c
10 changed files with 436 additions and 71 deletions
+2
View File
@@ -284,11 +284,13 @@ pub(super) async fn handle(
.admit_websocket( .admit_websocket(
session.profile_key(), session.profile_key(),
session.trace_session_id(), session.trace_session_id(),
session.token_hash(),
effective_ip, effective_ip,
kind, kind,
BASE_BUDGET_BYTES, BASE_BUDGET_BYTES,
Duration::from_secs(timeouts.long_poll_secs), Duration::from_secs(timeouts.long_poll_secs),
Duration::from_secs(timeouts.websocket_eviction_secs), Duration::from_secs(timeouts.websocket_eviction_secs),
session.carrier_cancellation(),
) )
.await .await
{ {
+72 -15
View File
@@ -30,7 +30,16 @@ pub(super) async fn run_upgraded(
trace: Option<TraceWebSocketContext>, trace: Option<TraceWebSocketContext>,
acknowledge_commit: bool, 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; return;
}; };
let Ok(parts) = upgraded.downcast::<TokioIo<ConnectionIo>>() else { let Ok(parts) = upgraded.downcast::<TokioIo<ConnectionIo>>() else {
@@ -51,7 +60,6 @@ pub(super) async fn run_upgraded(
.max_frame_size(Some(limits.carrier_batch_bytes)); .max_frame_size(Some(limits.carrier_batch_bytes));
let mut socket = WebSocketStream::from_raw_socket(io, Role::Server, Some(config)).await; let mut socket = WebSocketStream::from_raw_socket(io, Role::Server, Some(config)).await;
connection.mark_opened(); connection.mark_opened();
let cancellation = connection.cancellation();
if let Some(reservation) = lane_reservation.as_mut() { if let Some(reservation) = lane_reservation.as_mut() {
let _ = run_lane( let _ = run_lane(
&mut socket, &mut socket,
@@ -84,7 +92,9 @@ pub(super) async fn run_upgraded(
.timeouts .timeouts
.websocket_eviction_secs, .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 { if let Some(reservation) = lane_reservation {
session.close_websocket_lane(reservation.lane_id()); session.close_websocket_lane(reservation.lane_id());
drop(reservation); drop(reservation);
@@ -111,11 +121,22 @@ async fn run_multiplex(
let mut read_budget = None; let mut read_budget = None;
let liveness_interval = connection.liveness_interval(); let liveness_interval = connection.liveness_interval();
let mut next_ping = Instant::now() + 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 { loop {
let down = session.poll_down(cursor); let down = session.poll_down(cursor);
tokio::pin!(down); tokio::pin!(down);
let event = tokio::select! { let event = tokio::select! {
_ = cancellation.cancelled() => return Err(()), _ = cancellation.cancelled() => return Err(()),
_ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()),
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness, _ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
incoming = read_message( incoming = read_message(
socket, socket,
@@ -145,7 +166,13 @@ async fn run_multiplex(
result?; result?;
if acknowledge_commit && sequence == 1 && session.is_carrier_committed() { if acknowledge_commit && sequence == 1 && session.is_carrier_committed() {
let started = Instant::now(); let started = Instant::now();
send(socket, runtime, Message::Binary(Bytes::new())).await?; send(
socket,
runtime,
Message::Binary(Bytes::new()),
&cancellation,
)
.await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -155,6 +182,10 @@ async fn run_multiplex(
started, started,
); );
} }
if !active {
connection.mark_active();
active = true;
}
sequence = sequence.checked_add(1).ok_or(())?; sequence = sequence.checked_add(1).ok_or(())?;
connection.mark_peer_activity(); connection.mark_peer_activity();
next_ping = Instant::now() + liveness_interval; next_ping = Instant::now() + liveness_interval;
@@ -173,7 +204,7 @@ async fn run_multiplex(
} }
Message::Ping(payload) => { Message::Ping(payload) => {
let started = Instant::now(); let started = Instant::now();
flush(socket, runtime).await?; flush(socket, runtime, &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -220,7 +251,7 @@ async fn run_multiplex(
DriverEvent::Down(result) => { DriverEvent::Down(result) => {
if result.body.is_empty() { if result.body.is_empty() {
let started = Instant::now(); let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?; send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -241,7 +272,7 @@ async fn run_multiplex(
let body = result.body; let body = result.body;
let started = Instant::now(); let started = Instant::now();
if trace.is_some() { if trace.is_some() {
send(socket, runtime, Message::Binary(body.clone())).await?; send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -251,7 +282,7 @@ async fn run_multiplex(
started, started,
); );
} else { } else {
send(socket, runtime, Message::Binary(body)).await?; send(socket, runtime, Message::Binary(body), &cancellation).await?;
} }
connection.mark_progress(); connection.mark_progress();
} }
@@ -259,7 +290,7 @@ async fn run_multiplex(
} }
DriverEvent::Liveness => { DriverEvent::Liveness => {
let started = Instant::now(); let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?; send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -291,11 +322,22 @@ async fn run_lane(
let mut read_budget = None; let mut read_budget = None;
let liveness_interval = connection.liveness_interval(); let liveness_interval = connection.liveness_interval();
let mut next_ping = Instant::now() + 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 { loop {
let down = session.poll_down_lane(reservation.lane_id(), cursor); let down = session.poll_down_lane(reservation.lane_id(), cursor);
tokio::pin!(down); tokio::pin!(down);
let event = tokio::select! { let event = tokio::select! {
_ = cancellation.cancelled() => return Err(()), _ = cancellation.cancelled() => return Err(()),
_ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()),
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness, _ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
incoming = read_message( incoming = read_message(
socket, socket,
@@ -332,7 +374,18 @@ async fn run_lane(
result?; result?;
if acknowledge_commit && sequence == 1 && session.is_carrier_committed() { if acknowledge_commit && sequence == 1 && session.is_carrier_committed() {
let started = Instant::now(); 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( record_message(
runtime, runtime,
trace, trace,
@@ -342,6 +395,10 @@ async fn run_lane(
started, started,
); );
} }
if !active {
connection.mark_active();
active = true;
}
sequence = sequence.checked_add(1).ok_or(())?; sequence = sequence.checked_add(1).ok_or(())?;
connection.mark_peer_activity(); connection.mark_peer_activity();
next_ping = Instant::now() + liveness_interval; next_ping = Instant::now() + liveness_interval;
@@ -360,7 +417,7 @@ async fn run_lane(
} }
Message::Ping(payload) => { Message::Ping(payload) => {
let started = Instant::now(); let started = Instant::now();
flush(socket, runtime).await?; flush(socket, runtime, &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -410,7 +467,7 @@ async fn run_lane(
} }
if result.body.is_empty() { if result.body.is_empty() {
let started = Instant::now(); let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?; send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -431,7 +488,7 @@ async fn run_lane(
let body = result.body; let body = result.body;
let started = Instant::now(); let started = Instant::now();
if trace.is_some() { if trace.is_some() {
send(socket, runtime, Message::Binary(body.clone())).await?; send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
@@ -441,7 +498,7 @@ async fn run_lane(
started, started,
); );
} else { } else {
send(socket, runtime, Message::Binary(body)).await?; send(socket, runtime, Message::Binary(body), &cancellation).await?;
} }
connection.mark_progress(); connection.mark_progress();
} }
@@ -449,7 +506,7 @@ async fn run_lane(
} }
DriverEvent::Liveness => { DriverEvent::Liveness => {
let started = Instant::now(); let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?; send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message( record_message(
runtime, runtime,
trace, trace,
+14 -8
View File
@@ -160,6 +160,7 @@ pub(super) async fn send(
socket: &mut CarrierSocket, socket: &mut CarrierSocket,
runtime: &WebProcessRuntime, runtime: &WebProcessRuntime,
message: Message, message: Message,
cancellation: &CancellationToken,
) -> Result<(), ()> { ) -> Result<(), ()> {
let timeout = Duration::from_secs( let timeout = Duration::from_secs(
runtime runtime
@@ -169,15 +170,18 @@ pub(super) async fn send(
.timeouts .timeouts
.websocket_write_secs, .websocket_write_secs,
); );
tokio::time::timeout(timeout, socket.send(message)) tokio::select! {
.await _ = cancellation.cancelled() => Err(()),
.map_err(|_| ())? result = tokio::time::timeout(timeout, socket.send(message)) => {
.map_err(|_| ()) result.map_err(|_| ())?.map_err(|_| ())
}
}
} }
pub(super) async fn flush( pub(super) async fn flush(
socket: &mut CarrierSocket, socket: &mut CarrierSocket,
runtime: &WebProcessRuntime, runtime: &WebProcessRuntime,
cancellation: &CancellationToken,
) -> Result<(), ()> { ) -> Result<(), ()> {
let timeout = Duration::from_secs( let timeout = Duration::from_secs(
runtime runtime
@@ -187,10 +191,12 @@ pub(super) async fn flush(
.timeouts .timeouts
.websocket_write_secs, .websocket_write_secs,
); );
tokio::time::timeout(timeout, socket.flush()) tokio::select! {
.await _ = cancellation.cancelled() => Err(()),
.map_err(|_| ())? result = tokio::time::timeout(timeout, socket.flush()) => {
.map_err(|_| ()) result.map_err(|_| ())?.map_err(|_| ())
}
}
} }
pub(super) fn record_message( pub(super) fn record_message(
+19
View File
@@ -290,6 +290,21 @@ impl WebProcessRuntime {
Some((reader, body)) Some((reader, body))
} }
/// Reserves transient bytes while one downlink batch replaces queued frames.
pub(crate) fn try_downlink_staging_budget(
&self,
bytes: usize,
) -> Option<OwnedSemaphorePermit> {
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. /// Reserves bounded process-wide queue capacity for data or control traffic.
pub(crate) fn try_reserve_pending( pub(crate) fn try_reserve_pending(
&self, &self,
@@ -351,21 +366,25 @@ impl WebProcessRuntime {
self: &Arc<Self>, self: &Arc<Self>,
owner: ProfileKey, owner: ProfileKey,
session_id: u64, session_id: u64,
session_hash: TokenHash,
client_ip: IpAddr, client_ip: IpAddr,
kind: WebSocketKind, kind: WebSocketKind,
base_bytes: usize, base_bytes: usize,
liveness_interval: Duration, liveness_interval: Duration,
eviction_timeout: Duration, eviction_timeout: Duration,
parent_cancellation: CancellationToken,
) -> Result<WebSocketConnection, ManagerError> { ) -> Result<WebSocketConnection, ManagerError> {
websocket::admit( websocket::admit(
self, self,
owner, owner,
session_id, session_id,
session_hash,
client_ip, client_ip,
kind, kind,
base_bytes, base_bytes,
liveness_interval, liveness_interval,
eviction_timeout, eviction_timeout,
parent_cancellation,
) )
.await .await
} }
-2
View File
@@ -163,7 +163,6 @@ impl WebDataBudget {
} }
remove_owner(&mut state.owner_bytes, owner, bytes); remove_owner(&mut state.owner_bytes, owner, bytes);
drop(state); drop(state);
self.pressured.store(false, Ordering::Release);
self.notify.notify_waiters(); self.notify.notify_waiters();
} }
@@ -259,7 +258,6 @@ impl WebDataBudget {
state.websocket_bytes = state.websocket_bytes.saturating_sub(bytes); state.websocket_bytes = state.websocket_bytes.saturating_sub(bytes);
remove_owner(&mut state.owner_bytes, owner, bytes); remove_owner(&mut state.owner_bytes, owner, bytes);
drop(state); drop(state);
self.pressured.store(false, Ordering::Release);
self.notify.notify_waiters(); self.notify.notify_waiters();
} }
} }
+78 -22
View File
@@ -275,20 +275,10 @@ impl WebProcessRuntime {
pub(super) fn cleanup_websockets(&self) { pub(super) fn cleanup_websockets(&self) {
let now = self.websocket_tick(); let now = self.websocket_tick();
let mut victims = self let mut victims = claim_stale_victims(self, now);
.websockets
.lock()
.entries
.values()
.filter(|entry| {
now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire))
>= dead_after(entry)
})
.cloned()
.collect::<Vec<_>>();
if victims.is_empty() if victims.is_empty()
&& self.data_budget.take_pressure() && 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); victims.push(victim);
} }
@@ -315,16 +305,22 @@ fn select_victim(
session_id: u64, session_id: u64,
client_ip: IpAddr, client_ip: IpAddr,
excluded_id: Option<u64>, excluded_id: Option<u64>,
claim: bool,
) -> Option<Arc<WebSocketEntry>> { ) -> Option<Arc<WebSocketEntry>> {
let fair_share = runtime.data_budget.fair_share(Some(owner)); let fair_share = runtime.data_budget.fair_share(Some(owner));
let requester_usage = runtime.data_budget.owner_usage(owner); let requester_usage = runtime.data_budget.owner_usage(owner);
let now = runtime.websocket_tick(); let now = runtime.websocket_tick();
runtime let mut registry = runtime.websockets.lock();
.websockets if claim
.lock() && registry.evictions_in_flight >= runtime.limits.max_websocket_evictions_in_flight
{
return None;
}
let selected = registry
.entries .entries
.values() .values()
.filter(|entry| Some(entry.id) != excluded_id) .filter(|entry| Some(entry.id) != excluded_id)
.filter(|entry| !entry.closing.load(Ordering::Acquire))
.filter_map(|entry| { .filter_map(|entry| {
let owner_rank = if entry.session_id == session_id { let owner_rank = if entry.session_id == session_id {
0 0
@@ -353,15 +349,28 @@ fn select_victim(
)) ))
}) })
.min_by_key(|(key, _)| *key) .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<Arc<WebSocketEntry>> { fn select_pressure_victim(
runtime runtime: &WebProcessRuntime,
.websockets now: u64,
.lock() claim: bool,
) -> Option<Arc<WebSocketEntry>> {
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 .entries
.values() .values()
.filter(|entry| !entry.closing.load(Ordering::Acquire))
.map(|entry| { .map(|entry| {
( (
( (
@@ -374,11 +383,58 @@ fn select_pressure_victim(runtime: &WebProcessRuntime, now: u64) -> Option<Arc<W
) )
}) })
.min_by_key(|(key, _)| *key) .min_by_key(|(key, _)| *key)
.map(|(_, entry)| entry) .map(|(_, entry)| entry)?;
if claim && !claim_entry(&mut registry, &selected, runtime) {
return None;
}
Some(selected)
}
fn claim_stale_victims(runtime: &WebProcessRuntime, now: u64) -> Vec<Arc<WebSocketEntry>> {
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::<Vec<_>>();
candidates
.into_iter()
.filter(|entry| claim_entry(&mut registry, entry, runtime))
.collect()
}
fn claim_entry(
registry: &mut WebSocketRegistry,
entry: &Arc<WebSocketEntry>,
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 { 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) || now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire)) >= dead_after(entry)
{ {
0 0
+12
View File
@@ -21,6 +21,8 @@ use crate::web::manager::CarrierLearningContext;
mod backend; mod backend;
// Downlink queues own cursor replay, flow control, and memory reservations. // Downlink queues own cursor replay, flow control, and memory reservations.
mod downlink; 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. // Lane carrier state isolates request sequencing and downlink replay per logical stream.
mod lanes; mod lanes;
// WebSocket carrier state owns pre-OPEN lane reservations and failure isolation. // WebSocket carrier state owns pre-OPEN lane reservations and failure isolation.
@@ -71,6 +73,7 @@ struct QueuedFrame {
struct DownBatch { struct DownBatch {
body: Bytes, body: Bytes,
lease: Arc<resident::PendingResponseLease>,
base_cursor: u64, base_cursor: u64,
next_cursor: u64, next_cursor: u64,
data_bytes: usize, data_bytes: usize,
@@ -83,6 +86,7 @@ struct CarrierLane {
instance: u64, instance: u64,
pending_bytes: usize, pending_bytes: usize,
pending_items: usize, pending_items: usize,
resident: Arc<resident::ResidentCounters>,
pending_frames: VecDeque<QueuedFrame>, pending_frames: VecDeque<QueuedFrame>,
pending_windows: HashMap<u32, usize>, pending_windows: HashMap<u32, usize>,
unacked: Option<DownBatch>, unacked: Option<DownBatch>,
@@ -100,6 +104,7 @@ impl CarrierLane {
instance, instance,
pending_bytes: 0, pending_bytes: 0,
pending_items: 0, pending_items: 0,
resident: Arc::new(resident::ResidentCounters::default()),
pending_frames: VecDeque::new(), pending_frames: VecDeque::new(),
pending_windows: HashMap::new(), pending_windows: HashMap::new(),
unacked: None, unacked: None,
@@ -169,6 +174,7 @@ pub(crate) struct WebSession {
cancel: CancellationToken, cancel: CancellationToken,
tasks_live: AtomicUsize, tasks_live: AtomicUsize,
tasks_done: Arc<Notify>, tasks_done: Arc<Notify>,
resident: Arc<resident::ResidentCounters>,
finished: AtomicBool, finished: AtomicBool,
up_active: AtomicBool, up_active: AtomicBool,
} }
@@ -251,6 +257,7 @@ impl WebSession {
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
tasks_live: AtomicUsize::new(0), tasks_live: AtomicUsize::new(0),
tasks_done: Arc::new(Notify::new()), tasks_done: Arc::new(Notify::new()),
resident: Arc::new(resident::ResidentCounters::default()),
finished: AtomicBool::new(false), finished: AtomicBool::new(false),
up_active: AtomicBool::new(false), up_active: AtomicBool::new(false),
}) })
@@ -281,6 +288,11 @@ impl WebSession {
self.trace_session_id 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. /// Returns whether accepted carrier progress made this attempt immutable.
pub(crate) fn is_carrier_committed(&self) -> bool { pub(crate) fn is_carrier_committed(&self) -> bool {
self.state.lock().negotiation_phase == SessionNegotiationPhase::Committed self.state.lock().negotiation_phase == SessionNegotiationPhase::Committed
+54 -13
View File
@@ -1,7 +1,9 @@
use std::sync::Arc;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use bytes::{BufMut, Bytes, BytesMut}; use bytes::{BufMut, Bytes, BytesMut};
use super::resident::{OwnedBatchBody, PendingCounts, PendingResponseLease};
use super::{ use super::{
DownBatch, PendingClass, PollResult, QUEUE_ITEM_COST, QueuedFrame, SessionState, WebSession, DownBatch, PendingClass, PollResult, QUEUE_ITEM_COST, QueuedFrame, SessionState, WebSession,
}; };
@@ -60,6 +62,9 @@ impl WebSession {
if !state.pending_frames.is_empty() { if !state.pending_frames.is_empty() {
let batch = match self.take_down_batch_locked(&mut state, cursor) { let batch = match self.take_down_batch_locked(&mut state, cursor) {
Ok(batch) => batch, Ok(batch) => batch,
Err(ManagerError::Backpressure) => {
return Err(ManagerError::Backpressure);
}
Err(error) => { Err(error) => {
drop(state); drop(state);
self.close(); self.close();
@@ -121,6 +126,15 @@ impl WebSession {
.limits .limits
.pending_items_per_session .pending_items_per_session
.saturating_sub(item_reserve); .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 { if state.closed {
return false; return false;
} }
@@ -128,20 +142,16 @@ impl WebSession {
let fits = if control { let fits = if control {
bytes <= self.limits.control_bytes_per_session bytes <= self.limits.control_bytes_per_session
&& items <= item_reserve && items <= item_reserve
&& state.pending_bytes && pending_bytes
<= self.limits.pending_bytes_per_session.saturating_sub(bytes) <= self.limits.pending_bytes_per_session.saturating_sub(bytes)
&& state.pending_items && pending_items
<= self.limits.pending_items_per_session.saturating_sub(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) <= 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 { } else {
let data_bytes = state let data_bytes = pending_bytes.saturating_sub(pending_control_bytes);
.pending_bytes let data_items = pending_items.saturating_sub(pending_control_items);
.saturating_sub(state.pending_control_bytes);
let data_items = state
.pending_items
.saturating_sub(state.pending_control_items);
let (byte_limit, item_limit) = if class == PendingClass::Downlink { let (byte_limit, item_limit) = if class == PendingClass::Downlink {
let uplink_bytes = self.limits.max_body_bytes.saturating_add( let uplink_bytes = self.limits.max_body_bytes.saturating_add(
self.limits 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. /// Coalesces one flow-control update into the bounded control queue.
pub(super) fn queue_window_locked( pub(super) fn queue_window_locked(
&self, &self,
@@ -351,6 +376,12 @@ impl WebSession {
body_len += queued.encoded.len(); body_len += queued.encoded.len();
count += 1; 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 body = BytesMut::with_capacity(body_len);
let mut data_bytes = 0usize; let mut data_bytes = 0usize;
let mut data_items = 0usize; let mut data_items = 0usize;
@@ -383,8 +414,17 @@ impl WebSession {
*index = index.saturating_sub(count); *index = index.saturating_sub(count);
} }
state.down_cursor = next_cursor; 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 { Ok(DownBatch {
body: body.freeze(), body,
lease,
base_cursor: cursor, base_cursor: cursor,
next_cursor, next_cursor,
data_bytes, data_bytes,
@@ -398,8 +438,9 @@ impl WebSession {
let Some(batch) = state.unacked.take() else { let Some(batch) = state.unacked.take() else {
return; return;
}; };
self.release_locked(state, batch.data_bytes, batch.data_items, false); batch.lease.detach();
self.release_locked(state, batch.control_bytes, batch.control_items, true); 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() { for stream in state.streams.values_mut() {
if let Some(waker) = stream.write_waker.take() { if let Some(waker) = stream.write_waker.take() {
waker.wake(); waker.wake();
+38 -11
View File
@@ -6,6 +6,7 @@ use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq; use subtle::ConstantTimeEq;
use super::uplink::{inbound_reservation, validate_batch}; use super::uplink::{inbound_reservation, validate_batch};
use super::resident::{OwnedBatchBody, PendingCounts, PendingResponseLease};
use super::{ use super::{
CarrierLane, DownBatch, PendingClass, PollResult, QUEUE_ITEM_COST, QueuedFrame, SessionState, CarrierLane, DownBatch, PendingClass, PollResult, QUEUE_ITEM_COST, QueuedFrame, SessionState,
WebSession, insert_carrier_lane, remember_closed, WebSession, insert_carrier_lane, remember_closed,
@@ -241,8 +242,14 @@ impl WebSession {
.pending_items .pending_items
.saturating_sub(batch.data_items.saturating_add(batch.control_items)); .saturating_sub(batch.data_items.saturating_add(batch.control_items));
} }
self.release_locked(&mut state, batch.data_bytes, batch.data_items, false); batch.lease.detach();
self.release_locked(&mut state, batch.control_bytes, batch.control_items, true); 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) if let Some(stream) = state.streams.get_mut(&lane_id)
&& let Some(waker) = stream.write_waker.take() && let Some(waker) = stream.write_waker.take()
{ {
@@ -282,8 +289,11 @@ impl WebSession {
}); });
} }
if !lane.pending_frames.is_empty() { 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, Ok(batch) => batch,
Err(ManagerError::Backpressure) => {
return Err(ManagerError::Backpressure);
}
Err(error) => { Err(error) => {
drop(state); drop(state);
self.close(); self.close();
@@ -457,7 +467,8 @@ impl WebSession {
}); });
if can_coalesce { if can_coalesce {
if state.carrier_lanes.get(&stream_id).is_none_or(|lane| { 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 > self
.limits .limits
.pending_bytes_per_lane .pending_bytes_per_lane
@@ -491,9 +502,10 @@ impl WebSession {
PendingClass::Downlink PendingClass::Downlink
}; };
if state.carrier_lanes.get(&stream_id).is_none_or(|lane| { 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) > 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 >= self.limits.pending_items_per_lane
}) { }) {
return false; return false;
@@ -561,10 +573,9 @@ impl WebSession {
} }
} }
if let Some(batch) = lane.unacked.take() { if let Some(batch) = lane.unacked.take() {
data_bytes = data_bytes.saturating_add(batch.data_bytes); batch.lease.detach();
data_items = data_items.saturating_add(batch.data_items); self.release_local_locked(state, batch.data_bytes, batch.data_items, false);
control_bytes = control_bytes.saturating_add(batch.control_bytes); self.release_local_locked(state, batch.control_bytes, batch.control_items, true);
control_items = control_items.saturating_add(batch.control_items);
} }
self.release_locked(state, data_bytes, data_items, false); self.release_locked(state, data_bytes, data_items, false);
self.release_locked(state, control_bytes, control_items, true); 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( fn take_lane_down_batch(
session: &WebSession,
limits: &WebLimitsConfig, limits: &WebLimitsConfig,
lane: &mut CarrierLane, lane: &mut CarrierLane,
cursor: u64, cursor: u64,
@@ -613,6 +625,12 @@ fn take_lane_down_batch(
body_len += queued.encoded.len(); body_len += queued.encoded.len();
count += 1; 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 body = BytesMut::with_capacity(body_len);
let mut data_bytes = 0usize; let mut data_bytes = 0usize;
let mut data_items = 0usize; let mut data_items = 0usize;
@@ -645,8 +663,17 @@ fn take_lane_down_batch(
*index = index.saturating_sub(count); *index = index.saturating_sub(count);
} }
lane.down_cursor = next_cursor; 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 { Ok(DownBatch {
body: body.freeze(), body,
lease,
base_cursor: cursor, base_cursor: cursor,
next_cursor, next_cursor,
data_bytes, data_bytes,
+147
View File
@@ -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<WebProcessRuntime>,
owner: ProfileKey,
counts: PendingCounts,
session: Arc<ResidentCounters>,
lane: Option<Arc<ResidentCounters>>,
detached: AtomicBool,
}
impl PendingResponseLease {
pub(super) fn new(
session: &WebSession,
counts: PendingCounts,
lane: Option<Arc<ResidentCounters>>,
) -> Arc<Self> {
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<PendingResponseLease>,
}
impl OwnedBatchBody {
pub(super) fn new(bytes: Bytes, lease: Arc<PendingResponseLease>) -> Self {
Self {
bytes,
_lease: lease,
}
}
}
impl AsRef<[u8]> for OwnedBatchBody {
fn as_ref(&self) -> &[u8] {
&self.bytes
}
}