Runtime Ownership hardened

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

src/web/manager/lifecycle.rs

@@ -448,7 +448,10 @@ mod tests {
         tokio::task::yield_now().await;
         let drain = runtime.begin_shutdown();
 
-        assert!(runtime.try_http_connection().is_none());
+        assert_eq!(
+            runtime.try_http_connection().unwrap_err(),
+            super::super::HttpConnectionAdmissionError::Closed
+        );
         assert!(runtime.try_http_handler().is_none());
         assert!(runtime.try_lane_poll(false).is_none());
         assert_eq!(

src/web/manager/operator_lifecycle.rs

@@ -1,5 +1,5 @@
 use std::sync::Arc;
-use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
 
 use arc_swap::ArcSwap;
@@ -16,12 +16,14 @@ pub(crate) use status::{
     OperatorDrainOutcome, OperatorDrainState, OperatorDrainStatus, OperatorLifecycleState,
     OperatorLifecycleStatus,
 };
+// Admission fencing and registration-drain synchronization.
+mod admission;
+use admission::{OperatorAdmission, OperatorAdmissionRejection};
+pub(super) use admission::OperatorRegistration;
 // Mutable and published lifecycle state storage.
 mod state;
 use state::{ActiveDrain, OperatorLifecycleInner, OperatorSnapshot, WorkCounts};
 
-const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
-const OPERATOR_REGISTRATION_COUNT: usize = OPERATOR_ADMISSION_CLOSED - 1;
 const DRAIN_REF_VERSION: &str = "wd1";
 
 /// Stable operator-control rejection category.
@@ -33,77 +35,7 @@ pub(crate) enum OperatorLifecycleError {
     OperationInProgress,
 }
 
-struct OperatorAdmission {
-    state: AtomicUsize,
-    registrations_drained: Notify,
-}
-
-pub(super) struct OperatorRegistration<'a> {
-    admission: &'a OperatorAdmission,
-}
-
-impl OperatorAdmission {
-    fn new() -> Self {
-        Self {
-            state: AtomicUsize::new(0),
-            registrations_drained: Notify::new(),
-        }
-    }
-
-    fn try_register(&self) -> Option<OperatorRegistration<'_>> {
-        let mut state = self.state.load(Ordering::Acquire);
-        loop {
-            if state & OPERATOR_ADMISSION_CLOSED != 0
-                || state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT
-            {
-                return None;
-            }
-            match self.state.compare_exchange_weak(
-                state,
-                state + 1,
-                Ordering::AcqRel,
-                Ordering::Acquire,
-            ) {
-                Ok(_) => return Some(OperatorRegistration { admission: self }),
-                Err(observed) => state = observed,
-            }
-        }
-    }
-
-    fn close(&self) {
-        self.state
-            .fetch_or(OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel);
-    }
-
-    fn reopen(&self) {
-        self.state
-            .fetch_and(!OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel);
-    }
-
-    fn is_closed(&self) -> bool {
-        self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0
-    }
-
-    async fn wait_for_registrations(&self) {
-        loop {
-            let notified = self.registrations_drained.notified();
-            if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 {
-                return;
-            }
-            notified.await;
-        }
-    }
-}
-
-impl Drop for OperatorRegistration<'_> {
-    fn drop(&mut self) {
-        let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
-        if previous & OPERATOR_REGISTRATION_COUNT == 1 {
-            self.admission.registrations_drained.notify_waiters();
-        }
-    }
-}
-
+/// Process-owned reversible lifecycle and admission authority.
 pub(super) struct OperatorLifecycle {
     runtime_instance: Arc<str>,
     admission: OperatorAdmission,
@@ -115,6 +47,7 @@ pub(super) struct OperatorLifecycle {
 }
 
 impl OperatorLifecycle {
+    /// Creates a running lifecycle for one immutable process instance.
     pub(super) fn new(runtime_instance: Arc<str>) -> Self {
         let since = Instant::now();
         let snapshot = OperatorSnapshot {
@@ -142,16 +75,21 @@ impl OperatorLifecycle {
         }
     }
 
-    pub(super) fn try_register(&self) -> Option<OperatorRegistration<'_>> {
+    /// Registers one synchronous operator-fenced admission section.
+    pub(super) fn try_register(
+        &self,
+    ) -> Result<OperatorRegistration<'_>, crate::web::telemetry::WebRejectionReason> {
         self.admission.try_register()
     }
 
+    /// Wakes an active drain after tracked work ownership changes.
     pub(super) fn notify_work_changed(&self) {
         if self.admission.is_closed() {
             self.work_changed.notify_waiters();
         }
     }
 
+    /// Returns the lock-free lifecycle snapshot with effective config admission.
     pub(super) fn status(&self, config_enabled: bool) -> OperatorLifecycleStatus {
         let snapshot = self.published.load();
         let admission_open =
@@ -170,30 +108,6 @@ impl OperatorLifecycle {
         self.published.load().terminal
     }
 
-    fn rejection_reason(&self) -> crate::web::telemetry::WebRejectionReason {
-        let inner = self.inner.lock();
-        if inner.terminal {
-            return crate::web::telemetry::WebRejectionReason::RuntimeClosed;
-        }
-        match inner.state {
-            OperatorLifecycleState::Paused => {
-                crate::web::telemetry::WebRejectionReason::OperatorPaused
-            }
-            OperatorLifecycleState::Draining => {
-                crate::web::telemetry::WebRejectionReason::OperatorDraining
-            }
-            OperatorLifecycleState::ForceClosing => {
-                crate::web::telemetry::WebRejectionReason::OperatorForceClosing
-            }
-            OperatorLifecycleState::Drained => {
-                crate::web::telemetry::WebRejectionReason::OperatorDrained
-            }
-            OperatorLifecycleState::Running => {
-                crate::web::telemetry::WebRejectionReason::RuntimeClosed
-            }
-        }
-    }
-
     fn publish_locked(&self, inner: &OperatorLifecycleInner) {
         self.published.store(Arc::new(OperatorSnapshot {
             state: inner.state,
@@ -246,6 +160,8 @@ impl OperatorLifecycle {
             return false;
         }
         self.transition_locked(&mut inner, OperatorLifecycleState::ForceClosing);
+        self.admission
+            .close(OperatorAdmissionRejection::ForceClosing);
         let Some(drain) = inner.drain.as_mut() else {
             return false;
         };
@@ -269,6 +185,7 @@ impl OperatorLifecycle {
         }
         inner.active = None;
         self.transition_locked(&mut inner, OperatorLifecycleState::Drained);
+        self.admission.close(OperatorAdmissionRejection::Drained);
         if let Some(drain) = inner.drain.as_mut() {
             drain.state = OperatorDrainState::Completed;
             drain.outcome = Some(if forced {
@@ -286,7 +203,8 @@ impl OperatorLifecycle {
 
     fn close_terminal(&self) {
         let mut inner = self.inner.lock();
-        self.admission.close();
+        self.admission
+            .close(OperatorAdmissionRejection::RuntimeClosed);
         if inner.terminal {
             return;
         }
@@ -321,7 +239,6 @@ impl WebProcessRuntime {
             if inner.terminal || self.shutdown.is_cancelled() {
                 return Err(OperatorLifecycleError::Closed);
             }
-            self.operator_lifecycle.admission.close();
             if matches!(
                 inner.state,
                 OperatorLifecycleState::Running | OperatorLifecycleState::Drained
@@ -329,6 +246,9 @@ impl WebProcessRuntime {
                 self.operator_lifecycle
                     .transition_locked(&mut inner, OperatorLifecycleState::Paused);
             }
+            self.operator_lifecycle.admission.close(
+                OperatorAdmissionRejection::for_state(inner.state),
+            );
             self.operator_lifecycle.publish_locked(&inner);
         }
         self.operator_lifecycle
@@ -355,7 +275,7 @@ impl WebProcessRuntime {
         let started = TokioInstant::now();
         let deadline = started + timeout;
         let started_epoch_millis = crate::web::trace::store_epoch_millis();
-        {
+        let accepted = {
             let mut inner = self.operator_lifecycle.inner.lock();
             if inner.terminal || self.shutdown.is_cancelled() {
                 return Err(OperatorLifecycleError::Closed);
@@ -368,7 +288,6 @@ impl WebProcessRuntime {
             {
                 return Err(OperatorLifecycleError::OperationInProgress);
             }
-            self.operator_lifecycle.admission.close();
             inner.active = Some(ActiveDrain {
                 sequence,
                 cancellation: cancellation.clone(),
@@ -392,26 +311,20 @@ impl WebProcessRuntime {
             });
             self.operator_lifecycle
                 .transition_locked(&mut inner, OperatorLifecycleState::Draining);
+            self.operator_lifecycle
+                .admission
+                .close(OperatorAdmissionRejection::Draining);
+            let runtime = Arc::clone(self);
+            self.spawn_auxiliary(async move {
+                runtime
+                    .run_operator_drain(sequence, deadline, cancellation)
+                    .await;
+            });
             self.operator_lifecycle.publish_locked(&inner);
-        }
-        self.operator_lifecycle
-            .admission
-            .wait_for_registrations()
-            .await;
-        if self.shutdown.is_cancelled() || self.operator_lifecycle.is_terminal() {
-            return Err(OperatorLifecycleError::Closed);
-        }
-        let counts = self.operator_work_counts();
-        if !self.operator_lifecycle.update_counts(sequence, counts) {
-            return Err(OperatorLifecycleError::Closed);
-        }
-        let runtime = Arc::clone(self);
-        self.spawn_auxiliary(async move {
-            runtime
-                .run_operator_drain(sequence, deadline, cancellation)
-                .await;
-        });
-        Ok(self.operator_lifecycle_status())
+            let config_enabled = self.active_generation().config().web.enabled;
+            self.operator_lifecycle.status(config_enabled)
+        };
+        Ok(accepted)
     }
 
     /// Resumes operator admission and invalidates any active drain waiter.
@@ -444,23 +357,25 @@ impl WebProcessRuntime {
         Ok(self.operator_lifecycle_status())
     }
 
+    /// Registers one WEB manager admission commit under the operator fence.
     pub(super) fn try_operator_admission(
         &self,
     ) -> Result<OperatorRegistration<'_>, super::ManagerError> {
         match self.operator_lifecycle.try_register() {
-            Some(registration) => Ok(registration),
-            None => {
-                self.telemetry
-                    .record_rejection(self.operator_lifecycle.rejection_reason());
+            Ok(registration) => Ok(registration),
+            Err(reason) => {
+                self.telemetry.record_rejection(reason);
                 Err(super::ManagerError::AdmissionPaused)
             }
         }
     }
 
+    /// Notifies an active drain that tracked WEB work changed.
     pub(super) fn notify_operator_work_changed(&self) {
         self.operator_lifecycle.notify_work_changed();
     }
 
+    /// Terminally closes operator lifecycle during process shutdown.
     pub(super) fn close_operator_lifecycle(&self) {
         self.operator_lifecycle.close_terminal();
     }
@@ -482,9 +397,16 @@ impl WebProcessRuntime {
         deadline: TokioInstant,
         cancellation: CancellationToken,
     ) {
+        tokio::select! {
+            biased;
+            _ = cancellation.cancelled() => return,
+            _ = self.operator_lifecycle.admission.wait_for_registrations() => {}
+        }
         let mut forced = false;
         loop {
             let notified = self.operator_lifecycle.work_changed.notified();
+            tokio::pin!(notified);
+            notified.as_mut().enable();
             let counts = self.operator_work_counts();
             if !self.operator_lifecycle.update_counts(sequence, counts) {
                 return;
@@ -497,14 +419,14 @@ impl WebProcessRuntime {
                 tokio::select! {
                     biased;
                     _ = cancellation.cancelled() => return,
-                    _ = notified => {}
+                    _ = notified.as_mut() => {}
                 }
                 continue;
             }
             tokio::select! {
                 biased;
                 _ = cancellation.cancelled() => return,
-                _ = notified => {},
+                _ = notified.as_mut() => {},
                 _ = tokio::time::sleep_until(deadline) => {
                     let counts = self.operator_work_counts();
                     if counts.is_zero() {

src/web/manager/operator_lifecycle/admission.rs

@@ -0,0 +1,165 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use tokio::sync::Notify;
+
+use super::OperatorLifecycleState;
+
+const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
+const OPERATOR_REJECTION_BITS: u32 = 3;
+const OPERATOR_REJECTION_SHIFT: u32 = usize::BITS - 1 - OPERATOR_REJECTION_BITS;
+const OPERATOR_REJECTION_MASK: usize =
+    ((1usize << OPERATOR_REJECTION_BITS) - 1) << OPERATOR_REJECTION_SHIFT;
+const OPERATOR_REGISTRATION_COUNT: usize = (1usize << OPERATOR_REJECTION_SHIFT) - 1;
+
+#[derive(Clone, Copy)]
+#[repr(usize)]
+/// Stable rejection encoded into the closed admission-fence word.
+pub(super) enum OperatorAdmissionRejection {
+    /// Operator pause closed admission.
+    Paused = 1,
+    /// Graceful drain closed admission.
+    Draining = 2,
+    /// Deadline-triggered forced closure remains in progress.
+    ForceClosing = 3,
+    /// The latest drain reached confirmed zero.
+    Drained = 4,
+    /// Terminal process shutdown closed admission.
+    RuntimeClosed = 5,
+}
+
+impl OperatorAdmissionRejection {
+    /// Maps a closed lifecycle state to its admission rejection.
+    pub(super) fn for_state(state: OperatorLifecycleState) -> Self {
+        match state {
+            OperatorLifecycleState::Paused => Self::Paused,
+            OperatorLifecycleState::Draining => Self::Draining,
+            OperatorLifecycleState::ForceClosing => Self::ForceClosing,
+            OperatorLifecycleState::Drained => Self::Drained,
+            OperatorLifecycleState::Running => Self::RuntimeClosed,
+        }
+    }
+
+    fn from_admission_state(state: usize) -> Self {
+        match (state & OPERATOR_REJECTION_MASK) >> OPERATOR_REJECTION_SHIFT {
+            value if value == Self::Paused as usize => Self::Paused,
+            value if value == Self::Draining as usize => Self::Draining,
+            value if value == Self::ForceClosing as usize => Self::ForceClosing,
+            value if value == Self::Drained as usize => Self::Drained,
+            _ => Self::RuntimeClosed,
+        }
+    }
+
+    fn telemetry_reason(self) -> crate::web::telemetry::WebRejectionReason {
+        match self {
+            Self::Paused => crate::web::telemetry::WebRejectionReason::OperatorPaused,
+            Self::Draining => crate::web::telemetry::WebRejectionReason::OperatorDraining,
+            Self::ForceClosing => {
+                crate::web::telemetry::WebRejectionReason::OperatorForceClosing
+            }
+            Self::Drained => crate::web::telemetry::WebRejectionReason::OperatorDrained,
+            Self::RuntimeClosed => crate::web::telemetry::WebRejectionReason::RuntimeClosed,
+        }
+    }
+}
+
+/// Lock-free admission fence with bounded pre-cutover registration tracking.
+pub(super) struct OperatorAdmission {
+    state: AtomicUsize,
+    registrations_drained: Notify,
+}
+
+/// RAII ownership of one synchronous pre-cutover admission section.
+pub(in crate::web::manager) struct OperatorRegistration<'a> {
+    admission: &'a OperatorAdmission,
+}
+
+impl OperatorAdmission {
+    /// Creates an open admission fence.
+    pub(super) fn new() -> Self {
+        Self {
+            state: AtomicUsize::new(0),
+            registrations_drained: Notify::new(),
+        }
+    }
+
+    /// Registers one pre-cutover section or returns its stable rejection reason.
+    pub(super) fn try_register(
+        &self,
+    ) -> Result<OperatorRegistration<'_>, crate::web::telemetry::WebRejectionReason> {
+        let mut state = self.state.load(Ordering::Acquire);
+        loop {
+            if state & OPERATOR_ADMISSION_CLOSED != 0 {
+                return Err(
+                    OperatorAdmissionRejection::from_admission_state(state).telemetry_reason(),
+                );
+            }
+            if state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT {
+                return Err(crate::web::telemetry::WebRejectionReason::Concurrent);
+            }
+            match self.state.compare_exchange_weak(
+                state,
+                state + 1,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            ) {
+                Ok(_) => return Ok(OperatorRegistration { admission: self }),
+                Err(observed) => state = observed,
+            }
+        }
+    }
+
+    /// Atomically closes admission while preserving active registration count.
+    pub(super) fn close(&self, reason: OperatorAdmissionRejection) {
+        let reason = (reason as usize) << OPERATOR_REJECTION_SHIFT;
+        let mut state = self.state.load(Ordering::Acquire);
+        loop {
+            let next = (state & OPERATOR_REGISTRATION_COUNT)
+                | OPERATOR_ADMISSION_CLOSED
+                | reason;
+            match self.state.compare_exchange_weak(
+                state,
+                next,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            ) {
+                Ok(_) => return,
+                Err(observed) => state = observed,
+            }
+        }
+    }
+
+    /// Reopens admission without altering active registration ownership.
+    pub(super) fn reopen(&self) {
+        self.state.fetch_and(
+            !(OPERATOR_ADMISSION_CLOSED | OPERATOR_REJECTION_MASK),
+            Ordering::AcqRel,
+        );
+    }
+
+    /// Returns whether the admission fence is closed.
+    pub(super) fn is_closed(&self) -> bool {
+        self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0
+    }
+
+    /// Waits until every pre-cutover registration has been released.
+    pub(super) async fn wait_for_registrations(&self) {
+        loop {
+            let notified = self.registrations_drained.notified();
+            tokio::pin!(notified);
+            notified.as_mut().enable();
+            if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 {
+                return;
+            }
+            notified.await;
+        }
+    }
+}
+
+impl Drop for OperatorRegistration<'_> {
+    fn drop(&mut self) {
+        let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
+        if previous & OPERATOR_REGISTRATION_COUNT == 1 {
+            self.admission.registrations_drained.notify_waiters();
+        }
+    }
+}

src/web/manager/operator_lifecycle/tests.rs

@@ -133,6 +133,43 @@ async fn empty_drain_completes_gracefully_and_stays_closed_until_resume() {
     stop_runtime(runtime, generation).await;
 }
 
+#[tokio::test]
+async fn drain_request_returns_after_registering_its_worker() {
+    let (runtime, generation) = test_runtime();
+    let registration = runtime.try_operator_admission().unwrap();
+    let drain_runtime = Arc::clone(&runtime);
+    let drain = tokio::spawn(async move {
+        drain_runtime
+            .drain_operator(Duration::from_secs(30))
+            .await
+    });
+
+    tokio::time::timeout(Duration::from_secs(1), async {
+        while runtime.operator_lifecycle_status().state != OperatorLifecycleState::Draining {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .unwrap();
+    let accepted = tokio::time::timeout(Duration::from_secs(1), drain)
+        .await
+        .unwrap()
+        .unwrap()
+        .unwrap();
+    assert_eq!(accepted.state, OperatorLifecycleState::Draining);
+    drop(registration);
+
+    tokio::time::timeout(Duration::from_secs(1), async {
+        while runtime.operator_lifecycle_status().state != OperatorLifecycleState::Drained {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .unwrap();
+
+    stop_runtime(runtime, generation).await;
+}
+
 #[tokio::test]
 async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() {
     let (runtime, generation) = test_runtime();
@@ -140,7 +177,10 @@ async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() {
     let cancellation = CancellationToken::new();
     {
         let mut inner = runtime.operator_lifecycle.inner.lock();
-        runtime.operator_lifecycle.admission.close();
+        runtime
+            .operator_lifecycle
+            .admission
+            .close(OperatorAdmissionRejection::Draining);
         inner.active = Some(ActiveDrain {
             sequence,
             cancellation,