Skip to main content

fastapi_core/
shutdown.rs

1//! Graceful shutdown coordination for the server.
2//!
3//! This module provides graceful shutdown using asupersync's structured concurrency:
4//!
5//! # Shutdown Phases
6//!
7//! 1. **Stop accepting**: Server stops accepting new connections
8//! 2. **Shutdown flag**: New requests receive 503 Service Unavailable
9//! 3. **Grace period**: In-flight requests get remaining time to complete
10//! 4. **Cancellation**: Remaining requests cancelled after grace period
11//! 5. **Shutdown hooks**: Registered cleanup callbacks run
12//! 6. **Region close**: Server region fully closed
13//!
14//! # Signal Handling
15//!
16//! - SIGTERM/SIGINT triggers graceful shutdown
17//! - Second signal forces immediate exit
18//! - Custom signals supported via configuration
19//!
20//! # Example
21//!
22//! ```ignore
23//! use fastapi_core::shutdown::{ShutdownController, GracefulShutdown};
24//! use std::time::Duration;
25//!
26//! async fn run_server() {
27//!     let controller = ShutdownController::new();
28//!
29//!     // Register with signals
30//!     controller.listen_for_signals();
31//!
32//!     // Run with graceful shutdown
33//!     let shutdown = GracefulShutdown::new(controller.subscribe())
34//!         .grace_period(Duration::from_secs(30));
35//!
36//!     let result = shutdown.run(server_main()).await;
37//!     match result {
38//!         ShutdownOutcome::Completed(r) => println!("Server completed: {:?}", r),
39//!         ShutdownOutcome::GracefulShutdown => println!("Graceful shutdown"),
40//!         ShutdownOutcome::ForcedShutdown => println!("Forced shutdown"),
41//!     }
42//! }
43//! ```
44
45use asupersync::{Budget, CancelReason};
46use std::future::Future;
47use std::pin::Pin;
48use std::sync::Arc;
49use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
50use std::task::{Context, Poll, Waker};
51use std::time::Duration;
52
53// ============================================================================
54// Shutdown Phase State Machine
55// ============================================================================
56
57/// Current phase of the shutdown process.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59#[repr(u8)]
60pub enum ShutdownPhase {
61    /// Server is running normally.
62    Running = 0,
63    /// Stop accepting new connections.
64    StopAccepting = 1,
65    /// Shutdown flag set; new requests get 503.
66    ShutdownFlagged = 2,
67    /// Grace period for in-flight requests.
68    GracePeriod = 3,
69    /// Cancelling remaining requests.
70    Cancelling = 4,
71    /// Running shutdown hooks.
72    RunningHooks = 5,
73    /// Server fully stopped.
74    Stopped = 6,
75}
76
77impl ShutdownPhase {
78    /// Returns true if the server should stop accepting connections.
79    #[must_use]
80    pub fn should_reject_connections(self) -> bool {
81        self as u8 >= Self::StopAccepting as u8
82    }
83
84    /// Returns true if new requests should receive 503.
85    #[must_use]
86    pub fn should_reject_requests(self) -> bool {
87        self as u8 >= Self::ShutdownFlagged as u8
88    }
89
90    /// Returns true if the server is in a shutdown state.
91    #[must_use]
92    pub fn is_shutting_down(self) -> bool {
93        self as u8 >= Self::StopAccepting as u8
94    }
95
96    /// Returns true if cleanup is complete.
97    #[must_use]
98    pub fn is_stopped(self) -> bool {
99        self == Self::Stopped
100    }
101}
102
103impl From<u8> for ShutdownPhase {
104    fn from(value: u8) -> Self {
105        match value {
106            0 => Self::Running,
107            1 => Self::StopAccepting,
108            2 => Self::ShutdownFlagged,
109            3 => Self::GracePeriod,
110            4 => Self::Cancelling,
111            5 => Self::RunningHooks,
112            _ => Self::Stopped,
113        }
114    }
115}
116
117// ============================================================================
118// Shutdown State (Shared)
119// ============================================================================
120
121/// Shared state for shutdown coordination.
122struct ShutdownState {
123    /// Current shutdown phase.
124    phase: AtomicU8,
125    /// Flag for forced shutdown (second signal).
126    forced: AtomicBool,
127    /// Wakers waiting for shutdown.
128    wakers: parking_lot::Mutex<Vec<Waker>>,
129    /// Shutdown hooks to run.
130    hooks: parking_lot::Mutex<Vec<ShutdownHook>>,
131    /// In-flight request count.
132    in_flight: std::sync::atomic::AtomicUsize,
133}
134
135impl ShutdownState {
136    fn new() -> Self {
137        Self {
138            phase: AtomicU8::new(ShutdownPhase::Running as u8),
139            forced: AtomicBool::new(false),
140            wakers: parking_lot::Mutex::new(Vec::new()),
141            hooks: parking_lot::Mutex::new(Vec::new()),
142            in_flight: std::sync::atomic::AtomicUsize::new(0),
143        }
144    }
145
146    fn phase(&self) -> ShutdownPhase {
147        ShutdownPhase::from(self.phase.load(Ordering::Acquire))
148    }
149
150    fn set_phase(&self, phase: ShutdownPhase) {
151        self.phase.store(phase as u8, Ordering::Release);
152        self.wake_all();
153    }
154
155    fn try_advance_phase(&self, from: ShutdownPhase, to: ShutdownPhase) -> bool {
156        self.phase
157            .compare_exchange(from as u8, to as u8, Ordering::AcqRel, Ordering::Acquire)
158            .is_ok()
159    }
160
161    fn is_forced(&self) -> bool {
162        self.forced.load(Ordering::Acquire)
163    }
164
165    fn set_forced(&self) {
166        self.forced.store(true, Ordering::Release);
167        self.wake_all();
168    }
169
170    fn wake_all(&self) {
171        let wakers = std::mem::take(&mut *self.wakers.lock());
172        for waker in wakers {
173            waker.wake();
174        }
175    }
176
177    fn register_waker(&self, waker: &Waker) {
178        let mut wakers = self.wakers.lock();
179        if !wakers.iter().any(|w| w.will_wake(waker)) {
180            wakers.push(waker.clone());
181        }
182    }
183
184    fn increment_in_flight(&self) -> usize {
185        self.in_flight.fetch_add(1, Ordering::AcqRel) + 1
186    }
187
188    fn decrement_in_flight(&self) -> usize {
189        self.in_flight.fetch_sub(1, Ordering::AcqRel) - 1
190    }
191
192    fn in_flight_count(&self) -> usize {
193        self.in_flight.load(Ordering::Acquire)
194    }
195}
196
197// ============================================================================
198// Shutdown Controller
199// ============================================================================
200
201/// Controller for initiating and coordinating shutdown.
202///
203/// Create one controller per server and share it via [`subscribe()`](Self::subscribe).
204#[derive(Clone)]
205pub struct ShutdownController {
206    state: Arc<ShutdownState>,
207}
208
209impl ShutdownController {
210    /// Create a new shutdown controller.
211    #[must_use]
212    pub fn new() -> Self {
213        Self {
214            state: Arc::new(ShutdownState::new()),
215        }
216    }
217
218    /// Create a receiver to wait for shutdown signals.
219    #[must_use]
220    pub fn subscribe(&self) -> ShutdownReceiver {
221        ShutdownReceiver {
222            state: Arc::clone(&self.state),
223        }
224    }
225
226    /// Get the current shutdown phase.
227    #[must_use]
228    pub fn phase(&self) -> ShutdownPhase {
229        self.state.phase()
230    }
231
232    /// Check if shutdown has been initiated.
233    #[must_use]
234    pub fn is_shutting_down(&self) -> bool {
235        self.state.phase().is_shutting_down()
236    }
237
238    /// Check if forced shutdown was requested.
239    #[must_use]
240    pub fn is_forced(&self) -> bool {
241        self.state.is_forced()
242    }
243
244    /// Initiate graceful shutdown.
245    ///
246    /// This begins the shutdown sequence. If shutdown is already in progress,
247    /// calling this again triggers forced shutdown.
248    pub fn shutdown(&self) {
249        let current = self.state.phase();
250        if current == ShutdownPhase::Running {
251            self.state.set_phase(ShutdownPhase::StopAccepting);
252        } else if !self.state.is_forced() {
253            // Second signal = forced shutdown
254            self.state.set_forced();
255        }
256    }
257
258    /// Force immediate shutdown without grace period.
259    pub fn force_shutdown(&self) {
260        self.state.set_forced();
261        self.state.set_phase(ShutdownPhase::Cancelling);
262    }
263
264    /// Advance to the next shutdown phase.
265    ///
266    /// Returns `true` if the phase was advanced.
267    pub fn advance_phase(&self) -> bool {
268        let current = self.state.phase();
269        let next = match current {
270            ShutdownPhase::Running => ShutdownPhase::StopAccepting,
271            ShutdownPhase::StopAccepting => ShutdownPhase::ShutdownFlagged,
272            ShutdownPhase::ShutdownFlagged => ShutdownPhase::GracePeriod,
273            ShutdownPhase::GracePeriod => ShutdownPhase::Cancelling,
274            ShutdownPhase::Cancelling => ShutdownPhase::RunningHooks,
275            ShutdownPhase::RunningHooks => ShutdownPhase::Stopped,
276            ShutdownPhase::Stopped => return false,
277        };
278        self.state.try_advance_phase(current, next)
279    }
280
281    /// Register a shutdown hook to run during the RunningHooks phase.
282    ///
283    /// Hooks run in LIFO order (last registered runs first).
284    pub fn register_hook<F>(&self, hook: F)
285    where
286        F: FnOnce() + Send + 'static,
287    {
288        let mut hooks = self.state.hooks.lock();
289        hooks.push(ShutdownHook::Sync(Box::new(hook)));
290    }
291
292    /// Register an async shutdown hook.
293    pub fn register_async_hook<F, Fut>(&self, hook: F)
294    where
295        F: FnOnce() -> Fut + Send + 'static,
296        Fut: Future<Output = ()> + Send + 'static,
297    {
298        let mut hooks = self.state.hooks.lock();
299        hooks.push(ShutdownHook::AsyncFactory(Box::new(move || {
300            Box::pin(hook())
301        })));
302    }
303
304    /// Pop and return the next shutdown hook to run.
305    ///
306    /// Returns `None` when all hooks have been run.
307    pub fn pop_hook(&self) -> Option<ShutdownHook> {
308        let mut hooks = self.state.hooks.lock();
309        hooks.pop()
310    }
311
312    /// Returns the number of registered hooks.
313    #[must_use]
314    pub fn hook_count(&self) -> usize {
315        self.state.hooks.lock().len()
316    }
317
318    /// Track a new in-flight request.
319    ///
320    /// Returns a guard that decrements the count when dropped.
321    #[must_use]
322    pub fn track_request(&self) -> InFlightGuard {
323        self.state.increment_in_flight();
324        InFlightGuard {
325            state: Arc::clone(&self.state),
326        }
327    }
328
329    /// Get the current in-flight request count.
330    #[must_use]
331    pub fn in_flight_count(&self) -> usize {
332        self.state.in_flight_count()
333    }
334}
335
336impl Default for ShutdownController {
337    fn default() -> Self {
338        Self::new()
339    }
340}
341
342// ============================================================================
343// Shutdown Receiver
344// ============================================================================
345
346/// Receiver for shutdown signals.
347///
348/// Obtained from [`ShutdownController::subscribe()`].
349#[derive(Clone)]
350pub struct ShutdownReceiver {
351    state: Arc<ShutdownState>,
352}
353
354impl ShutdownReceiver {
355    /// Wait for shutdown to be initiated.
356    ///
357    /// Returns immediately if shutdown is already in progress.
358    pub async fn wait(&self) {
359        ShutdownWaitFuture { state: &self.state }.await
360    }
361
362    /// Get the current shutdown phase.
363    #[must_use]
364    pub fn phase(&self) -> ShutdownPhase {
365        self.state.phase()
366    }
367
368    /// Check if shutdown has been initiated.
369    #[must_use]
370    pub fn is_shutting_down(&self) -> bool {
371        self.state.phase().is_shutting_down()
372    }
373
374    /// Check if forced shutdown was requested.
375    #[must_use]
376    pub fn is_forced(&self) -> bool {
377        self.state.is_forced()
378    }
379}
380
381/// Future for waiting on shutdown.
382struct ShutdownWaitFuture<'a> {
383    state: &'a ShutdownState,
384}
385
386impl Future for ShutdownWaitFuture<'_> {
387    type Output = ();
388
389    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
390        if self.state.phase().is_shutting_down() {
391            Poll::Ready(())
392        } else {
393            self.state.register_waker(cx.waker());
394            // Double-check after registering
395            if self.state.phase().is_shutting_down() {
396                Poll::Ready(())
397            } else {
398                Poll::Pending
399            }
400        }
401    }
402}
403
404// ============================================================================
405// In-Flight Guard
406// ============================================================================
407
408/// RAII guard for tracking in-flight requests.
409///
410/// Decrements the in-flight count when dropped.
411pub struct InFlightGuard {
412    state: Arc<ShutdownState>,
413}
414
415impl Drop for InFlightGuard {
416    fn drop(&mut self) {
417        self.state.decrement_in_flight();
418    }
419}
420
421// ============================================================================
422// Shutdown Hooks
423// ============================================================================
424
425/// A shutdown hook to run during cleanup.
426pub enum ShutdownHook {
427    /// Synchronous cleanup function.
428    Sync(Box<dyn FnOnce() + Send>),
429    /// Factory for async cleanup future.
430    AsyncFactory(Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send>),
431}
432
433impl ShutdownHook {
434    /// Run the hook synchronously.
435    ///
436    /// For async hooks, this returns the future to await.
437    pub fn run(self) -> Option<Pin<Box<dyn Future<Output = ()> + Send>>> {
438        match self {
439            Self::Sync(f) => {
440                f();
441                None
442            }
443            Self::AsyncFactory(f) => Some(f()),
444        }
445    }
446}
447
448// ============================================================================
449// Graceful Shutdown Builder
450// ============================================================================
451
452/// Configuration for graceful shutdown.
453#[derive(Clone)]
454pub struct GracefulConfig {
455    /// Grace period for in-flight requests.
456    pub grace_period: Duration,
457    /// Budget allocated to cleanup operations.
458    pub cleanup_budget: Budget,
459    /// Log shutdown events.
460    pub log_events: bool,
461}
462
463impl Default for GracefulConfig {
464    fn default() -> Self {
465        Self {
466            grace_period: Duration::from_secs(30),
467            cleanup_budget: Budget::new()
468                .with_poll_quota(500)
469                .with_deadline(asupersync::Time::from_secs(5)),
470            log_events: true,
471        }
472    }
473}
474
475/// Builder for graceful shutdown.
476pub struct GracefulShutdown {
477    receiver: ShutdownReceiver,
478    config: GracefulConfig,
479}
480
481impl GracefulShutdown {
482    /// Create a new graceful shutdown builder.
483    #[must_use]
484    pub fn new(receiver: ShutdownReceiver) -> Self {
485        Self {
486            receiver,
487            config: GracefulConfig::default(),
488        }
489    }
490
491    /// Set the grace period for in-flight requests.
492    #[must_use]
493    pub fn grace_period(mut self, duration: Duration) -> Self {
494        self.config.grace_period = duration;
495        self
496    }
497
498    /// Set the cleanup budget.
499    #[must_use]
500    pub fn cleanup_budget(mut self, budget: Budget) -> Self {
501        self.config.cleanup_budget = budget;
502        self
503    }
504
505    /// Enable or disable event logging.
506    #[must_use]
507    pub fn log_events(mut self, enabled: bool) -> Self {
508        self.config.log_events = enabled;
509        self
510    }
511
512    /// Run a future with graceful shutdown support.
513    ///
514    /// Returns when either:
515    /// - The future completes normally
516    /// - Graceful shutdown completes
517    /// - Forced shutdown is triggered
518    pub async fn run<F, T>(self, fut: F) -> ShutdownOutcome<T>
519    where
520        F: Future<Output = T>,
521    {
522        use std::pin::pin;
523        use std::task::Poll;
524
525        let mut fut = pin!(fut);
526
527        std::future::poll_fn(|cx| {
528            // Drive the main future forward first; if it completes we return immediately.
529            if let Poll::Ready(v) = fut.as_mut().poll(cx) {
530                return Poll::Ready(ShutdownOutcome::Completed(v));
531            }
532
533            // If shutdown has started, return the appropriate outcome.
534            //
535            // Note: `forced` is an orthogonal flag; if it is set we treat it as higher priority.
536            if self.receiver.state.is_forced() {
537                return Poll::Ready(ShutdownOutcome::ForcedShutdown);
538            }
539            if self.receiver.state.phase().is_shutting_down() {
540                return Poll::Ready(ShutdownOutcome::GracefulShutdown);
541            }
542
543            // Register to be woken on any shutdown phase/forced transition.
544            self.receiver.state.register_waker(cx.waker());
545            // Double-check after registering to avoid missing a transition.
546            if self.receiver.state.is_forced() {
547                Poll::Ready(ShutdownOutcome::ForcedShutdown)
548            } else if self.receiver.state.phase().is_shutting_down() {
549                Poll::Ready(ShutdownOutcome::GracefulShutdown)
550            } else {
551                Poll::Pending
552            }
553        })
554        .await
555    }
556
557    /// Get the configuration.
558    #[must_use]
559    pub fn config(&self) -> &GracefulConfig {
560        &self.config
561    }
562}
563
564/// Outcome of running with graceful shutdown.
565#[derive(Debug)]
566pub enum ShutdownOutcome<T> {
567    /// The future completed normally.
568    Completed(T),
569    /// Graceful shutdown was triggered.
570    GracefulShutdown,
571    /// Forced shutdown was triggered (second signal).
572    ForcedShutdown,
573}
574
575impl<T> ShutdownOutcome<T> {
576    /// Returns `true` if the future completed normally.
577    #[must_use]
578    pub fn is_completed(&self) -> bool {
579        matches!(self, Self::Completed(_))
580    }
581
582    /// Returns `true` if shutdown was triggered.
583    #[must_use]
584    pub fn is_shutdown(&self) -> bool {
585        matches!(self, Self::GracefulShutdown | Self::ForcedShutdown)
586    }
587
588    /// Returns `true` if forced shutdown was triggered.
589    #[must_use]
590    pub fn is_forced(&self) -> bool {
591        matches!(self, Self::ForcedShutdown)
592    }
593
594    /// Extract the completed value, if any.
595    #[must_use]
596    pub fn into_completed(self) -> Option<T> {
597        match self {
598            Self::Completed(v) => Some(v),
599            _ => None,
600        }
601    }
602}
603
604// ============================================================================
605// Request Budget Subdivision
606// ============================================================================
607
608/// Calculate subdivided budget for a request during grace period.
609///
610/// During graceful shutdown, in-flight requests get a proportional share
611/// of the remaining grace period.
612///
613/// # Arguments
614///
615/// * `grace_remaining` - Time remaining in the grace period
616/// * `in_flight_count` - Number of in-flight requests
617/// * `original_budget` - The request's original budget (if any)
618///
619/// # Returns
620///
621/// A budget that is the minimum of:
622/// - The remaining grace period divided by in-flight count
623/// - The original budget (if provided)
624#[must_use]
625pub fn subdivide_grace_budget(
626    grace_remaining: Duration,
627    in_flight_count: usize,
628    original_budget: Option<Budget>,
629) -> Budget {
630    use asupersync::Time;
631
632    let count = in_flight_count.max(1);
633    let per_request = grace_remaining / count as u32;
634
635    // Convert Duration to Time (nanoseconds since epoch, but used as a relative deadline)
636    let deadline_nanos = per_request.as_nanos() as u64;
637    let grace_budget = Budget::new().with_deadline(Time::from_nanos(deadline_nanos));
638
639    match original_budget {
640        Some(original) => original.meet(grace_budget),
641        None => grace_budget,
642    }
643}
644
645// ============================================================================
646// Cancel Reason for Shutdown
647// ============================================================================
648
649/// Create a cancel reason for server shutdown.
650#[must_use]
651pub fn shutdown_cancel_reason() -> CancelReason {
652    CancelReason::shutdown()
653}
654
655/// Create a cancel reason for grace period expiry.
656#[must_use]
657pub fn grace_expired_cancel_reason() -> CancelReason {
658    CancelReason::timeout()
659}
660
661// ============================================================================
662// Request Context Integration
663// ============================================================================
664
665/// Extension trait for checking shutdown status from request context.
666pub trait ShutdownAware {
667    /// Check if the server is shutting down.
668    fn is_shutting_down(&self) -> bool;
669
670    /// Get the current shutdown phase, if available.
671    fn shutdown_phase(&self) -> Option<ShutdownPhase>;
672}
673
674// ============================================================================
675// Tests
676// ============================================================================
677
678#[cfg(test)]
679mod tests {
680    use super::*;
681
682    #[test]
683    fn shutdown_phase_transitions() {
684        assert!(!ShutdownPhase::Running.should_reject_connections());
685        assert!(ShutdownPhase::StopAccepting.should_reject_connections());
686        assert!(ShutdownPhase::ShutdownFlagged.should_reject_requests());
687        assert!(ShutdownPhase::GracePeriod.is_shutting_down());
688        assert!(ShutdownPhase::Stopped.is_stopped());
689    }
690
691    #[test]
692    fn controller_basic() {
693        let controller = ShutdownController::new();
694        assert_eq!(controller.phase(), ShutdownPhase::Running);
695        assert!(!controller.is_shutting_down());
696
697        controller.shutdown();
698        assert_eq!(controller.phase(), ShutdownPhase::StopAccepting);
699        assert!(controller.is_shutting_down());
700    }
701
702    #[test]
703    fn controller_double_shutdown_forces() {
704        let controller = ShutdownController::new();
705        controller.shutdown();
706        assert!(!controller.is_forced());
707
708        controller.shutdown();
709        assert!(controller.is_forced());
710    }
711
712    #[test]
713    fn controller_advance_phase() {
714        let controller = ShutdownController::new();
715
716        assert!(controller.advance_phase());
717        assert_eq!(controller.phase(), ShutdownPhase::StopAccepting);
718
719        assert!(controller.advance_phase());
720        assert_eq!(controller.phase(), ShutdownPhase::ShutdownFlagged);
721
722        assert!(controller.advance_phase());
723        assert_eq!(controller.phase(), ShutdownPhase::GracePeriod);
724
725        assert!(controller.advance_phase());
726        assert_eq!(controller.phase(), ShutdownPhase::Cancelling);
727
728        assert!(controller.advance_phase());
729        assert_eq!(controller.phase(), ShutdownPhase::RunningHooks);
730
731        assert!(controller.advance_phase());
732        assert_eq!(controller.phase(), ShutdownPhase::Stopped);
733
734        // Can't advance past Stopped
735        assert!(!controller.advance_phase());
736    }
737
738    #[test]
739    fn in_flight_tracking() {
740        let controller = ShutdownController::new();
741        assert_eq!(controller.in_flight_count(), 0);
742
743        let guard1 = controller.track_request();
744        assert_eq!(controller.in_flight_count(), 1);
745
746        let guard2 = controller.track_request();
747        assert_eq!(controller.in_flight_count(), 2);
748
749        drop(guard1);
750        assert_eq!(controller.in_flight_count(), 1);
751
752        drop(guard2);
753        assert_eq!(controller.in_flight_count(), 0);
754    }
755
756    #[test]
757    fn shutdown_hooks_lifo() {
758        let controller = ShutdownController::new();
759        let order: Arc<parking_lot::Mutex<Vec<i32>>> =
760            Arc::new(parking_lot::Mutex::new(Vec::new()));
761
762        let order1 = Arc::clone(&order);
763        controller.register_hook(move || order1.lock().push(1));
764
765        let order2 = Arc::clone(&order);
766        controller.register_hook(move || order2.lock().push(2));
767
768        let order3 = Arc::clone(&order);
769        controller.register_hook(move || order3.lock().push(3));
770
771        assert_eq!(controller.hook_count(), 3);
772
773        // Pop and run in LIFO order
774        while let Some(hook) = controller.pop_hook() {
775            hook.run();
776        }
777
778        assert_eq!(*order.lock(), vec![3, 2, 1]);
779    }
780
781    #[test]
782    fn subdivide_grace_budget_basic() {
783        let grace = Duration::from_secs(30);
784        let budget = subdivide_grace_budget(grace, 3, None);
785
786        // Each request gets 10 seconds
787        assert!(budget.deadline.is_some());
788    }
789
790    #[test]
791    fn subdivide_grace_budget_respects_original() {
792        use asupersync::Time;
793
794        let grace = Duration::from_secs(30);
795        let original = Budget::new().with_deadline(Time::from_secs(5));
796        let budget = subdivide_grace_budget(grace, 3, Some(original));
797
798        // Should respect the tighter (original) deadline
799        assert!(budget.deadline.is_some());
800    }
801
802    #[test]
803    fn receiver_is_shutting_down() {
804        let controller = ShutdownController::new();
805        let receiver = controller.subscribe();
806
807        assert!(!receiver.is_shutting_down());
808
809        controller.shutdown();
810        assert!(receiver.is_shutting_down());
811    }
812
813    #[test]
814    fn graceful_shutdown_run_completed() {
815        let controller = ShutdownController::new();
816        let shutdown = GracefulShutdown::new(controller.subscribe());
817
818        let rt = asupersync::runtime::RuntimeBuilder::current_thread()
819            .build()
820            .expect("runtime must build");
821        let out = rt.block_on(async { shutdown.run(async { 42i32 }).await });
822
823        assert!(matches!(out, ShutdownOutcome::Completed(42)));
824    }
825
826    #[test]
827    fn graceful_shutdown_run_graceful_shutdown() {
828        let controller = ShutdownController::new();
829        controller.shutdown();
830        let shutdown = GracefulShutdown::new(controller.subscribe());
831
832        let rt = asupersync::runtime::RuntimeBuilder::current_thread()
833            .build()
834            .expect("runtime must build");
835        let out = rt.block_on(async { shutdown.run(std::future::pending::<i32>()).await });
836
837        assert!(matches!(out, ShutdownOutcome::GracefulShutdown));
838    }
839
840    #[test]
841    fn graceful_shutdown_run_forced_shutdown() {
842        let controller = ShutdownController::new();
843        controller.force_shutdown();
844        let shutdown = GracefulShutdown::new(controller.subscribe());
845
846        let rt = asupersync::runtime::RuntimeBuilder::current_thread()
847            .build()
848            .expect("runtime must build");
849        let out = rt.block_on(async { shutdown.run(std::future::pending::<i32>()).await });
850
851        assert!(matches!(out, ShutdownOutcome::ForcedShutdown));
852    }
853
854    #[test]
855    fn shutdown_outcome_accessors() {
856        let completed: ShutdownOutcome<i32> = ShutdownOutcome::Completed(42);
857        assert!(completed.is_completed());
858        assert!(!completed.is_shutdown());
859        assert_eq!(completed.into_completed(), Some(42));
860
861        let graceful: ShutdownOutcome<i32> = ShutdownOutcome::GracefulShutdown;
862        assert!(!graceful.is_completed());
863        assert!(graceful.is_shutdown());
864        assert!(!graceful.is_forced());
865
866        let forced: ShutdownOutcome<i32> = ShutdownOutcome::ForcedShutdown;
867        assert!(forced.is_shutdown());
868        assert!(forced.is_forced());
869    }
870
871    #[test]
872    fn cancel_reasons() {
873        let shutdown = shutdown_cancel_reason();
874        assert!(shutdown.is_shutdown());
875
876        let grace = grace_expired_cancel_reason();
877        assert!(!grace.is_shutdown());
878    }
879}