use core::future::{Future, poll_fn};
use core::sync::atomic::{AtomicBool, Ordering};
use core::task::Poll;
use futures::task::AtomicWaker;
pub struct ShutdownSignal {
cancelled: AtomicBool,
waker: AtomicWaker,
}
impl ShutdownSignal {
pub const fn new() -> Self {
Self {
cancelled: AtomicBool::new(false),
waker: AtomicWaker::new(),
}
}
pub fn is_shutdown(&self) -> bool {
self.cancelled.load(Ordering::Acquire)
}
pub fn cancelled(&self) -> impl Future<Output = ()> + '_ {
poll_fn(|cx| {
if self.is_shutdown() {
return Poll::Ready(());
}
self.waker.register(cx.waker());
if self.is_shutdown() {
Poll::Ready(())
} else {
Poll::Pending
}
})
}
#[doc(hidden)]
pub fn cancel(&self) {
if !self.cancelled.swap(true, Ordering::AcqRel) {
self.waker.wake();
}
}
}
impl Default for ShutdownSignal {
fn default() -> Self {
Self::new()
}
}