strut_core/
spindown.rs

1use self::registry::SpindownRegistry;
2use crate::AppSpindownToken;
3use parking_lot::Mutex;
4use std::sync::OnceLock;
5use std::time::Duration;
6
7mod registry;
8pub mod token;
9
10// Global singleton spindown registry
11static GLOBAL: OnceLock<SpindownRegistry> = OnceLock::new();
12
13// Spindown timeout (stored statically to allow customizing)
14const DEFAULT_TIMEOUT_SECS: u64 = 2;
15static TIMEOUT_SECS: Mutex<u64> = Mutex::new(DEFAULT_TIMEOUT_SECS);
16
17/// A facade for interacting with the application’s global spindown registry.
18///
19/// Allows [registering](AppSpindown::register) arbitrary workloads and later
20/// [waiting](AppSpindown::completed) for all registered workloads to signal
21/// their graceful completion (within a configurable timeout).
22///
23/// ## Problem statement
24///
25/// Every application may choose to run asynchronous background tasks that hold
26/// some kind of resource. An example here would be a pool of database connections
27/// that is owned by a background task that lends access to the pool to any task
28/// that requests it. However, when the application is shut down, all background
29/// tasks are unceremoniously killed, which prevents proper clean-up, such as
30/// closing the database connections.
31///
32/// ## Spindown
33///
34/// To solve the problem, this crate enables the following flow.
35///
36/// ### Main thread
37///
38/// The main thread:
39///
40/// - spawns background tasks,
41/// - waits until the global [context](AppContext) is
42///   [terminated](AppContext::terminated)(e.g., in a [`select`](tokio::select)
43///   block, while also waiting for the main logic to finish),
44/// - [waits](AppSpindown::completed) for all background tasks to signal
45///   completion,
46/// - returns from the `main` function.
47///
48/// ### Background tasks
49///
50/// Meanwhile, each spawned background task:
51///
52/// - [registers](AppSpindown::register) with the global spindown registry,
53/// - also waits until the global [context](AppContext) is
54///   [terminated](AppContext::terminated) (e.g., in a [`select`](tokio::select)
55///   block, while also doing other work),
56/// - performs clean-up procedures (e.g., closes connections, etc.),
57/// - [punches out](AppSpindownToken::punch_out) the spindown token to signal
58///   completion.
59pub struct AppSpindown;
60
61impl AppSpindown {
62    /// Informs the application’s global spindown registry that a workload with
63    /// the given name (an arbitrary human-readable string) will need to be
64    /// awaited during the application spindown, giving it some time to perform
65    /// clean-up.
66    ///
67    /// The returned [`AppSpindownToken`] must be used by the registering
68    /// workload to signal back to the registry once it has gracefully completed.
69    pub fn register(name: impl AsRef<str>) -> AppSpindownToken {
70        // Retrieve global registry
71        let registry = Self::global_registry();
72
73        // Register workload
74        registry.register(name.as_ref())
75    }
76
77    /// Allows customizing the spindown timeout for the
78    /// [global singleton registry](Self::global_registry). Importantly, this
79    /// method must be called early on, before any interaction with the global
80    /// spindown registry, such as [registering](Self::register) a workload. If
81    /// called later, this method will have no effect.
82    pub fn set_timeout_secs(timeout_secs: impl Into<u64>) {
83        *TIMEOUT_SECS.lock() = timeout_secs.into();
84    }
85
86    /// Collects all previously [registered](AppSpindown::register) workloads,
87    /// and then waits (within a [timeout](Self::set_timeout_secs)) for them to
88    /// signal completion.
89    ///
90    /// This function is destructive, as it consumes the internally stored list
91    /// of workloads.
92    ///
93    /// The spindown is performed in repeated cycles (within a single shared
94    /// timeout). If new workloads are registered while previous ones are being
95    /// spun down, a new cycle is initiated to wait for the next batch. This is
96    /// repeated until no more registered workloads are found.
97    ///
98    /// Importantly, this function does **not** signal to the workloads to begin
99    /// their spindown. This is the job of the global
100    /// [`AppContext`](crate::AppContext).
101    pub async fn completed() {
102        // Retrieve global registry
103        let registry = Self::global_registry();
104
105        // Repeatedly await all workloads
106        let _ = registry.spun_down().await;
107    }
108
109    /// Retrieves the global (singleton) [`SpindownRegistry`], lazily
110    /// initialized.
111    fn global_registry() -> &'static SpindownRegistry {
112        GLOBAL.get_or_init(|| {
113            let timeout = Self::deduce_timeout();
114
115            SpindownRegistry::new(timeout)
116        })
117    }
118
119    /// Infers the timeout duration.
120    fn deduce_timeout() -> Duration {
121        // Grab lock
122        let timeout_secs = *TIMEOUT_SECS.lock();
123
124        Duration::from_secs(timeout_secs)
125    }
126}