use tokio_util::sync::CancellationToken;
pub struct CancelSignal {
inner: CancellationToken,
}
impl CancelSignal {
#[must_use]
pub fn new() -> Self {
Self {
inner: CancellationToken::new(),
}
}
pub fn cancel(&self) {
self.inner.cancel();
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.inner.is_cancelled()
}
pub async fn notified(&self) {
self.inner.cancelled().await;
}
}
impl Default for CancelSignal {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
#[tokio::test]
async fn test_cancel_wakes_instantly() {
let signal = Arc::new(CancelSignal::new());
let signal_clone = signal.clone();
let handle = tokio::spawn(async move {
signal_clone.notified().await;
"woke up"
});
tokio::task::yield_now().await;
signal.cancel();
let result = handle.await.unwrap();
assert_eq!(result, "woke up");
}
#[tokio::test]
async fn test_notified_returns_immediately_if_already_cancelled() {
let signal = CancelSignal::new();
signal.cancel();
signal.notified().await;
}
#[test]
fn test_is_cancelled() {
let signal = CancelSignal::new();
assert!(!signal.is_cancelled());
signal.cancel();
assert!(signal.is_cancelled());
}
#[tokio::test]
async fn test_cancel_is_idempotent() {
let signal = Arc::new(CancelSignal::new());
signal.cancel();
signal.cancel();
signal.cancel();
assert!(signal.is_cancelled());
signal.notified().await;
}
#[tokio::test]
async fn test_multiple_waiters_all_wake() {
let signal = Arc::new(CancelSignal::new());
let mut handles = Vec::new();
for _ in 0..10 {
let s = signal.clone();
handles.push(tokio::spawn(async move {
s.notified().await;
true
}));
}
tokio::task::yield_now().await;
signal.cancel();
for handle in handles {
assert!(handle.await.unwrap());
}
}
}