1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
use {
    crate::future::{CatchUnwind, FutureExt},
    futures_channel::oneshot::{self, Receiver, Sender},
    futures_core::{
        future::Future,
        ready,
        task::{Context, Poll},
    },
    pin_project_lite::pin_project,
    std::{
        any::Any,
        fmt,
        panic::{self, AssertUnwindSafe},
        pin::Pin,
        sync::{
            atomic::{AtomicBool, Ordering},
            Arc,
        },
        thread,
    },
};

/// The handle to a remote future returned by
/// [`remote_handle`](crate::future::FutureExt::remote_handle). When you drop this,
/// the remote future will be woken up to be dropped by the executor.
///
/// ## Unwind safety
///
/// When the remote future panics, [Remote] will catch the unwind and transfer it to
/// the thread where `RemoteHandle` is being awaited. This is good for the common
/// case where [Remote] is spawned on a threadpool. It is unlikely that other code
/// in the executor working thread shares mutable data with the spawned future and we
/// preserve the executor from losing its working threads.
///
/// If you run the future locally and send the handle of to be awaited elsewhere, you
/// must be careful with regard to unwind safety because the thread in which the future
/// is polled will keep running after the panic and the thread running the [RemoteHandle]
/// will unwind.
#[must_use = "dropping a remote handle cancels the underlying future"]
#[derive(Debug)]
#[cfg_attr(docsrs, doc(cfg(feature = "channel")))]
pub struct RemoteHandle<T> {
    rx: Receiver<thread::Result<T>>,
    keep_running: Arc<AtomicBool>,
}

impl<T> RemoteHandle<T> {
    /// Drops this handle *without* canceling the underlying future.
    ///
    /// This method can be used if you want to drop the handle, but let the
    /// execution continue.
    pub fn forget(self) {
        self.keep_running.store(true, Ordering::SeqCst);
    }
}

impl<T: 'static> Future for RemoteHandle<T> {
    type Output = T;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
        match ready!(self.rx.poll_unpin(cx)) {
            Ok(Ok(output)) => Poll::Ready(output),
            // the remote future panicked.
            Ok(Err(e)) => panic::resume_unwind(e),
            // The oneshot sender was dropped.
            Err(e) => panic::resume_unwind(Box::new(e)),
        }
    }
}

type SendMsg<Fut> = Result<<Fut as Future>::Output, Box<(dyn Any + Send + 'static)>>;

pin_project! {
    /// A future which sends its output to the corresponding `RemoteHandle`.
    /// Created by [`remote_handle`](crate::future::FutureExt::remote_handle).
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    #[cfg_attr(docsrs, doc(cfg(feature = "channel")))]
    pub struct Remote<Fut: Future> {
        tx: Option<Sender<SendMsg<Fut>>>,
        keep_running: Arc<AtomicBool>,
        #[pin]
        future: CatchUnwind<AssertUnwindSafe<Fut>>,
    }
}

impl<Fut: Future + fmt::Debug> fmt::Debug for Remote<Fut> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("Remote").field(&self.future).finish()
    }
}

impl<Fut: Future> Future for Remote<Fut> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let this = self.project();

        if this.tx.as_mut().unwrap().poll_canceled(cx).is_ready()
            && !this.keep_running.load(Ordering::SeqCst)
        {
            // Cancelled, bail out
            return Poll::Ready(());
        }

        let output = ready!(this.future.poll(cx));

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(this.tx.take().unwrap().send(output));
        Poll::Ready(())
    }
}

pub(super) fn remote_handle<Fut: Future>(future: Fut) -> (Remote<Fut>, RemoteHandle<Fut::Output>) {
    let (tx, rx) = oneshot::channel();
    let keep_running = Arc::new(AtomicBool::new(false));

    // Unwind Safety: See the docs for RemoteHandle.
    let wrapped = Remote {
        future: AssertUnwindSafe(future).catch_unwind(),
        tx: Some(tx),
        keep_running: keep_running.clone(),
    };

    (wrapped, RemoteHandle { rx, keep_running })
}