Skip to main content

runmat_runtime/
interrupt.rs

1use runmat_thread_local::runmat_thread_local;
2use std::cell::RefCell;
3use std::sync::{
4    atomic::{AtomicBool, Ordering},
5    Arc,
6};
7
8runmat_thread_local! {
9    static INTERRUPT_HANDLE: RefCell<Option<Arc<AtomicBool>>> = const { RefCell::new(None) };
10}
11
12pub struct InterruptGuard {
13    previous: Option<Arc<AtomicBool>>,
14    state: Option<std::rc::Rc<crate::context::RuntimeContextState>>,
15}
16
17impl InterruptGuard {
18    pub fn install(handle: Option<Arc<AtomicBool>>) -> Self {
19        if let Some(state) = active_state() {
20            let replacement = handle.unwrap_or_else(|| Arc::new(AtomicBool::new(false)));
21            let previous = state.cancellation.replace(replacement);
22            Self {
23                previous: Some(previous),
24                state: Some(state),
25            }
26        } else {
27            let previous = INTERRUPT_HANDLE.with(|slot| slot.replace(handle));
28            Self {
29                previous,
30                state: None,
31            }
32        }
33    }
34}
35
36impl Drop for InterruptGuard {
37    fn drop(&mut self) {
38        if let Some(state) = &self.state {
39            if let Some(previous) = self.previous.take() {
40                state.cancellation.replace(previous);
41            }
42        } else {
43            INTERRUPT_HANDLE.with(|slot| {
44                slot.replace(self.previous.take());
45            });
46        }
47    }
48}
49
50pub fn replace_interrupt(handle: Option<Arc<AtomicBool>>) -> InterruptGuard {
51    InterruptGuard::install(handle)
52}
53
54pub fn is_cancelled() -> bool {
55    if let Some(state) = active_state() {
56        return state.is_cancelled();
57    }
58    INTERRUPT_HANDLE.with(|slot| {
59        slot.borrow()
60            .as_ref()
61            .map(|flag| flag.load(Ordering::Relaxed))
62            .unwrap_or(false)
63    })
64}
65
66pub fn current_interrupt() -> Option<Arc<AtomicBool>> {
67    if let Some(state) = active_state() {
68        return Some(Arc::clone(&state.cancellation.borrow()));
69    }
70    INTERRUPT_HANDLE.with(|slot| slot.borrow().clone())
71}
72
73fn active_state() -> Option<std::rc::Rc<crate::context::RuntimeContextState>> {
74    crate::context::legacy::active().map(|context| std::rc::Rc::clone(context.state()))
75}