WEB: Lifecycle + Lane ownership + Diag fixes

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

src/web/session/websocket.rs

@@ -4,7 +4,10 @@ use std::time::Instant;
 use sha2::{Digest, Sha256};
 
 use super::uplink::{AppliedProgress, inbound_reservation, validate_batch};
-use super::{PendingClass, WebSession, inbound_queue_cost, insert_carrier_lane};
+use super::{
+    CarrierLaneIdentity, PendingClass, StreamIdentity, WebSession, WebSocketLaneClaim,
+    inbound_queue_cost, insert_carrier_lane,
+};
 use crate::config::WebCarrier;
 use crate::web::frame;
 use crate::web::manager::ManagerError;
@@ -12,9 +15,19 @@ use crate::web::manager::ManagerError;
 /// Pre-OPEN stream quota and synthetic tuple ownership for one WebSocket lane.
 pub(crate) struct WebSocketLaneReservation {
     session: Arc<WebSession>,
-    lane_id: u32,
-    peer_port: u16,
-    transferred: bool,
+    claim: WebSocketLaneClaim,
+    stream: Option<StreamIdentity>,
+    phase: WebSocketLaneReservationPhase,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum WebSocketLaneReservationPhase {
+    Reserved,
+    Bound,
+    Transferred,
+    StreamOwned,
+    Closing,
+    Released,
 }
 
 /// Session-wide ownership of the only automatic WebSocket carrier probe.
@@ -57,28 +70,99 @@ impl Drop for WebSocketProbeReservation {
 impl WebSocketLaneReservation {
     /// Returns the logical stream owned by this connection.
     pub(crate) fn lane_id(&self) -> u32 {
-        self.lane_id
+        self.claim.lane.lane_id
     }
 
-    fn transfer_to_stream(&mut self) {
-        let removed = self
-            .session
-            .state
-            .lock()
-            .websocket_lane_reservations
-            .remove(&self.lane_id);
-        if removed == Some(self.peer_port) {
-            self.transferred = true;
+    /// Returns the exact lane incarnation owned by this connection.
+    pub(crate) fn lane_identity(&self) -> CarrierLaneIdentity {
+        self.claim.lane
+    }
+
+    /// Binds this pre-upgrade reservation to one admitted process connection.
+    pub(crate) fn bind(&mut self, connection_id: u64) -> Result<(), ManagerError> {
+        if self.phase != WebSocketLaneReservationPhase::Reserved {
+            return Err(ManagerError::Concurrent);
         }
+        let mut state = self.session.state.lock();
+        if state.closed
+            || state
+                .carrier_lanes
+                .get(&self.claim.lane.lane_id)
+                .is_none_or(|lane| lane.instance != self.claim.lane.instance)
+        {
+            return Err(ManagerError::Closed);
+        }
+        let Some(current) = state
+            .websocket_lane_reservations
+            .get_mut(&self.claim.lane.lane_id)
+            .filter(|current| **current == self.claim && current.connection_id.is_none())
+        else {
+            return Err(ManagerError::Closed);
+        };
+        current.connection_id = Some(connection_id);
+        self.claim.connection_id = Some(connection_id);
+        self.phase = WebSocketLaneReservationPhase::Bound;
+        Ok(())
+    }
+
+    fn transfer_to_stream(&mut self, stream: StreamIdentity) -> Result<(), ManagerError> {
+        if self.phase != WebSocketLaneReservationPhase::Bound
+            || stream.id != self.claim.lane.lane_id
+        {
+            return Err(ManagerError::Protocol);
+        }
+        let mut state = self.session.state.lock();
+        if state
+            .carrier_lanes
+            .get(&self.claim.lane.lane_id)
+            .is_none_or(|lane| lane.instance != self.claim.lane.instance)
+            || state
+                .streams
+                .get(&stream.id)
+                .is_none_or(|current| current.instance != stream.instance)
+            || state
+                .websocket_lane_reservations
+                .get(&self.claim.lane.lane_id)
+                != Some(&self.claim)
+        {
+            return Err(ManagerError::Closed);
+        }
+        state
+            .websocket_lane_reservations
+            .remove(&self.claim.lane.lane_id);
+        self.stream = Some(stream);
+        self.phase = WebSocketLaneReservationPhase::Transferred;
+        Ok(())
+    }
+
+    fn mark_stream_owned(&mut self, stream: StreamIdentity) -> Result<(), ManagerError> {
+        if self.phase != WebSocketLaneReservationPhase::Transferred || self.stream != Some(stream) {
+            return Err(ManagerError::Protocol);
+        }
+        self.phase = WebSocketLaneReservationPhase::StreamOwned;
+        Ok(())
+    }
+
+    fn retain_after_rejected_spawn(&mut self) {
+        debug_assert_eq!(self.phase, WebSocketLaneReservationPhase::StreamOwned);
+        self.phase = WebSocketLaneReservationPhase::Transferred;
+    }
+
+    fn release(&mut self) {
+        if self.phase == WebSocketLaneReservationPhase::Released {
+            return;
+        }
+        let stream_owned = self.phase == WebSocketLaneReservationPhase::StreamOwned;
+        self.phase = WebSocketLaneReservationPhase::Closing;
+        self.session
+            .release_websocket_lane_claim(self.claim, self.stream, stream_owned);
+        self.phase = WebSocketLaneReservationPhase::Released;
     }
 }
 
 impl Drop for WebSocketLaneReservation {
     fn drop(&mut self) {
-        if !self.transferred {
-            self.session
-                .release_websocket_lane_reservation(self.lane_id, self.peer_port);
-        }
+        self.release();
     }
 }
 
@@ -162,9 +246,7 @@ impl WebSession {
             );
             return Err(ManagerError::Limit);
         }
-        state.websocket_lane_reservations.insert(lane_id, peer_port);
-        if insert_carrier_lane(&mut state, lane_id).is_none() {
-            state.websocket_lane_reservations.remove(&lane_id);
+        let Some(lane) = insert_carrier_lane(&mut state, lane_id) else {
             state.active_peer_ports.remove(&peer_port);
             manager.release_stream(
                 self.profile_key,
@@ -173,14 +255,37 @@ impl WebSession {
                 peer_port,
             );
             return Err(ManagerError::Protocol);
+        };
+        let claim = WebSocketLaneClaim {
+            lane,
+            peer_port,
+            connection_id: None,
+        };
+        let inserted = match state.websocket_lane_reservations.entry(lane_id) {
+            std::collections::hash_map::Entry::Vacant(entry) => {
+                entry.insert(claim);
+                true
+            }
+            std::collections::hash_map::Entry::Occupied(_) => false,
+        };
+        if !inserted {
+            self.release_lane_locked(&mut state, lane_id);
+            state.active_peer_ports.remove(&peer_port);
+            manager.release_stream(
+                self.profile_key,
+                self.client_ip,
+                self.profile.public_addr,
+                peer_port,
+            );
+            return Err(ManagerError::Concurrent);
         }
         drop(state);
         self.lane_open_notify.notify_waiters();
         Ok(WebSocketLaneReservation {
             session: Arc::clone(self),
-            lane_id,
-            peer_port,
-            transferred: false,
+            claim,
+            stream: None,
+            phase: WebSocketLaneReservationPhase::Reserved,
         })
     }
 
@@ -192,12 +297,16 @@ impl WebSession {
         body: &[u8],
     ) -> Result<bool, ManagerError> {
         if !Arc::ptr_eq(self, &reservation.session)
-            || reservation.lane_id == 0
-            || reservation.lane_id > frame::MAX_STREAM_ID
+            || reservation.lane_id() == 0
+            || reservation.lane_id() > frame::MAX_STREAM_ID
+            || !matches!(
+                reservation.phase,
+                WebSocketLaneReservationPhase::Bound | WebSocketLaneReservationPhase::StreamOwned
+            )
         {
             return Err(ManagerError::Protocol);
         }
-        let lane_id = reservation.lane_id;
+        let lane_id = reservation.lane_id();
         let frames = frame::parse_all(body, &self.limits).map_err(|_| ManagerError::Protocol)?;
         if frames
             .iter()
@@ -216,8 +325,19 @@ impl WebSession {
                 return Err(ManagerError::Closed);
             }
             self.ensure_carrier_active_locked(&state)?;
-            if !reservation.transferred
-                && state.websocket_lane_reservations.get(&lane_id) != Some(&reservation.peer_port)
+            if state
+                .carrier_lanes
+                .get(&lane_id)
+                .is_none_or(|lane| lane.instance != reservation.claim.lane.instance)
+                || (reservation.phase == WebSocketLaneReservationPhase::Bound
+                    && state.websocket_lane_reservations.get(&lane_id) != Some(&reservation.claim))
+                || (reservation.phase == WebSocketLaneReservationPhase::StreamOwned
+                    && reservation.stream.is_none_or(|stream| {
+                        state
+                            .streams
+                            .get(&stream.id)
+                            .is_none_or(|current| current.instance != stream.instance)
+                    }))
             {
                 return Err(ManagerError::Closed);
             }
@@ -252,8 +372,8 @@ impl WebSession {
             let mut unused_bytes = reserve_bytes;
             let mut unused_items = reserve_items;
             let mut progress = AppliedProgress::default();
-            let mut reserved_open =
-                (!reservation.transferred).then_some((lane_id, reservation.peer_port));
+            let mut reserved_open = (reservation.phase == WebSocketLaneReservationPhase::Bound)
+                .then_some((lane_id, reservation.claim.peer_port));
             let applied = self.apply_batch_locked(
                 &mut state,
                 &frames,
@@ -287,15 +407,18 @@ impl WebSession {
             self.finish_carrier_health();
         }
         for completion in opened {
-            if completion.stream.id != lane_id || completion.peer_port != reservation.peer_port {
+            let stream = completion.stream;
+            if stream.id != lane_id || completion.peer_port != reservation.claim.peer_port {
                 return Err(ManagerError::Protocol);
             }
+            reservation.transfer_to_stream(stream)?;
+            reservation.mark_stream_owned(stream)?;
             if !self.spawn_stream(completion, true) {
+                reservation.retain_after_rejected_spawn();
                 return Err(ManagerError::Limit);
             }
-            reservation.transfer_to_stream();
         }
-        if !reservation.transferred {
+        if reservation.phase != WebSocketLaneReservationPhase::StreamOwned {
             return Err(ManagerError::Protocol);
         }
         if let Some(manager) = self.manager.upgrade() {
@@ -304,49 +427,87 @@ impl WebSession {
         Ok(progressed)
     }
 
-    /// Ends one failed or disconnected lane without closing its parent session.
-    pub(crate) fn close_websocket_lane(&self, lane_id: u32) {
-        let reserved = {
-            let mut state = self.state.lock();
-            let reserved = state.websocket_lane_reservations.remove(&lane_id);
-            if let Some(stream) = state.streams.remove(&lane_id) {
-                state.closing_streams.insert(lane_id, stream.instance);
-                let (bytes, items) = inbound_queue_cost(&stream.inbound);
-                self.release_locked(&mut state, bytes, items, false);
-                if let Some(waker) = stream.read_waker {
-                    waker.wake();
-                }
-                if let Some(waker) = stream.write_waker {
-                    waker.wake();
-                }
-            }
-            self.remember_closed_locked(&mut state, lane_id);
-            self.release_lane_locked(&mut state, lane_id);
-            reserved
-        };
-        if let Some(peer_port) = reserved {
-            self.release_websocket_lane_reservation(lane_id, peer_port);
+    /// Ends one exact failed or disconnected lane without closing its parent session.
+    pub(crate) fn close_websocket_lane(&self, mut reservation: WebSocketLaneReservation) {
+        if std::ptr::eq(self, Arc::as_ptr(&reservation.session)) {
+            reservation.release();
         }
-        self.lane_open_notify.notify_waiters();
     }
 
-    fn release_websocket_lane_reservation(&self, lane_id: u32, peer_port: u16) {
-        let removed = {
+    fn release_websocket_lane_claim(
+        &self,
+        claim: WebSocketLaneClaim,
+        stream: Option<StreamIdentity>,
+        stream_owned: bool,
+    ) {
+        let release_port = {
             let mut state = self.state.lock();
-            if state.websocket_lane_reservations.get(&lane_id) == Some(&peer_port) {
-                state.websocket_lane_reservations.remove(&lane_id);
+            let lane_matches = state
+                .carrier_lanes
+                .get(&claim.lane.lane_id)
+                .is_some_and(|lane| lane.instance == claim.lane.instance);
+            if !lane_matches && !state.closed {
+                return;
             }
-            self.release_lane_locked(&mut state, lane_id);
-            state.active_peer_ports.remove(&peer_port)
+            let release_port = if let Some(stream) = stream {
+                if stream.id != claim.lane.lane_id {
+                    return;
+                }
+                let current_stream = state
+                    .streams
+                    .get(&stream.id)
+                    .is_some_and(|current| current.instance == stream.instance);
+                if current_stream {
+                    let Some(stream_state) = state.streams.remove(&stream.id) else {
+                        return;
+                    };
+                    state
+                        .closing_streams
+                        .insert(claim.lane.lane_id, stream.instance);
+                    let (bytes, items) = inbound_queue_cost(&stream_state.inbound);
+                    self.release_locked(&mut state, bytes, items, false);
+                    if let Some(waker) = stream_state.read_waker {
+                        waker.wake();
+                    }
+                    if let Some(waker) = stream_state.write_waker {
+                        waker.wake();
+                    }
+                    false
+                } else if stream_owned {
+                    false
+                } else {
+                    state.active_peer_ports.remove(&claim.peer_port)
+                }
+            } else {
+                if state.websocket_lane_reservations.get(&claim.lane.lane_id) != Some(&claim) {
+                    return;
+                }
+                state
+                    .websocket_lane_reservations
+                    .remove(&claim.lane.lane_id);
+                state.active_peer_ports.remove(&claim.peer_port)
+            };
+            if lane_matches {
+                self.remember_closed_locked(&mut state, claim.lane.lane_id);
+                if state
+                    .carrier_lanes
+                    .get(&claim.lane.lane_id)
+                    .is_some_and(|lane| lane.instance == claim.lane.instance)
+                {
+                    self.release_lane_locked(&mut state, claim.lane.lane_id);
+                }
+            }
+            release_port
         };
-        if removed && let Some(manager) = self.manager.upgrade() {
+        if release_port && let Some(manager) = self.manager.upgrade() {
             manager.release_stream(
                 self.profile_key,
                 self.client_ip,
                 self.profile.public_addr,
-                peer_port,
+                claim.peer_port,
             );
         }
+        self.lane_open_notify.notify_waiters();
     }
 }