astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-09-05 18:16:06 +03:00
Code Issues Packages Projects Releases Wiki Activity

Session Residence + Ownership

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
src/web/http/websocket.rs
+2
View File
@@ -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
{
src/web/http/websocket/driver.rs
+71 -14
View File
@@ -30,7 +30,16 @@ pub(super) async fn run_upgraded(
trace: Option<TraceWebSocketContext>,
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::<TokioIo<ConnectionIo>>() 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,
);
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,
src/web/http/websocket/driver/io.rs
+14 -8
View File
@@ -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(
src/web/manager.rs
+19
View File
@@ -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<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.
pub(crate) fn try_reserve_pending(
&self,
@@ -351,21 +366,25 @@ impl WebProcessRuntime {
self: &Arc<Self>,
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<WebSocketConnection, ManagerError> {
websocket::admit(
self,
owner,
session_id,
session_hash,
client_ip,
kind,
base_bytes,
liveness_interval,
eviction_timeout,
parent_cancellation,
)
.await
}
src/web/manager/budget.rs
-2
View File
@@ -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();
}
}
src/web/manager/websocket.rs
+78 -22
View File
@@ -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::<Vec<_>>();
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<u64>,
claim: bool,
) -> Option<Arc<WebSocketEntry>> {
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<Arc<WebSocketEntry>> {
runtime
.websockets
.lock()
fn select_pressure_victim(
runtime: &WebProcessRuntime,
now: u64,
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
.values()
.filter(|entry| !entry.closing.load(Ordering::Acquire))
.map(|entry| {
(
(
@@ -374,11 +383,58 @@ fn select_pressure_victim(runtime: &WebProcessRuntime, now: u64) -> Option<Arc<W
)
})
.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 {
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
src/web/session.rs
+12
View File
@@ -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<resident::PendingResponseLease>,
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<resident::ResidentCounters>,
pending_frames: VecDeque<QueuedFrame>,
pending_windows: HashMap<u32, usize>,
unacked: Option<DownBatch>,
@@ -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<Notify>,
resident: Arc<resident::ResidentCounters>,
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
src/web/session/downlink.rs
+54 -13
View File
@@ -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();
src/web/session/lanes.rs
+38 -11
View File
@@ -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,
src/web/session/resident.rs
+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
}
}