napparent_tabular/
cancel.rs

1//! Cooperative cancellation checked between pipeline batches.
2
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::Arc;
5
6pub const INTERRUPT_MSG: &str = "interrupted by user";
7
8type CancelHook = Arc<dyn Fn() -> Result<(), String> + Send + Sync>;
9
10/// Checked at the start of each batch iteration; safe to share across threads.
11pub struct CancelToken {
12    stop: Arc<AtomicBool>,
13    hook: Option<CancelHook>,
14}
15
16impl CancelToken {
17    /// Never triggers cancellation.
18    pub fn none() -> Self {
19        Self {
20            stop: Arc::new(AtomicBool::new(false)),
21            hook: None,
22        }
23    }
24
25    pub fn from_flag(stop: Arc<AtomicBool>) -> Self {
26        Self { stop, hook: None }
27    }
28
29    pub fn with_hook(hook: CancelHook) -> Self {
30        Self {
31            stop: Arc::new(AtomicBool::new(false)),
32            hook: Some(hook),
33        }
34    }
35
36    pub fn flag(&self) -> Arc<AtomicBool> {
37        Arc::clone(&self.stop)
38    }
39
40    pub fn request_stop(&self) {
41        self.stop.store(true, Ordering::Relaxed);
42    }
43
44    pub fn check(&self) -> Result<(), String> {
45        if let Some(hook) = &self.hook {
46            if let Err(e) = hook() {
47                self.stop.store(true, Ordering::Relaxed);
48                return Err(if e.contains("interrupted") {
49                    e
50                } else {
51                    format!("{INTERRUPT_MSG}: {e}")
52                });
53            }
54        }
55        if self.stop.load(Ordering::Relaxed) {
56            return Err(INTERRUPT_MSG.into());
57        }
58        Ok(())
59    }
60}
61
62/// Restores a no-op SIGINT handler when dropped.
63pub struct CtrlcGuard;
64
65impl CancelToken {
66    /// Install a process-wide SIGINT handler that sets the token stop flag.
67    #[cfg(feature = "progress")]
68    pub fn with_ctrlc_handler() -> Result<(Self, CtrlcGuard), String> {
69        let stop = Arc::new(AtomicBool::new(false));
70        let stop_handler = Arc::clone(&stop);
71        ctrlc::set_handler(move || {
72            stop_handler.store(true, Ordering::Relaxed);
73        })
74        .map_err(|e| format!("ctrlc handler: {e}"))?;
75        Ok((Self::from_flag(stop), CtrlcGuard))
76    }
77
78    /// Install a process-wide SIGINT handler that sets the token stop flag.
79    #[cfg(not(feature = "progress"))]
80    pub fn with_ctrlc_handler() -> Result<(Self, CtrlcGuard), String> {
81        Err("enable the `progress` feature for SIGINT handling".into())
82    }
83}
84
85impl Drop for CtrlcGuard {
86    fn drop(&mut self) {
87        #[cfg(feature = "progress")]
88        let _ = ctrlc::set_handler(|| {});
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn check_returns_err_when_flag_set() {
98        let token = CancelToken::none();
99        token.request_stop();
100        let err = token.check().unwrap_err();
101        assert!(err.contains("interrupted"));
102    }
103}
napparent_tabular/cancel.rs

napparent_tabular/
cancel.rs
