use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use futures_core::Stream;
pub type MsgStream<Msg> = Pin<Box<dyn Stream<Item = Msg> + Send>>;
pub struct Task {
cancel: Option<Arc<Cancel>>,
}
impl Task {
pub fn detach(mut self) {
self.cancel = None;
}
}
impl Drop for Task {
fn drop(&mut self) {
if let Some(cancel) = &self.cancel {
cancel.cancel();
}
}
}
pub struct Cancel {
cancelled: AtomicBool,
waker: Mutex<Option<Waker>>,
}
impl Cancel {
pub(crate) fn new() -> Self {
Self {
cancelled: AtomicBool::new(false),
waker: Mutex::new(None),
}
}
pub(crate) fn cancel(&self) {
self.cancelled.store(true, Ordering::SeqCst);
if let Some(waker) = self.waker.lock().unwrap().take() {
waker.wake();
}
}
pub fn is_cancelled(&self) -> bool {
self.cancelled.load(Ordering::SeqCst)
}
pub fn cancelled(self: &Arc<Self>) -> Cancelled {
Cancelled {
cancel: Arc::clone(self),
}
}
}
pub struct Cancelled {
cancel: Arc<Cancel>,
}
impl Future for Cancelled {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.cancel.is_cancelled() {
return Poll::Ready(());
}
*self.cancel.waker.lock().unwrap() = Some(cx.waker().clone());
if self.cancel.is_cancelled() {
Poll::Ready(())
} else {
Poll::Pending
}
}
}
pub enum Effect<Msg> {
Spawn {
stream: MsgStream<Msg>,
cancel: Arc<Cancel>,
},
}
pub(crate) fn spawn_effect<Msg>(
stream: impl Stream<Item = Msg> + Send + 'static,
) -> (Effect<Msg>, Task) {
let cancel = Arc::new(Cancel::new());
(
Effect::Spawn {
stream: Box::pin(stream),
cancel: Arc::clone(&cancel),
},
Task {
cancel: Some(cancel),
},
)
}
pub(crate) fn spawn_once_effect<Msg>(
future: impl Future<Output = Msg> + Send + 'static,
) -> (Effect<Msg>, Task) {
spawn_effect(FutureStream {
future: Some(Box::pin(future)),
})
}
struct FutureStream<F: Future> {
future: Option<Pin<Box<F>>>,
}
impl<F: Future> Stream for FutureStream<F> {
type Item = F::Output;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<F::Output>> {
let this = self.get_mut();
match &mut this.future {
Some(future) => match future.as_mut().poll(cx) {
Poll::Ready(value) => {
this.future = None;
Poll::Ready(Some(value))
}
Poll::Pending => Poll::Pending,
},
None => Poll::Ready(None),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
struct Pending;
impl Stream for Pending {
type Item = ();
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<()>> {
Poll::Pending
}
}
#[test]
fn drop_cancels() {
let (effect, task) = spawn_effect(Pending);
let Effect::Spawn { cancel, .. } = effect;
assert!(!cancel.is_cancelled());
drop(task);
assert!(cancel.is_cancelled());
}
#[test]
fn detach_does_not_cancel() {
let (effect, task) = spawn_effect(Pending);
let Effect::Spawn { cancel, .. } = effect;
task.detach();
assert!(!cancel.is_cancelled());
}
}