saluki_app/
util.rs

1//! General-purpose application utilities.
2
3use std::{ffi::OsStr, num::NonZeroUsize};
4
5use saluki_error::{generic_error, GenericError};
6use tracing::info;
7
8/// Environment variable that overrides the parallelism of the async runtime.
9const WORKER_THREADS_ENV_VAR: &str = "SALUKI_WORKER_THREADS";
10
11/// Minimum parallelism of the async runtime when the parallelism isn't overridden.
12///
13/// With a single worker thread, any task doing a long stretch of synchronous work stalls every other task on the
14/// runtime, including those responsible for health checks. A second worker thread lets those tasks keep running, even
15/// when the system reports a parallelism of one.
16const MIN_ASYNC_RUNTIME_PARALLELISM: usize = 2;
17
18/// Determines the parallelism -- the number of worker threads -- to use for the primary async runtime.
19///
20/// If `SALUKI_WORKER_THREADS` is set, its value is used as-is. Otherwise, the
21/// [available parallelism][std::thread::available_parallelism] of the system is used, or a minimum of two worker
22/// threads, whichever is greater.
23///
24/// # Errors
25///
26/// If `SALUKI_WORKER_THREADS` is set but isn't a positive integer, an error is returned.
27pub fn get_async_runtime_parallelism() -> Result<usize, GenericError> {
28    get_async_runtime_parallelism_inner(
29        std::env::var_os(WORKER_THREADS_ENV_VAR).as_deref(),
30        std::thread::available_parallelism().ok(),
31    )
32}
33
34fn get_async_runtime_parallelism_inner(
35    parallelism_override: Option<&OsStr>, available_parallelism: Option<NonZeroUsize>,
36) -> Result<usize, GenericError> {
37    match parallelism_override {
38        Some(raw_override) => raw_override
39            .to_str()
40            .and_then(|s| s.parse::<NonZeroUsize>().ok())
41            .map(NonZeroUsize::get)
42            .ok_or_else(|| {
43                generic_error!(
44                    "`{}` must be a positive integer, got {:?}.",
45                    WORKER_THREADS_ENV_VAR,
46                    raw_override
47                )
48            }),
49        None => Ok(available_parallelism
50            .map_or(1, NonZeroUsize::get)
51            .max(MIN_ASYNC_RUNTIME_PARALLELISM)),
52    }
53}
54
55/// Waits for a shutdown signal.
56///
57/// On Unix, this waits for either `SIGINT` or `SIGTERM`, either of which are used to request a graceful shutdown:
58/// `SIGINT` interactively (`Ctrl+C`), and `SIGTERM` by process supervisors (systemd, container runtimes,
59/// Kubernetes) during rollouts, evictions, node drains, and container shutdown.
60///
61/// On Windows, this waits for either `CTRL_C_EVENT` (interactively) or `CTRL_BREAK_EVENT`, the latter being what a
62/// parent process supervisor sends via `GenerateConsoleCtrlEvent` to request a graceful stop.
63pub async fn wait_for_shutdown_signal() {
64    #[cfg(unix)]
65    {
66        use tokio::signal::unix::{signal, SignalKind};
67
68        let mut sigterm = signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
69
70        tokio::select! {
71            _ = tokio::signal::ctrl_c() => info!("Received SIGINT, shutting down..."),
72            _ = sigterm.recv() => info!("Received SIGTERM, shutting down..."),
73        }
74    }
75
76    #[cfg(windows)]
77    {
78        let mut ctrl_break = tokio::signal::windows::ctrl_break().expect("failed to install CTRL_BREAK handler");
79
80        tokio::select! {
81            _ = tokio::signal::ctrl_c() => info!("Received CTRL_C, shutting down..."),
82            _ = ctrl_break.recv() => info!("Received CTRL_BREAK, shutting down..."),
83        }
84    }
85
86    #[cfg(not(any(unix, windows)))]
87    {
88        let _ = tokio::signal::ctrl_c().await;
89
90        info!("Received SIGINT, shutting down...");
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use std::{ffi::OsStr, num::NonZeroUsize};
97
98    use super::{get_async_runtime_parallelism_inner, MIN_ASYNC_RUNTIME_PARALLELISM, WORKER_THREADS_ENV_VAR};
99
100    #[test]
101    fn async_runtime_parallelism_uses_available_parallelism_above_minimum() {
102        let parallelism = get_async_runtime_parallelism_inner(None, NonZeroUsize::new(8)).expect("no override");
103        assert_eq!(parallelism, 8);
104    }
105
106    #[test]
107    fn async_runtime_parallelism_enforces_minimum_without_override() {
108        for available_parallelism in [None, NonZeroUsize::new(1)] {
109            let parallelism = get_async_runtime_parallelism_inner(None, available_parallelism).expect("no override");
110            assert_eq!(parallelism, MIN_ASYNC_RUNTIME_PARALLELISM);
111        }
112    }
113
114    #[test]
115    fn async_runtime_parallelism_uses_override_as_is() {
116        for (raw_override, expected) in [("1", 1), ("3", 3), ("16", 16)] {
117            let parallelism = get_async_runtime_parallelism_inner(Some(OsStr::new(raw_override)), NonZeroUsize::new(8))
118                .expect("valid override");
119            assert_eq!(parallelism, expected);
120        }
121    }
122
123    #[test]
124    fn async_runtime_parallelism_rejects_invalid_override() {
125        for raw_override in ["", "0", "-1", "four", " 4"] {
126            let error = get_async_runtime_parallelism_inner(Some(OsStr::new(raw_override)), NonZeroUsize::new(8))
127                .expect_err("invalid override is rejected");
128            assert!(
129                error.to_string().contains(WORKER_THREADS_ENV_VAR),
130                "error should name the environment variable: {error}"
131            );
132        }
133    }
134}
saluki_app/util.rs

saluki_app/
util.rs
