merge 036, everything broken

core/startos/src/util/future.rs

@@ -0,0 +1,119 @@
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use futures::future::abortable;
+use futures::stream::{AbortHandle, Abortable};
+use futures::Future;
+use tokio::sync::watch;
+
+#[pin_project::pin_project(PinnedDrop)]
+pub struct DropSignaling<F> {
+    #[pin]
+    fut: F,
+    on_drop: watch::Sender<bool>,
+}
+impl<F> DropSignaling<F> {
+    pub fn new(fut: F) -> Self {
+        Self {
+            fut,
+            on_drop: watch::channel(false).0,
+        }
+    }
+    pub fn subscribe(&self) -> DropHandle {
+        DropHandle(self.on_drop.subscribe())
+    }
+}
+impl<F> Future for DropSignaling<F>
+where
+    F: Future,
+{
+    type Output = F::Output;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        this.fut.poll(cx)
+    }
+}
+#[pin_project::pinned_drop]
+impl<F> PinnedDrop for DropSignaling<F> {
+    fn drop(self: Pin<&mut Self>) {
+        let _ = self.on_drop.send(true);
+    }
+}
+
+#[derive(Clone)]
+pub struct DropHandle(watch::Receiver<bool>);
+impl DropHandle {
+    pub async fn wait(&mut self) {
+        let _ = self.0.wait_for(|a| *a).await;
+    }
+}
+
+#[pin_project::pin_project]
+pub struct RemoteCancellable<F> {
+    #[pin]
+    fut: Abortable<DropSignaling<F>>,
+    on_drop: DropHandle,
+    handle: AbortHandle,
+}
+impl<F: Future> RemoteCancellable<F> {
+    pub fn new(fut: F) -> Self {
+        let sig_fut = DropSignaling::new(fut);
+        let on_drop = sig_fut.subscribe();
+        let (fut, handle) = abortable(sig_fut);
+        Self {
+            fut,
+            on_drop,
+            handle,
+        }
+    }
+}
+impl<F> RemoteCancellable<F> {
+    pub fn cancellation_handle(&self) -> CancellationHandle {
+        CancellationHandle {
+            on_drop: self.on_drop.clone(),
+            handle: self.handle.clone(),
+        }
+    }
+}
+impl<F> Future for RemoteCancellable<F>
+where
+    F: Future,
+{
+    type Output = Option<F::Output>;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        this.fut.poll(cx).map(|a| a.ok())
+    }
+}
+
+#[derive(Clone)]
+pub struct CancellationHandle {
+    on_drop: DropHandle,
+    handle: AbortHandle,
+}
+impl CancellationHandle {
+    pub fn cancel(&mut self) {
+        self.handle.abort();
+    }
+
+    pub async fn cancel_and_wait(&mut self) {
+        self.handle.abort();
+        self.on_drop.wait().await
+    }
+}
+
+#[tokio::test]
+async fn test_cancellable() {
+    use std::sync::Arc;
+
+    let arc = Arc::new(());
+    let weak = Arc::downgrade(&arc);
+    let cancellable = RemoteCancellable::new(async move {
+        futures::future::pending::<()>().await;
+        drop(arc)
+    });
+    let mut handle = cancellable.cancellation_handle();
+    tokio::spawn(cancellable);
+    handle.cancel_and_wait().await;
+    assert!(weak.strong_count() == 0);
+}