use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct CancellationToken {
inner: Arc<AtomicBool>,
}
impl CancellationToken {
pub fn new() -> Self {
Self {
inner: Arc::new(AtomicBool::new(false)),
}
}
pub fn cancel(&self) {
self.inner.store(true, Ordering::SeqCst);
}
pub fn is_cancelled(&self) -> bool {
self.inner.load(Ordering::SeqCst)
}
pub async fn cancelled(&self) {
while !self.inner.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
}
}
impl Default for CancellationToken {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_token_is_not_cancelled() {
let token = CancellationToken::new();
assert!(!token.is_cancelled());
}
#[test]
fn test_cancel_sets_is_cancelled() {
let token = CancellationToken::new();
token.cancel();
assert!(token.is_cancelled());
}
#[test]
fn test_clone_shares_cancellation_state() {
let token = CancellationToken::new();
let clone = token.clone();
assert!(!token.is_cancelled());
assert!(!clone.is_cancelled());
clone.cancel();
assert!(token.is_cancelled());
assert!(clone.is_cancelled());
}
#[test]
fn test_multiple_clones_all_cancelled() {
let token = CancellationToken::new();
let c1 = token.clone();
let c2 = token.clone();
let c3 = token.clone();
token.cancel();
assert!(c1.is_cancelled());
assert!(c2.is_cancelled());
assert!(c3.is_cancelled());
}
#[tokio::test]
async fn test_cancelled_future_resolves() {
let token = CancellationToken::new();
let cloned = token.clone();
tokio::spawn(async move {
cloned.cancel();
});
token.cancelled().await;
assert!(token.is_cancelled());
}
}