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}