fix(websocket): retain WsPumps handle on the server path (WS-08)

src/websocket/upgrade.rs

@@ -10,6 +10,8 @@
 //! `CallConnection` (the identity rides the channels-layer
 //! `Connection`) and runs the shared `Dispatcher::run_loop_single_stream`.
 
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use alkcall::channels::adapter::ChannelsAdapter;
@@ -20,19 +22,137 @@ use alkcall::registry::registration::OperationRegistry;
 use axum::extract::ws::WebSocketUpgrade;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
+use parking_lot::Mutex;
 
-use super::byte_adapter::{split_ws_to_bytes, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP};
+use super::byte_adapter::{
+    split_ws_to_bytes, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
+};
+
+/// Registry of live WS session pump handles (WS-08). The upgrade
+/// handler retains each session's [`WsPumps`] handle here for the
+/// session's lifetime so a stuck session is evictable in-crate:
+/// [`WsSessions::abort`] forces the pumps' tasks to end (the socket
+/// halves close and the channels session unwinds). The handle for a
+/// session is removed when the session task finishes (self-removing
+/// guard), so the registry only holds live sessions. Assembly layers
+/// share one instance via [`crate::server::state::RouterState`] (the
+/// upgrade handler registers against it) or per-route request
+/// extensions (an extension clone takes precedence).
+///
+/// Un-registered (default) — the upgrade runs fine and simply keeps no
+/// eviction lever, matching the pre-WS-08 behavior.
+#[derive(Clone, Default)]
+pub struct WsSessions {
+    counter: Arc<AtomicU64>,
+    sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
+}
+
+/// The upgrade handler's state slice: what it needs beyond the request
+/// itself. Axum lifts it via `FromRef` from either full router state —
+/// the adapter's [`crate::server::state::RouterState`] (carrying the
+/// shared [`WsSessions`] instance) or a bare `Arc<OperationRegistry>`
+/// (custom upgrade routes / integration tests get a handler-private
+/// registry; eviction still works in-crate, just not shared).
+#[derive(Clone)]
+pub struct SessionState {
+    registry: Arc<OperationRegistry>,
+    sessions: Arc<WsSessions>,
+}
+
+impl SessionState {
+    /// From the plain registry state (custom upgrade routes /
+    /// integration tests): sessions default to a fresh [`WsSessions`]
+    /// private to the handler (eviction still functional in-crate but
+    /// not shared with the assembly layer).
+    pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
+        Self {
+            registry: Arc::clone(registry),
+            sessions: Arc::new(WsSessions::new()),
+        }
+    }
+
+    /// From the adapter's router state: the shared [`WsSessions`]
+    /// instance the assembly layer can hold for eviction.
+    pub(crate) fn new(registry: Arc<OperationRegistry>, sessions: Arc<WsSessions>) -> Self {
+        Self { registry, sessions }
+    }
+
+    pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
+        &self.registry
+    }
+
+    pub(crate) fn sessions(&self) -> &Arc<WsSessions> {
+        &self.sessions
+    }
+}
+
+impl axum::extract::FromRef<SessionState> for Arc<OperationRegistry> {
+    fn from_ref(state: &SessionState) -> Self {
+        Arc::clone(&state.registry)
+    }
+}
+
+/// `FromRef` chain: a bare `Arc<OperationRegistry>` router state lifts
+/// into the handler's [`SessionState`]; a full [`RouterState`] carries
+/// the shared [`WsSessions`] instance and lifts through its own impl.
+impl axum::extract::FromRef<Arc<OperationRegistry>> for SessionState {
+    fn from_ref(registry: &Arc<OperationRegistry>) -> Self {
+        SessionState::from_registry(registry)
+    }
+}
+
+impl WsSessions {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Abort every live session's pump tasks (forced teardown).
+    pub fn abort(&self) {
+        for (_, pumps) in self.sessions.lock().drain() {
+            pumps.abort();
+        }
+    }
+
+    /// Number of live sessions currently tracked.
+    pub fn len(&self) -> usize {
+        self.sessions.lock().len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.sessions.lock().is_empty()
+    }
+
+    fn insert(&self, pumps: Arc<WsPumps>) -> u64 {
+        let id = self.counter.fetch_add(1, Ordering::Relaxed);
+        self.sessions.lock().insert(id, pumps);
+        id
+    }
+
+    fn remove(&self, id: u64) {
+        self.sessions.lock().remove(&id);
+    }
+}
 
 /// The channels session for an upgraded socket: adapt → `Connection`
 /// (identity attached) → `ChannelsAdapter::handle`. `policy` gates
-/// data-channel opens (ADR-041).
+/// data-channel opens (ADR-041). When `sessions` is `Some`, the
+/// session's pump handle is registered for the session's lifetime —
+/// the WS-08 eviction lever ([`WsSessions::abort`]).
 pub async fn run_channels_session(
     socket: axum::extract::ws::WebSocket,
     registry: Arc<OperationRegistry>,
     identity: Identity,
     policy: Arc<dyn ChannelLifecyclePolicy>,
+    sessions: Option<WsSessions>,
 ) {
-    let (byte_stream, _pumps) = split_ws_to_bytes(socket);
+    let (byte_stream, pumps) = split_ws_to_bytes(socket);
+    let pumps = Arc::new(pumps);
+
+    let _guard = sessions.map(|sessions| {
+        let id = sessions.insert(Arc::clone(&pumps));
+        SessionGuard { sessions, id }
+    });
+
     let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None);
     let _ = conn.set_identity(identity.clone());
 
@@ -48,6 +168,17 @@ pub async fn run_channels_session(
     }
 }
 
+struct SessionGuard {
+    sessions: WsSessions,
+    id: u64,
+}
+
+impl Drop for SessionGuard {
+    fn drop(&mut self) {
+        self.sessions.remove(self.id);
+    }
+}
+
 /// The `install_channel_zero` hook: split channel 0's `BiStream` into
 /// the shared writer + reader, construct the single-stream
 /// `CallConnection`, and run the shared `Dispatcher`'s single-stream
@@ -131,20 +262,31 @@ pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
 ///
 /// The channel lifecycle policy comes from the
 /// [`ChannelsPolicy`] request extension when present, else `NoCap`.
+/// The session pump handles are retained in the [`WsSessions`]
+/// registry — the shared instance from the router state
+/// (`RouterState::ws_sessions`) unless a request extension carries
+/// one, so a stuck session stays evictable (WS-08).
 pub async fn ws_upgrade_handler(
-    axum::extract::State(registry): axum::extract::State<Arc<OperationRegistry>>,
+    sessions: Option<axum::Extension<WsSessions>>,
+    axum::extract::State(state): axum::extract::State<SessionState>,
     axum::Extension(identity): axum::Extension<Identity>,
     policy: Option<axum::Extension<ChannelsPolicy>>,
     ws_upgrade: WebSocketUpgrade,
 ) -> Response {
+    let sessions = Some(
+        sessions
+            .map(|axum::Extension(s)| s)
+            .unwrap_or_else(|| WsSessions::clone(state.sessions())),
+    );
     let policy = policy
         .map(|axum::Extension(p)| p.0)
         .unwrap_or_else(|| Arc::new(NoCap));
+    let registry = Arc::clone(state.registry());
     ws_upgrade
         .max_frame_size(INBOUND_WS_FRAME_CAP)
         .max_message_size(INBOUND_WS_MESSAGE_CAP)
         .on_upgrade(move |socket| async move {
-            run_channels_session(socket, registry, identity, policy).await
+            run_channels_session(socket, registry, identity, policy, sessions).await
         })
 }