Session Residence + Ownership

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

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();

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,

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<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
+    }
+}