WEB: Lifecycle + Lane ownership + Diag fixes

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

src/web/manager/lifecycle.rs

@@ -2,13 +2,30 @@ use std::net::IpAddr;
 use std::sync::atomic::Ordering;
 use std::time::{Duration, Instant};
 
-use tracing::info;
+use tokio::time::Instant as TokioInstant;
+use tracing::{info, warn};
 
 use super::state::{
     decrement_map, remember_closed_token_locked, remove_bootstrap_locked, remove_expired_locked,
 };
 use super::{ProfileKey, TokenHash, WebProcessRuntime};
 
+/// Result of draining all process-owned WEB work under one absolute deadline.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum WebShutdownOutcome {
+    /// Every registered session and auxiliary task completed.
+    Drained,
+    /// Cancellation was asserted, but registered work remained at the deadline.
+    DeadlineExceeded,
+}
+
+/// Owned shutdown snapshot retained after sessions leave the live registry.
+pub(crate) struct WebShutdownDrain {
+    runtime: std::sync::Arc<WebProcessRuntime>,
+    sessions: Vec<std::sync::Arc<crate::web::session::WebSession>>,
+    started: TokioInstant,
+}
+
 impl WebProcessRuntime {
     /// Removes one closed session and retains a bounded host-bound replay marker.
     pub(crate) fn session_finished(
@@ -49,11 +66,20 @@ impl WebProcessRuntime {
         self.sessions_closed.fetch_add(1, Ordering::Relaxed);
     }
 
-    /// Stops issuance, closes all sessions, and joins bounded child work.
-    pub(crate) async fn shutdown(&self) {
+    /// Closes every WEB authority gate before any graceful wait begins.
+    pub(crate) fn begin_shutdown(self: &std::sync::Arc<Self>) -> WebShutdownDrain {
+        let started = TokioInstant::now();
         self.shutdown.cancel();
         self.close_websockets();
         self.data_budget.close();
+        self.http_connections.close();
+        self.http_handlers.close();
+        self.lane_polls.close();
+        self.lane_aux_polls.close();
+        self.body_readers.close();
+        self.body_bytes.close();
+        self.stream_handshakes.close();
+        self.websocket_connections.close();
         let sessions = {
             let mut state = self.state.lock();
             state.closed = true;
@@ -65,6 +91,24 @@ impl WebProcessRuntime {
         for session in &sessions {
             session.close();
         }
+        self.tasks.close();
+        WebShutdownDrain {
+            runtime: std::sync::Arc::clone(self),
+            sessions,
+            started,
+        }
+    }
+
+    /// Stops issuance and drains all WEB work until one absolute deadline.
+    pub(crate) async fn shutdown_until(
+        self: &std::sync::Arc<Self>,
+        deadline: TokioInstant,
+    ) -> WebShutdownOutcome {
+        self.begin_shutdown().wait_until(deadline).await
+    }
+
+    /// Stops issuance and drains WEB work under the currently configured budget.
+    pub(crate) async fn shutdown(self: &std::sync::Arc<Self>) -> WebShutdownOutcome {
         let timeout_secs = self
             .active_runtime
             .load()
@@ -72,34 +116,11 @@ impl WebProcessRuntime {
             .web
             .timeouts
             .shutdown_secs;
-        let waits = async {
-            for session in sessions {
-                session.wait().await;
-            }
-        };
-        let _ = tokio::time::timeout(Duration::from_secs(timeout_secs), waits).await;
-        self.tasks.close();
-        let _ = tokio::time::timeout(Duration::from_secs(timeout_secs), self.tasks.wait()).await;
-        let sessions_live = self.state.lock().sessions.len();
-        let streams_live = self.stream_admission.lock().streams_live;
-        let budget = self.data_budget.snapshot();
-        info!(
-            target: "telemt::web",
-            sessions_created = self.sessions_created.load(Ordering::Relaxed),
-            sessions_closed = self.sessions_closed.load(Ordering::Relaxed),
-            sessions_live,
-            streams_opened = self.streams_opened.load(Ordering::Relaxed),
-            streams_rejected = self.streams_rejected.load(Ordering::Relaxed),
-            streams_live,
-            pending_bytes = budget.queue_bytes,
-            pending_items = budget.queue_items,
-            websocket_bytes = budget.websocket_bytes,
-            data_high_water_bytes = budget.high_water_bytes,
-            bytes_up = self.bytes_up.load(Ordering::Relaxed),
-            bytes_down = self.bytes_down.load(Ordering::Relaxed),
-            limit_hits = self.limit_hits.load(Ordering::Relaxed),
-            "WEB runtime stopped"
-        );
+        let now = TokioInstant::now();
+        let deadline = now
+            .checked_add(Duration::from_secs(timeout_secs))
+            .unwrap_or(now);
+        self.shutdown_until(deadline).await
     }
 
     /// Expires credentials and closes idle sessions without holding locks across callbacks.
@@ -156,3 +177,228 @@ impl WebProcessRuntime {
         }
     }
 }
+
+impl WebShutdownDrain {
+    /// Waits for frozen session ownership and process auxiliary tasks concurrently.
+    pub(crate) async fn wait_until(self, deadline: TokioInstant) -> WebShutdownOutcome {
+        let sessions = &self.sessions;
+        let session_waits = async {
+            for session in sessions {
+                session.wait().await;
+            }
+        };
+        let outcome = wait_for_drain(deadline, session_waits, self.runtime.tasks.wait()).await;
+        self.log_outcome(outcome, deadline);
+        outcome
+    }
+
+    fn log_outcome(&self, outcome: WebShutdownOutcome, deadline: TokioInstant) {
+        let sessions_live = self.runtime.state.lock().sessions.len();
+        let streams_live = self.runtime.stream_admission.lock().streams_live;
+        let session_tasks_live = self.sessions.iter().fold(0usize, |total, session| {
+            total.saturating_add(session.tasks_live())
+        });
+        let sessions_pending = self
+            .sessions
+            .iter()
+            .filter(|session| session.tasks_live() != 0)
+            .count();
+        let auxiliary_tasks_live = self.runtime.tasks.len();
+        let budget = self.runtime.data_budget.snapshot();
+        let budget_ms = deadline
+            .saturating_duration_since(self.started)
+            .as_millis()
+            .min(u128::from(u64::MAX)) as u64;
+        let elapsed_ms = TokioInstant::now()
+            .saturating_duration_since(self.started)
+            .as_millis()
+            .min(u128::from(u64::MAX)) as u64;
+        match outcome {
+            WebShutdownOutcome::Drained => info!(
+                target: "telemt::web",
+                shutdown_drained = true,
+                shutdown_budget_ms = budget_ms,
+                shutdown_elapsed_ms = elapsed_ms,
+                sessions_created = self.runtime.sessions_created.load(Ordering::Relaxed),
+                sessions_closed = self.runtime.sessions_closed.load(Ordering::Relaxed),
+                sessions_live,
+                sessions_pending,
+                session_tasks_live,
+                auxiliary_tasks_live,
+                streams_opened = self.runtime.streams_opened.load(Ordering::Relaxed),
+                streams_rejected = self.runtime.streams_rejected.load(Ordering::Relaxed),
+                streams_live,
+                pending_bytes = budget.queue_bytes,
+                pending_items = budget.queue_items,
+                websocket_bytes = budget.websocket_bytes,
+                data_high_water_bytes = budget.high_water_bytes,
+                bytes_up = self.runtime.bytes_up.load(Ordering::Relaxed),
+                bytes_down = self.runtime.bytes_down.load(Ordering::Relaxed),
+                limit_hits = self.runtime.limit_hits.load(Ordering::Relaxed),
+                "WEB runtime stopped"
+            ),
+            WebShutdownOutcome::DeadlineExceeded => warn!(
+                target: "telemt::web",
+                shutdown_drained = false,
+                shutdown_budget_ms = budget_ms,
+                shutdown_elapsed_ms = elapsed_ms,
+                sessions_created = self.runtime.sessions_created.load(Ordering::Relaxed),
+                sessions_closed = self.runtime.sessions_closed.load(Ordering::Relaxed),
+                sessions_live,
+                sessions_pending,
+                session_tasks_live,
+                auxiliary_tasks_live,
+                streams_opened = self.runtime.streams_opened.load(Ordering::Relaxed),
+                streams_rejected = self.runtime.streams_rejected.load(Ordering::Relaxed),
+                streams_live,
+                pending_bytes = budget.queue_bytes,
+                pending_items = budget.queue_items,
+                websocket_bytes = budget.websocket_bytes,
+                data_high_water_bytes = budget.high_water_bytes,
+                bytes_up = self.runtime.bytes_up.load(Ordering::Relaxed),
+                bytes_down = self.runtime.bytes_down.load(Ordering::Relaxed),
+                limit_hits = self.runtime.limit_hits.load(Ordering::Relaxed),
+                "WEB runtime shutdown deadline exceeded"
+            ),
+        }
+    }
+}
+
+async fn wait_for_drain<S, T>(deadline: TokioInstant, sessions: S, tasks: T) -> WebShutdownOutcome
+where
+    S: std::future::Future<Output = ()>,
+    T: std::future::Future<Output = ()>,
+{
+    let waits = async {
+        tokio::join!(sessions, tasks);
+    };
+    if tokio::time::timeout_at(deadline, waits).await.is_ok() {
+        WebShutdownOutcome::Drained
+    } else {
+        WebShutdownOutcome::DeadlineExceeded
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::future::Future;
+    use std::pin::Pin;
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::task::{Context, Poll};
+
+    use arc_swap::ArcSwap;
+    use tokio::sync::Notify;
+
+    use super::*;
+    use crate::config::ProxyConfig;
+    use crate::maestro::generation::test_runtime_generation;
+
+    struct DropProbe {
+        polls: Arc<AtomicUsize>,
+        drops: Arc<AtomicUsize>,
+    }
+
+    impl Future for DropProbe {
+        type Output = ();
+
+        fn poll(self: Pin<&mut Self>, _context: &mut Context<'_>) -> Poll<Self::Output> {
+            self.polls.fetch_add(1, Ordering::AcqRel);
+            Poll::Pending
+        }
+    }
+
+    impl Drop for DropProbe {
+        fn drop(&mut self) {
+            self.drops.fetch_add(1, Ordering::AcqRel);
+        }
+    }
+
+    fn runtime() -> (
+        Arc<WebProcessRuntime>,
+        Arc<crate::maestro::generation::RuntimeGeneration>,
+    ) {
+        let generation = test_runtime_generation(1, ProxyConfig::default());
+        let runtime = WebProcessRuntime::start(Arc::new(ArcSwap::from(Arc::clone(&generation))));
+        (runtime, generation)
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn drain_uses_one_absolute_deadline_for_both_wait_groups() {
+        let started = TokioInstant::now();
+        let outcome = wait_for_drain(
+            started + Duration::from_secs(5),
+            tokio::time::sleep(Duration::from_secs(4)),
+            tokio::time::sleep(Duration::from_secs(9)),
+        )
+        .await;
+
+        assert_eq!(outcome, WebShutdownOutcome::DeadlineExceeded);
+        assert_eq!(TokioInstant::now() - started, Duration::from_secs(5));
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn drain_returns_when_both_wait_groups_finish() {
+        let started = TokioInstant::now();
+        let outcome = wait_for_drain(
+            started + Duration::from_secs(5),
+            tokio::time::sleep(Duration::from_secs(3)),
+            tokio::time::sleep(Duration::from_secs(2)),
+        )
+        .await;
+
+        assert_eq!(outcome, WebShutdownOutcome::Drained);
+        assert_eq!(TokioInstant::now() - started, Duration::from_secs(3));
+    }
+
+    #[tokio::test]
+    async fn post_shutdown_auxiliary_is_dropped_without_polling() {
+        let (runtime, generation) = runtime();
+        let polls = Arc::new(AtomicUsize::new(0));
+        let drops = Arc::new(AtomicUsize::new(0));
+        let drain = runtime.begin_shutdown();
+
+        runtime.spawn_auxiliary(DropProbe {
+            polls: Arc::clone(&polls),
+            drops: Arc::clone(&drops),
+        });
+
+        assert_eq!(polls.load(Ordering::Acquire), 0);
+        assert_eq!(drops.load(Ordering::Acquire), 1);
+        assert_eq!(
+            drain
+                .wait_until(TokioInstant::now() + Duration::from_secs(1))
+                .await,
+            WebShutdownOutcome::Drained
+        );
+        generation.stop_sessions().await;
+        generation.stop_background_tasks().await;
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn expired_deadline_still_closes_every_runtime_gate() {
+        let (runtime, generation) = runtime();
+        let existing = runtime.try_http_connection().unwrap();
+        let release = Arc::new(Notify::new());
+        let release_task = Arc::clone(&release);
+        runtime.tasks.spawn(async move {
+            release_task.notified().await;
+        });
+        tokio::task::yield_now().await;
+        let drain = runtime.begin_shutdown();
+
+        assert!(runtime.try_http_connection().is_none());
+        assert!(runtime.try_http_handler().is_none());
+        assert!(runtime.try_lane_poll(false).is_none());
+        assert_eq!(
+            drain.wait_until(TokioInstant::now()).await,
+            WebShutdownOutcome::DeadlineExceeded
+        );
+
+        drop(existing);
+        release.notify_waiters();
+        runtime.tasks.wait().await;
+        generation.stop_sessions().await;
+        generation.stop_background_tasks().await;
+    }
+}