use std::cell::RefCell;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
thread_local! {
static CURRENT: RefCell<Option<CancelToken>> = const { RefCell::new(None) };
}
#[derive(Debug, Clone, Default)]
pub struct CancelToken {
flag: Arc<AtomicBool>,
}
impl CancelToken {
pub fn new() -> Self {
Self::default()
}
pub fn cancel(&self) {
self.flag.store(true, Ordering::Relaxed);
}
pub fn is_cancelled(&self) -> bool {
self.flag.load(Ordering::Relaxed)
}
pub fn reset(&self) {
self.flag.store(false, Ordering::Relaxed);
}
pub fn scope<R>(&self, run: impl FnOnce() -> R) -> R {
struct Restore(Option<CancelToken>);
impl Drop for Restore {
fn drop(&mut self) {
let previous = self.0.take();
let _ = CURRENT.try_with(|current| {
if let Ok(mut current) = current.try_borrow_mut() {
*current = previous;
}
});
}
}
let previous = CURRENT
.try_with(|current| {
current
.try_borrow_mut()
.ok()
.and_then(|mut current| current.replace(self.clone()))
})
.ok()
.flatten();
let _restore = Restore(previous);
run()
}
pub(crate) fn current() -> Option<CancelToken> {
CURRENT
.try_with(|current| current.try_borrow().ok().and_then(|c| c.clone()))
.ok()
.flatten()
}
}