Pending Bytes&Items in Carrier Lanes

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

src/web/manager/websocket.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::net::IpAddr;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
 use std::time::Duration;
 
 use tokio::sync::OwnedSemaphorePermit;
@@ -10,7 +10,7 @@ use tokio_util::sync::CancellationToken;
 use super::{ManagerError, ProfileKey, WebProcessRuntime, WebSocketBudgetLease};
 
 /// One process-owned WebSocket carrier class used for eviction priority.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub(crate) enum WebSocketKind {
     /// One connection multiplexes every logical stream in a session.
     Multiplex,
@@ -18,23 +18,42 @@ pub(crate) enum WebSocketKind {
     Lane(u32),
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+struct WebSocketClaimKey {
+    session_hash: super::TokenHash,
+    kind: WebSocketKind,
+}
+
+#[repr(u8)]
+enum WebSocketPhase {
+    Claimed,
+    Upgraded,
+    Active,
+    Closing,
+}
+
 pub(super) struct WebSocketEntry {
     id: u64,
     owner: ProfileKey,
     session_id: u64,
+    claim: WebSocketClaimKey,
     client_ip: IpAddr,
     kind: WebSocketKind,
     liveness_interval_ms: u64,
     created_tick: u64,
     last_peer_tick: AtomicU64,
     last_progress_tick: AtomicU64,
-    opened: AtomicBool,
+    phase: AtomicU8,
+    closing: AtomicBool,
     cancel: CancellationToken,
+    released: CancellationToken,
 }
 
 #[derive(Default)]
 pub(super) struct WebSocketRegistry {
     entries: HashMap<u64, Arc<WebSocketEntry>>,
+    claims: HashMap<WebSocketClaimKey, u64>,
+    evictions_in_flight: usize,
     closed: bool,
 }
 
@@ -64,10 +83,20 @@ impl WebSocketConnection {
 
     /// Marks successful ownership transfer from HTTP to the WebSocket codec.
     pub(crate) fn mark_opened(&self) {
-        self.entry.opened.store(true, Ordering::Release);
+        self.entry
+            .phase
+            .store(WebSocketPhase::Upgraded as u8, Ordering::Release);
         self.mark_progress();
     }
 
+    /// Marks the first validated carrier binary message as active progress.
+    pub(crate) fn mark_active(&self) {
+        self.entry
+            .phase
+            .store(WebSocketPhase::Active as u8, Ordering::Release);
+        self.mark_peer_activity();
+    }
+
     /// Refreshes the peer-liveness deadline after any received WebSocket message.
     pub(crate) fn mark_peer_activity(&self) {
         if let Some(runtime) = self.runtime.upgrade() {
@@ -90,7 +119,16 @@ impl WebSocketConnection {
 impl Drop for WebSocketConnection {
     fn drop(&mut self) {
         if let Some(runtime) = self.runtime.upgrade() {
-            runtime.websockets.lock().entries.remove(&self.entry.id);
+            let mut registry = runtime.websockets.lock();
+            registry.entries.remove(&self.entry.id);
+            if registry.claims.get(&self.entry.claim) == Some(&self.entry.id) {
+                registry.claims.remove(&self.entry.claim);
+            }
+            if self.entry.closing.load(Ordering::Acquire) {
+                registry.evictions_in_flight = registry.evictions_in_flight.saturating_sub(1);
+            }
+            drop(registry);
+            self.entry.released.cancel();
             drop(self.base_budget.take());
             drop(self.slot.take());
             runtime.websocket_notify.notify_waiters();
@@ -103,81 +141,126 @@ pub(super) async fn admit(
     runtime: &Arc<WebProcessRuntime>,
     owner: ProfileKey,
     session_id: u64,
+    session_hash: super::TokenHash,
     client_ip: IpAddr,
     kind: WebSocketKind,
     base_bytes: usize,
     liveness_interval: Duration,
     eviction_timeout: Duration,
+    parent_cancellation: CancellationToken,
 ) -> Result<WebSocketConnection, ManagerError> {
     let liveness_interval_ms = liveness_interval.as_millis().min(u128::from(u64::MAX)) as u64;
-    if let Some(connection) = try_admit(
+    match try_admit(
         runtime,
         owner,
         session_id,
+        session_hash,
         client_ip,
         kind,
         base_bytes,
         liveness_interval_ms,
+        &parent_cancellation,
     ) {
-        return Ok(connection);
+        Ok(connection) => return Ok(connection),
+        Err(TryAdmitError::Conflict) => return Err(ManagerError::Concurrent),
+        Err(TryAdmitError::Closed) => return Err(ManagerError::Closed),
+        Err(TryAdmitError::Capacity) => {}
     }
-    let Some(victim) = select_victim(runtime, owner, session_id, client_ip, None) else {
+    let Some(victim) = select_victim(runtime, owner, session_id, client_ip, None, true) else {
         runtime.record_limit_hit();
         return Err(ManagerError::Limit);
     };
-    let released = runtime.websocket_notify.notified();
+    let released = victim.released.cancelled();
     victim.cancel.cancel();
     let _ = tokio::time::timeout(eviction_timeout, released).await;
-    try_admit(
+    match try_admit(
         runtime,
         owner,
         session_id,
+        session_hash,
         client_ip,
         kind,
         base_bytes,
         liveness_interval_ms,
-    )
-    .ok_or_else(|| {
-        runtime.record_limit_hit();
-        ManagerError::Limit
-    })
+        &parent_cancellation,
+    ) {
+        Ok(connection) => Ok(connection),
+        Err(TryAdmitError::Conflict) => Err(ManagerError::Concurrent),
+        Err(TryAdmitError::Closed) => Err(ManagerError::Closed),
+        Err(TryAdmitError::Capacity) => {
+            runtime.record_limit_hit();
+            Err(ManagerError::Limit)
+        }
+    }
+}
+
+enum TryAdmitError {
+    Capacity,
+    Conflict,
+    Closed,
 }
 
 fn try_admit(
     runtime: &Arc<WebProcessRuntime>,
     owner: ProfileKey,
     session_id: u64,
+    session_hash: super::TokenHash,
     client_ip: IpAddr,
     kind: WebSocketKind,
     base_bytes: usize,
     liveness_interval_ms: u64,
-) -> Option<WebSocketConnection> {
+    parent_cancellation: &CancellationToken,
+) -> Result<WebSocketConnection, TryAdmitError> {
+    let claim = WebSocketClaimKey { session_hash, kind };
+    {
+        let registry = runtime.websockets.lock();
+        if registry.closed {
+            return Err(TryAdmitError::Closed);
+        }
+        if registry.claims.contains_key(&claim) {
+            return Err(TryAdmitError::Conflict);
+        }
+    }
     let slot = Arc::clone(&runtime.websocket_connections)
         .try_acquire_owned()
-        .ok()?;
-    let base_budget = runtime.try_websocket_base_budget(owner, base_bytes)?;
-    let id = runtime.websocket_next_id.fetch_add(1, Ordering::Relaxed);
+        .map_err(|_| TryAdmitError::Capacity)?;
+    let base_budget = runtime
+        .try_websocket_base_budget(owner, base_bytes)
+        .ok_or(TryAdmitError::Capacity)?;
+    let id = runtime
+        .websocket_next_id
+        .fetch_update(Ordering::AcqRel, Ordering::Acquire, |value| {
+            value.checked_add(1)
+        })
+        .map_err(|_| TryAdmitError::Capacity)?;
     let now = runtime.websocket_tick();
     let entry = Arc::new(WebSocketEntry {
         id,
         owner,
         session_id,
+        claim,
         client_ip,
         kind,
         liveness_interval_ms,
         created_tick: now,
         last_peer_tick: AtomicU64::new(now),
         last_progress_tick: AtomicU64::new(now),
-        opened: AtomicBool::new(false),
-        cancel: CancellationToken::new(),
+        phase: AtomicU8::new(WebSocketPhase::Claimed as u8),
+        closing: AtomicBool::new(false),
+        cancel: parent_cancellation.child_token(),
+        released: CancellationToken::new(),
     });
     let mut registry = runtime.websockets.lock();
     if registry.closed {
-        return None;
+        return Err(TryAdmitError::Closed);
     }
+    if registry.claims.contains_key(&claim) {
+        return Err(TryAdmitError::Conflict);
+    }
+    registry.claims.insert(claim, id);
     registry.entries.insert(id, Arc::clone(&entry));
     drop(registry);
-    Some(WebSocketConnection {
+    Ok(WebSocketConnection {
         runtime: Arc::downgrade(runtime),
         entry,
         slot: Some(slot),