use crate::platform::SharedPtr;
use futures::future::poll_fn;
use std::future::Future;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Poll, Waker};
pub type StreamResult<T> = Result<T, super::error::StreamError>;
#[derive(Clone, Default)]
pub struct WakerSet(SharedPtr<parking_lot::Mutex<Vec<Waker>>>);
impl WakerSet {
pub fn new() -> Self {
WakerSet(SharedPtr::new(parking_lot::Mutex::new(Vec::new())))
}
pub fn register(&self, waker: &Waker) {
let mut wakers = self.0.lock();
if !wakers.iter().any(|w| w.will_wake(waker)) {
wakers.push(waker.clone());
}
}
pub fn wake_all(&self) {
let mut wakers = self.0.lock();
for waker in wakers.drain(..) {
waker.wake();
}
}
}
#[derive(Default)]
struct AbortInner {
aborted: AtomicBool,
reason: parking_lot::Mutex<Option<String>>,
wakers: WakerSet,
}
#[derive(Clone, Default)]
pub struct AbortSignal {
inner: SharedPtr<AbortInner>,
}
impl AbortSignal {
pub fn aborted(&self) -> bool {
self.inner.aborted.load(Ordering::Acquire)
}
pub fn reason(&self) -> Option<String> {
self.inner.reason.lock().clone()
}
pub fn aborted_future(&self) -> impl Future<Output = ()> + 'static {
let inner = self.inner.clone();
poll_fn(move |cx| {
if inner.aborted.load(Ordering::Acquire) {
return Poll::Ready(());
}
inner.wakers.register(cx.waker());
if inner.aborted.load(Ordering::Acquire) {
Poll::Ready(())
} else {
Poll::Pending
}
})
}
fn trigger(&self, reason: Option<String>) {
let mut guard = self.inner.reason.lock();
if self.inner.aborted.load(Ordering::Acquire) {
return;
}
*guard = reason;
self.inner.aborted.store(true, Ordering::Release);
drop(guard);
self.inner.wakers.wake_all();
}
}
#[derive(Clone, Default)]
pub struct AbortController {
signal: AbortSignal,
}
impl AbortController {
pub fn new() -> Self {
Self::default()
}
pub fn signal(&self) -> AbortSignal {
self.signal.clone()
}
pub fn abort(&self, reason: Option<String>) {
self.signal.trigger(reason);
}
}