1 use {
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")))]
42 pub struct RemoteHandle<T> {
43 rx: Receiver<thread::Result<T>>,
44 keep_running: Arc<AtomicBool>,
45 }
46
47 impl<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.
forget(self)52 pub fn forget(self) {
53 self.keep_running.store(true, Ordering::SeqCst);
54 }
55 }
56
57 impl<T: 'static> Future for RemoteHandle<T> {
58 type Output = T;
59
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T>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
71 type SendMsg<Fut> = Result<<Fut as Future>::Output, Box<(dyn Any + Send + 'static)>>;
72
73 pin_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
86 impl<Fut: Future + fmt::Debug> fmt::Debug for Remote<Fut> {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result87 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88 f.debug_tuple("Remote").field(&self.future).finish()
89 }
90 }
91
92 impl<Fut: Future> Future for Remote<Fut> {
93 type Output = ();
94
poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()>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
remote_handle<Fut: Future>(future: Fut) -> (Remote<Fut>, RemoteHandle<Fut::Output>)114 pub(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 }
127