diff --git a/pkg/procmgr/rust/src/manager/mod.rs b/pkg/procmgr/rust/src/manager/mod.rs index 1de3c622e9d1..ac64418a5976 100644 --- a/pkg/procmgr/rust/src/manager/mod.rs +++ b/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; diff --git a/pkg/procmgr/rust/src/manager/process_manager.rs b/pkg/procmgr/rust/src/manager/process_manager.rs index ccc53485e289..2e133926dfae 100644 --- a/pkg/procmgr/rust/src/manager/process_manager.rs +++ b/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(); diff --git a/pkg/procmgr/rust/src/manager/runtime.rs b/pkg/procmgr/rust/src/manager/runtime.rs index 8e2afe102b69..56a56fbe3453 100644 --- a/pkg/procmgr/rust/src/manager/runtime.rs +++ b/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(); diff --git a/pkg/procmgr/rust/src/platform/unix/mod.rs b/pkg/procmgr/rust/src/platform/unix/mod.rs index 93f3646ee5fe..6a6891a17a03 100644 --- a/pkg/procmgr/rust/src/platform/unix/mod.rs +++ b/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)] diff --git a/pkg/procmgr/rust/src/platform/windows/mod.rs b/pkg/procmgr/rust/src/platform/windows/mod.rs index 27ef67260c2f..69978230a0aa 100644 --- a/pkg/procmgr/rust/src/platform/windows/mod.rs +++ b/pkg/procmgr/rust/src/platform/windows/mod.rs @@ -406,13 +406,19 @@ fn default_stable_fleet_policies_dir() -> Option { } 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)] diff --git a/pkg/procmgr/rust/src/platform/windows/scm_service.rs b/pkg/procmgr/rust/src/platform/windows/scm_service.rs index bbe189f90352..61ba7b470d7d 100644 --- a/pkg/procmgr/rust/src/platform/windows/scm_service.rs +++ b/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(); diff --git a/pkg/procmgr/rust/src/process.rs b/pkg/procmgr/rust/src/process.rs index bfae0bfc21fb..2b30698a7dd6 100644 --- a/pkg/procmgr/rust/src/process.rs +++ b/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( ctx: &mut C, - budget: ShutdownBudget, + mut budget: ShutdownBudget, ) -> Option { 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 { diff --git a/pkg/procmgr/rust/src/shutdown.rs b/pkg/procmgr/rust/src/shutdown.rs index 598476c5be9a..82aef4c82247 100644 --- a/pkg/procmgr/rust/src/shutdown.rs +++ b/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 { + 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>( + graceful_budget: Duration, + fut: F, +) -> Result { + #[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::()).await; + } + tokio::select! { + biased; + _ = notified => { + tokio::time::timeout(Duration::ZERO, std::future::pending::()).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);