futures_util/future/future/
remote_handle.rs

1use {
2    crate::future::{CatchUnwind, FutureExt},
3    futures_channel::oneshot::{self, Receiver, Sender},
4    futures_core::{
5        future::Future,
6        ready,
7        task::{Context, Poll},
8    },
9    pin_project_lite::pin_project,
10    std::{
11        any::Any,
12        fmt,
13        panic::{self, AssertUnwindSafe},
14        pin::Pin,
15        sync::{
16            atomic::{AtomicBool, Ordering},
17            Arc,
18        },
19        thread,
20    },
21};
22
23/// The handle to a remote future returned by
24/// [`remote_handle`](crate::future::FutureExt::remote_handle). When you drop this,
25/// the remote future will be woken up to be dropped by the executor.
26///
27/// ## Unwind safety
28///
29/// When the remote future panics, [Remote] will catch the unwind and transfer it to
30/// the thread where `RemoteHandle` is being awaited. This is good for the common
31/// case where [Remote] is spawned on a threadpool. It is unlikely that other code
32/// in the executor working thread shares mutable data with the spawned future and we
33/// preserve the executor from losing its working threads.
34///
35/// If you run the future locally and send the handle of to be awaited elsewhere, you
36/// must be careful with regard to unwind safety because the thread in which the future
37/// is polled will keep running after the panic and the thread running the [RemoteHandle]
38/// will unwind.
39#[must_use = "dropping a remote handle cancels the underlying future"]
40#[derive(Debug)]
41#[cfg_attr(docsrs, doc(cfg(feature = "channel")))]
42pub struct RemoteHandle<T> {
43    rx: Receiver<thread::Result<T>>,
44    keep_running: Arc<AtomicBool>,
45}
46
47impl<T> RemoteHandle<T> {
48    /// Drops this handle *without* canceling the underlying future.
49    ///
50    /// This method can be used if you want to drop the handle, but let the
51    /// execution continue.
52    pub fn forget(self) {
53        self.keep_running.store(true, Ordering::SeqCst);
54    }
55}
56
57impl<T: 'static> Future for RemoteHandle<T> {
58    type Output = T;
59
60    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
61        match ready!(self.rx.poll_unpin(cx)) {
62            Ok(Ok(output)) => Poll::Ready(output),
63            // the remote future panicked.
64            Ok(Err(e)) => panic::resume_unwind(e),
65            // The oneshot sender was dropped.
66            Err(e) => panic::resume_unwind(Box::new(e)),
67        }
68    }
69}
70
71type SendMsg<Fut> = Result<<Fut as Future>::Output, Box<(dyn Any + Send + 'static)>>;
72
73pin_project! {
74    /// A future which sends its output to the corresponding `RemoteHandle`.
75    /// Created by [`remote_handle`](crate::future::FutureExt::remote_handle).
76    #[must_use = "futures do nothing unless you `.await` or poll them"]
77    #[cfg_attr(docsrs, doc(cfg(feature = "channel")))]
78    pub struct Remote<Fut: Future> {
79        tx: Option<Sender<SendMsg<Fut>>>,
80        keep_running: Arc<AtomicBool>,
81        #[pin]
82        future: CatchUnwind<AssertUnwindSafe<Fut>>,
83    }
84}
85
86impl<Fut: Future + fmt::Debug> fmt::Debug for Remote<Fut> {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        f.debug_tuple("Remote").field(&self.future).finish()
89    }
90}
91
92impl<Fut: Future> Future for Remote<Fut> {
93    type Output = ();
94
95    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
96        let this = self.project();
97
98        if this.tx.as_mut().unwrap().poll_canceled(cx).is_ready()
99            && !this.keep_running.load(Ordering::SeqCst)
100        {
101            // Cancelled, bail out
102            return Poll::Ready(());
103        }
104
105        let output = ready!(this.future.poll(cx));
106
107        // if the receiving end has gone away then that's ok, we just ignore the
108        // send error here.
109        drop(this.tx.take().unwrap().send(output));
110        Poll::Ready(())
111    }
112}
113
114pub(super) fn remote_handle<Fut: Future>(future: Fut) -> (Remote<Fut>, RemoteHandle<Fut::Output>) {
115    let (tx, rx) = oneshot::channel();
116    let keep_running = Arc::new(AtomicBool::new(false));
117
118    // Unwind Safety: See the docs for RemoteHandle.
119    let wrapped = Remote {
120        future: AssertUnwindSafe(future).catch_unwind(),
121        tx: Some(tx),
122        keep_running: keep_running.clone(),
123    };
124
125    (wrapped, RemoteHandle { rx, keep_running })
126}