[procmgr] Anchor in-flight Stop RPCs to SCM shutdown budget

pkg/procmgr/rust/src/manager/mod.rs

@@ -19,6 +19,8 @@ mod tracked_join;
 #[cfg(test)]
 mod tests;
 
+#[cfg(windows)]
+pub(crate) use deferred_cleanup::join_deferred_spawn_tasks;
 pub(crate) use lifecycle::Lifecycle;
 pub(crate) use runtime::RuntimeContext;
 

pkg/procmgr/rust/src/manager/process_manager.rs

@@ -253,10 +253,7 @@ impl ProcessManager {
         };
 
         self.catalog
-            .wait_for_process_stop(
-                idx,
-                crate::shutdown::ShutdownBudget::unlimited(std::time::Instant::now()),
-            )
+            .wait_for_process_stop(idx, crate::shutdown::ShutdownBudget::for_single_stop())
             .await;
 
         let state = self.catalog.read_processes().await[idx].state();

pkg/procmgr/rust/src/manager/runtime.rs

@@ -746,6 +746,21 @@ mod tests {
         platform::reset_shutdown_state_for_test();
     }
 
+    #[tokio::test]
+    async fn shutdown_signal_completes_when_requested_before_first_poll() {
+        let _guard = platform::test_shutdown_lock().await;
+        platform::reset_shutdown_state_for_test();
+        platform::signal_shutdown_for_test();
+
+        let shutdown = platform::shutdown_signal();
+        tokio::pin!(shutdown);
+        tokio::time::timeout(Duration::from_millis(50), shutdown.as_mut())
+            .await
+            .expect("shutdown should complete when the flag was set before the first poll");
+
+        platform::reset_shutdown_state_for_test();
+    }
+
     #[tokio::test]
     async fn drain_exits_during_work_drains_beyond_channel_capacity() {
         let manager = empty_manager();

pkg/procmgr/rust/src/platform/unix/mod.rs

@@ -110,6 +110,12 @@ pub fn stderr_inheritable() -> bool {
 }
 
 pub(crate) async fn wait_for_shutdown() {
+    let notified = shutdown_notify().notified();
+    tokio::pin!(notified);
+    notified.as_mut().enable();
+    if shutdown_requested() {
+        return;
+    }
     use tokio::signal::unix::{SignalKind, signal};
     let mut sigterm = signal(SignalKind::terminate()).expect("failed to register SIGTERM handler");
     let mut sigint = signal(SignalKind::interrupt()).expect("failed to register SIGINT handler");
@@ -122,7 +128,7 @@ pub(crate) async fn wait_for_shutdown() {
             mark_shutdown_requested();
             log::info!("received SIGINT");
         }
-        _ = shutdown_notify().notified() => {
+        _ = notified => {
             mark_shutdown_requested();
         }
     }
@@ -135,7 +141,7 @@ pub async fn shutdown_signal() {
 #[cfg(test)]
 pub(crate) fn signal_shutdown_for_test() {
     mark_shutdown_requested();
-    shutdown_notify().notify_one();
+    shutdown_notify().notify_waiters();
 }
 
 #[cfg(test)]

pkg/procmgr/rust/src/platform/windows/mod.rs

@@ -406,13 +406,19 @@ fn default_stable_fleet_policies_dir() -> Option<PathBuf> {
 }
 
 pub(crate) async fn wait_for_shutdown() {
+    let notified = shutdown_notify().notified();
+    tokio::pin!(notified);
+    notified.as_mut().enable();
+    if shutdown_requested() {
+        return;
+    }
     tokio::select! {
         result = tokio::signal::ctrl_c() => {
             result.expect("failed to register Ctrl+C handler");
             SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
             log::info!("received Ctrl+C");
         }
-        _ = shutdown_notify().notified() => {
+        _ = notified => {
             SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
             log::info!("received service stop request");
         }
@@ -426,7 +432,7 @@ pub async fn shutdown_signal() {
 #[cfg(test)]
 pub(crate) fn signal_shutdown_for_test() {
     SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
-    shutdown_notify().notify_one();
+    shutdown_notify().notify_waiters();
 }
 
 #[cfg(test)]

pkg/procmgr/rust/src/platform/windows/scm_service.rs

@@ -24,7 +24,7 @@ use windows_sys::Win32::System::Services::{
 
 use crate::config::YamlConfigLoader;
 use crate::manager::ProcessManager;
-use crate::manager::deferred_cleanup::join_deferred_spawn_tasks;
+use crate::manager::join_deferred_spawn_tasks;
 use crate::shutdown::ShutdownBudget;
 use crate::uuid_gen::V4UuidGenerator;
 
@@ -82,7 +82,7 @@ unsafe extern "system" fn ctrl_handler(
                 SCM_STOP_WAIT_HINT.as_millis() as u32,
             );
             super::record_service_stop_signal();
-            super::shutdown_notify().notify_one();
+            super::shutdown_notify().notify_waiters();
             NO_ERROR
         }
         SERVICE_CONTROL_INTERROGATE => NO_ERROR,
@@ -162,12 +162,34 @@ fn run_service_inner() -> Result<()> {
         .unwrap_or_else(|| ShutdownBudget::unlimited(Instant::now()));
     runtime.block_on(join_deferred_spawn_tasks(shutdown_budget));
 
-    let runtime_shutdown_cap = shutdown_budget.remaining_cap(Duration::from_secs(30));
-    runtime.shutdown_timeout(runtime_shutdown_cap);
-
+    shutdown_runtime_after_supervisor(runtime);
     result
 }
 
+/// Tear down the service Tokio runtime without blocking past the SCM stop budget.
+///
+/// After `block_on` returns, outstanding `spawn_blocking` work (for example a
+/// stuck Windows auto-start spawn) would otherwise keep `Runtime` drop waiting
+/// indefinitely and delay reporting `SERVICE_STOPPED`.
+fn shutdown_runtime_after_supervisor(runtime: tokio::runtime::Runtime) {
+    let signal_time = super::service_stop_signal_time();
+    let timeout = signal_time
+        .map(|signal| service_shutdown_deadline(signal).saturating_duration_since(Instant::now()))
+        .unwrap_or(Duration::from_secs(60));
+
+    if signal_time.is_some() {
+        if timeout.is_zero() {
+            warn!(
+                "SCM shutdown budget exhausted before runtime teardown; not waiting for outstanding tasks"
+            );
+        } else {
+            info!("waiting up to {timeout:?} for runtime teardown (SCM budget)");
+        }
+    }
+
+    runtime.shutdown_timeout(timeout);
+}
+
 pub fn run_as_service() -> Result<()> {
     let name = service_name_wide();
 

pkg/procmgr/rust/src/process.rs

@@ -9,7 +9,7 @@ use crate::handle::ProcessHandle;
 #[cfg(windows)]
 use crate::handle::ProcessWaitControl;
 use crate::platform;
-use crate::shutdown::ShutdownBudget;
+use crate::shutdown::{self, ShutdownBudget};
 use crate::spawn::{SpawnProfile, profile_for};
 use crate::state::ProcessState;
 use anyhow::{Context, Result, bail};
@@ -236,7 +236,7 @@ impl StopWaitPlan {
                 mut handle,
                 graceful_budget,
                 ..
-            } => match time::timeout(graceful_budget, &mut handle).await {
+            } => match shutdown::wait_graceful_or_shutdown(graceful_budget, &mut handle).await {
                 Ok(Ok(status)) => StopWaitResult::Exited(status),
                 Ok(Err(error)) => {
                     warn!("[{name}] watcher join failed during stop: {error:#}");
@@ -249,7 +249,7 @@ impl StopWaitPlan {
                 mut handle,
                 graceful_budget,
                 ..
-            } => match time::timeout(graceful_budget, handle.wait()).await {
+            } => match shutdown::wait_graceful_or_shutdown(graceful_budget, handle.wait()).await {
                 Ok(Ok(status)) => StopWaitResult::Exited(Some(status)),
                 Ok(Err(error)) => {
                     warn!("[{name}] wait failed during stop: {error:#}");
@@ -953,6 +953,7 @@ impl ManagedProcess {
             return false;
         }
 
+        let budget = budget.refresh();
         let graceful_budget = budget.graceful_budget(self.stop_timeout());
         match result {
             StopWaitResult::Exited(status) => self.apply_stop_exit(status),
@@ -999,7 +1000,7 @@ impl ManagedProcess {
             self.mark_stopped();
         }
         self.stop_wait_generation.clear();
-        self.release_stop_wait_resources(budget).await;
+        self.release_stop_wait_resources(budget.refresh()).await;
     }
 
     pub(crate) fn has_orphaned_stop_wait(&self) -> bool {
@@ -1133,7 +1134,7 @@ impl ManagedProcess {
     }
 
     pub async fn wait_for_stop(&mut self) {
-        self.wait_for_stop_since(ShutdownBudget::unlimited(Instant::now().into()))
+        self.wait_for_stop_since(ShutdownBudget::unlimited(std::time::Instant::now()))
             .await;
     }
 
@@ -1267,12 +1268,13 @@ pub(crate) trait StopWaitContext {
 
 pub(crate) async fn run_stop_wait<C: StopWaitContext>(
     ctx: &mut C,
-    budget: ShutdownBudget,
+    mut budget: ShutdownBudget,
 ) -> Option<StopWaitOwner> {
     let mut owner_needing_complete = None;
     while let Some(plan) = ctx.plan_stop(budget).await {
         let owner = plan.owner();
         let result = plan.execute().await;
+        budget = budget.refresh();
         if ctx.finalize_stop(owner, result, budget).await {
             owner_needing_complete = Some(owner);
         } else {

pkg/procmgr/rust/src/shutdown.rs

@@ -26,7 +26,24 @@ impl ShutdownBudget {
         }
     }
 
-    /// Budget for ordered shutdown after a service stop signal.
+    pub(crate) fn prefer_service_stop(fallback: Self) -> Self {
+        #[cfg(windows)]
+        {
+            if let Some(signal_time) = crate::platform::service_stop_signal_time() {
+                return Self::service_stop(signal_time);
+            }
+        }
+        fallback
+    }
+
+    pub(crate) fn refresh(self) -> Self {
+        Self::prefer_service_stop(self)
+    }
+
+    pub(crate) fn for_single_stop() -> Self {
+        Self::prefer_service_stop(Self::unlimited(Instant::now()))
+    }
+
     pub(crate) fn service_stop(signal_time: Instant) -> Self {
         #[cfg(windows)]
         {
@@ -54,7 +71,6 @@ impl ShutdownBudget {
         self.remaining_cap(from_stop_timeout)
     }
 
-    /// Remaining time for a phase, capped by `cap` and the service deadline.
     pub(crate) fn remaining_cap(&self, cap: Duration) -> Duration {
         self.deadline
             .map(|deadline| cap.min(deadline.saturating_duration_since(Instant::now())))
@@ -64,6 +80,44 @@ impl ShutdownBudget {
     pub(crate) fn is_bounded(&self) -> bool {
         self.deadline.is_some()
     }
+
+    /// Remaining SCM shutdown budget for a caller-defined cap, if a service stop is in progress.
+    #[cfg(windows)]
+    pub(crate) fn remaining_service_stop_cap(cap: Duration) -> Option<Duration> {
+        crate::platform::service_stop_signal_time()
+            .map(|signal_time| Self::service_stop(signal_time).remaining_cap(cap))
+    }
+}
+
+pub(crate) async fn wait_graceful_or_shutdown<T, F: std::future::Future<Output = T>>(
+    graceful_budget: Duration,
+    fut: F,
+) -> Result<T, tokio::time::error::Elapsed> {
+    #[cfg(windows)]
+    {
+        let notified = crate::platform::shutdown_notify().notified();
+        tokio::pin!(notified);
+        notified.as_mut().enable();
+        let budget = if crate::platform::shutdown_requested() {
+            ShutdownBudget::remaining_service_stop_cap(graceful_budget).unwrap_or(graceful_budget)
+        } else {
+            graceful_budget
+        };
+        if budget.is_zero() {
+            return tokio::time::timeout(Duration::ZERO, std::future::pending::<T>()).await;
+        }
+        tokio::select! {
+            biased;
+            _ = notified => {
+                tokio::time::timeout(Duration::ZERO, std::future::pending::<T>()).await
+            }
+            result = tokio::time::timeout(budget, fut) => result,
+        }
+    }
+    #[cfg(not(windows))]
+    {
+        tokio::time::timeout(graceful_budget, fut).await
+    }
 }
 
 pub async fn shutdown_ordered(processes: &mut [ManagedProcess], order: &[usize]) {
@@ -95,6 +149,38 @@ mod tests {
         test_helpers::make_config(cmd, args)
     }
 
+    #[tokio::test]
+    async fn test_already_requested_shutdown_still_waits_graceful_budget() {
+        let _guard = crate::platform::test_shutdown_lock().await;
+        crate::platform::reset_shutdown_state_for_test();
+        crate::platform::signal_shutdown_for_test();
+
+        let started = Instant::now();
+        let result =
+            wait_graceful_or_shutdown(Duration::from_millis(150), std::future::pending::<()>())
+                .await;
+        crate::platform::reset_shutdown_state_for_test();
+
+        assert!(result.is_err(), "pending future should time out");
+        let elapsed = started.elapsed();
+        assert!(
+            elapsed >= Duration::from_millis(100),
+            "already-requested shutdown must not skip the graceful budget, elapsed {elapsed:?}"
+        );
+    }
+
+    #[test]
+    fn test_prefer_service_stop_without_signal_preserves_fallback() {
+        let signal_time = Instant::now();
+        let fallback = ShutdownBudget::unlimited(signal_time);
+        let budget = ShutdownBudget::prefer_service_stop(fallback);
+        let graceful = budget.graceful_budget(Duration::from_secs(90));
+        assert!(
+            graceful >= Duration::from_secs(89) && graceful <= Duration::from_secs(90),
+            "expected ~90s graceful budget without SCM signal, got {graceful:?}"
+        );
+    }
+
     #[tokio::test]
     async fn test_shutdown_all_graceful() {
         let cfg1 = sleep_config();
@@ -148,7 +234,6 @@ mod tests {
         p3.spawn().unwrap();
 
         let mut procs = vec![p1, p2, p3];
-        // Reverse order: p3, p2, p1
         shutdown_ordered(&mut procs, &[2, 1, 0]).await;
 
         assert_eq!(procs[0].state(), ProcessState::Stopped);