Skip to main content

inferlab_runtime/
interrupt.rs

1use std::sync::Arc;
2use std::sync::OnceLock;
3use std::sync::atomic::{AtomicBool, Ordering};
4
5static INTERRUPTED: AtomicBool = AtomicBool::new(false);
6static CTRL_C_HANDLER: OnceLock<Result<(), InterruptInstallError>> = OnceLock::new();
7
8#[derive(Clone, Debug, thiserror::Error)]
9pub enum InterruptInstallError {
10    #[error(
11        "failed to install Ctrl-C handler: Ctrl-C error: Signal could not be found from the system ({signal})"
12    )]
13    NoSuchSignal { signal: String },
14    #[error(
15        "failed to install Ctrl-C handler: Ctrl-C error: Ctrl-C signal handler already registered"
16    )]
17    MultipleHandlers,
18    #[error("failed to install Ctrl-C handler: Ctrl-C error: Unexpected system error")]
19    System {
20        #[source]
21        source: Arc<std::io::Error>,
22    },
23}
24
25impl From<ctrlc::Error> for InterruptInstallError {
26    fn from(source: ctrlc::Error) -> Self {
27        match source {
28            ctrlc::Error::NoSuchSignal(signal) => Self::NoSuchSignal {
29                signal: format!("{signal:?}"),
30            },
31            ctrlc::Error::MultipleHandlers => Self::MultipleHandlers,
32            ctrlc::Error::System(source) => Self::System {
33                source: Arc::new(source),
34            },
35        }
36    }
37}
38
39pub fn prepare() -> Result<(), InterruptInstallError> {
40    INTERRUPTED.store(false, Ordering::SeqCst);
41    CTRL_C_HANDLER
42        .get_or_init(|| {
43            ctrlc::set_handler(|| INTERRUPTED.store(true, Ordering::SeqCst))
44                .map_err(InterruptInstallError::from)
45        })
46        .clone()
47}
48
49pub fn received() -> bool {
50    INTERRUPTED.load(Ordering::SeqCst)
51}
52
53#[cfg(test)]
54mod tests {
55    use super::InterruptInstallError;
56
57    #[test]
58    fn ctrlc_categories_remain_typed() {
59        let missing =
60            InterruptInstallError::from(ctrlc::Error::NoSuchSignal(ctrlc::SignalType::Ctrlc));
61        assert!(matches!(
62            missing,
63            InterruptInstallError::NoSuchSignal { .. }
64        ));
65
66        let duplicate = InterruptInstallError::from(ctrlc::Error::MultipleHandlers);
67        assert!(matches!(duplicate, InterruptInstallError::MultipleHandlers));
68
69        let system = InterruptInstallError::from(ctrlc::Error::System(std::io::Error::other(
70            "handler installation failed",
71        )));
72        assert!(matches!(system, InterruptInstallError::System { .. }));
73    }
74}