async_rs/implementors/
tokio.rs

1//! tokio implementation of async runtime definition traits
2
3use crate::{
4    Runtime,
5    sys::AsSysFd,
6    traits::{Executor, Reactor, RuntimeKit},
7    util::Task,
8};
9use async_compat::{Compat, CompatExt};
10use futures_core::Stream;
11use futures_io::{AsyncRead, AsyncWrite};
12use std::{
13    future::Future,
14    io::{self, Read, Write},
15    net::SocketAddr,
16    pin::Pin,
17    sync::Arc,
18    task::{Context, Poll},
19    time::{Duration, Instant},
20};
21use tokio::{
22    net::TcpStream,
23    runtime::{EnterGuard, Handle, Runtime as TokioRT},
24    time::Sleep,
25};
26use tokio_stream::{StreamExt, wrappers::IntervalStream};
27
28use task::TTask;
29
30/// Type alias for the tokio runtime
31pub type TokioRuntime = Runtime<Tokio>;
32
33impl TokioRuntime {
34    /// Create a new TokioRuntime backed by a freshly created tokio multi-threaded runtime.
35    pub fn tokio() -> io::Result<Self> {
36        Ok(Self::tokio_with_runtime(TokioRT::new()?))
37    }
38
39    /// Create a new TokioRuntime and bind it to the current tokio runtime by default.
40    #[must_use]
41    pub fn tokio_current() -> Self {
42        Self::new(Tokio::current())
43    }
44
45    /// Create a new TokioRuntime and bind it to the tokio runtime associated to this handle by default.
46    #[must_use]
47    pub fn tokio_with_handle(handle: Handle) -> Self {
48        Self::new(Tokio::default().with_handle(handle))
49    }
50
51    /// Create a new TokioRuntime and bind it to this tokio runtime.
52    #[must_use]
53    pub fn tokio_with_runtime(runtime: TokioRT) -> Self {
54        Self::new(Tokio::default().with_runtime(runtime))
55    }
56}
57
58/// What every entry point says when it cannot find a runtime to work with.
59///
60/// The ones returning an `io::Result` — `register` and `tcp_connect_addr` — report it, the rest
61/// panic with it, but none of them may leave the caller with tokio's own "there is no reactor
62/// running" thrown from somewhere further in, which says nothing about how to fix it.
63const NO_RUNTIME: &str = "no tokio runtime: use Runtime::tokio() or Runtime::tokio_with_handle()";
64
65/// The [`RuntimeKit`] implementation backed by the tokio async runtime
66#[derive(Default, Clone, Debug)]
67pub struct Tokio {
68    handle: Option<Handle>,
69    runtime: Option<Arc<TokioRT>>,
70}
71
72impl Tokio {
73    /// Bind to the tokio Runtime associated to this handle by default.
74    ///
75    /// A runtime given to [`with_runtime`](Self::with_runtime) wins over this one whichever order
76    /// the two are called in: only the owned runtime can be driven by
77    /// [`block_on`](crate::traits::Executor::block_on), so letting a handle override it would bind
78    /// half the kit to one runtime and half to another.
79    #[must_use]
80    pub fn with_handle(mut self, handle: Handle) -> Self {
81        self.handle = Some(handle);
82        self
83    }
84
85    /// Bind to this tokio runtime by default.
86    #[must_use]
87    pub fn with_runtime(mut self, runtime: TokioRT) -> Self {
88        let handle = runtime.handle().clone();
89        self.runtime = Some(Arc::new(runtime));
90        self.with_handle(handle)
91    }
92
93    /// Bind to the current tokio Runtime by default.
94    #[must_use]
95    pub fn current() -> Self {
96        Self::default().with_handle(Handle::current())
97    }
98
99    /// The runtime this kit is bound to, if any.
100    ///
101    /// Every entry point resolves through this so the kit cannot end up straddling two runtimes:
102    /// `with_runtime` also records a handle, but `with_handle` may be called afterwards, and
103    /// `block_on` can only drive the owned one.
104    fn bound_handle(&self) -> Option<&Handle> {
105        self.runtime
106            .as_ref()
107            .map(|r| r.handle())
108            .or(self.handle.as_ref())
109    }
110
111    fn handle(&self) -> Option<Handle> {
112        self.bound_handle()
113            .cloned()
114            .or_else(|| Handle::try_current().ok())
115    }
116
117    /// Enter the runtime this kit is bound to, if any.
118    ///
119    /// `None` is not a failure: an unbound kit runs on whichever runtime the caller is already in,
120    /// and entering that one again would be a no-op. Whether there is one at all is a separate
121    /// question, which [`has_runtime`](Self::has_runtime) answers.
122    fn enter(&self) -> Option<EnterGuard<'_>> {
123        self.bound_handle().map(Handle::enter)
124    }
125
126    /// Whether anything will be there to serve the call: our own runtime, or the caller's.
127    fn has_runtime(&self) -> bool {
128        self.bound_handle().is_some() || Handle::try_current().is_ok()
129    }
130
131    /// [`enter`](Self::enter), for the entry points which have nowhere to report a failure.
132    ///
133    /// `sleep` and `interval` capture their handle as they are constructed, so tokio would panic
134    /// from inside them with a message which does not mention this crate. Fail with ours first.
135    fn require_enter(&self) -> Option<EnterGuard<'_>> {
136        assert!(self.has_runtime(), "{NO_RUNTIME}");
137        self.enter()
138    }
139
140    fn require_handle(&self) -> Handle {
141        self.handle().expect(NO_RUNTIME)
142    }
143}
144
145impl RuntimeKit for Tokio {}
146
147impl Executor for Tokio {
148    type Task<T: Send + 'static> = TTask<T>;
149
150    fn block_on<T, F: Future<Output = T>>(&self, f: F) -> T {
151        if let Some(runtime) = self.runtime.as_ref() {
152            runtime.block_on(f)
153        } else {
154            // handle() already falls back to the ambient runtime, so there is nowhere left to
155            // look once it comes back empty.
156            self.require_handle().block_on(f)
157        }
158    }
159
160    fn spawn<T: Send + 'static, F: Future<Output = T> + Send + 'static>(
161        &self,
162        f: F,
163    ) -> Task<Self::Task<T>> {
164        TTask(Some(self.require_handle().spawn(f))).into()
165    }
166
167    fn spawn_blocking<T: Send + 'static, F: FnOnce() -> T + Send + 'static>(
168        &self,
169        f: F,
170    ) -> Task<Self::Task<T>> {
171        TTask(Some(self.require_handle().spawn_blocking(f))).into()
172    }
173}
174
175impl Reactor for Tokio {
176    type TcpStream = Compat<TcpStream>;
177    type Sleep = Sleep;
178
179    fn register<H: Read + Write + AsSysFd + Send + 'static>(
180        &self,
181        socket: H,
182    ) -> io::Result<impl AsyncRead + AsyncWrite + Send + Unpin + 'static> {
183        // AsyncFd::new reaches for the current runtime and panics when there is none. We return an
184        // io::Result, so answer the question ourselves rather than letting it unwind from in there.
185        if !self.has_runtime() {
186            return Err(io::Error::other(NO_RUNTIME));
187        }
188        let _enter = self.enter();
189        #[cfg(unix)]
190        {
191            Ok(unix::AsyncFdWrapper(tokio::io::unix::AsyncFd::new(socket)?))
192        }
193        #[cfg(not(unix))]
194        {
195            let _ = socket;
196            Err::<crate::util::DummyIO, _>(io::Error::other(
197                "Registering FD on tokio reactor is only supported on unix",
198            ))
199        }
200    }
201
202    fn sleep(&self, dur: Duration) -> Self::Sleep {
203        let _enter = self.require_enter();
204        tokio::time::sleep(dur)
205    }
206
207    fn interval(&self, dur: Duration) -> impl Stream<Item = Instant> + Send + 'static {
208        let _enter = self.require_enter();
209        IntervalStream::new(tokio::time::interval(dur)).map(tokio::time::Instant::into_std)
210    }
211
212    fn tcp_connect_addr(
213        &self,
214        addr: SocketAddr,
215    ) -> impl Future<Output = io::Result<Self::TcpStream>> + Send + 'static {
216        // Unlike sleep and interval, which grab their handle as they are constructed, connecting
217        // only touches the reactor once the future is polled, which can be from anywhere. Carry
218        // the context along instead of entering it here, where it would be gone by then.
219        //
220        // Only the kit's own binding is resolved now, so the future binds to the kit's runtime
221        // rather than to whichever one happens to poll it later. An unbound kit has nothing to
222        // carry and falls back to the ambient runtime -- but at poll time, which is the only
223        // moment there is one to find: deciding here would condemn a future built on a plain
224        // thread even when it is later polled inside a perfectly good runtime.
225        InTokioContext::new(self.bound_handle().cloned(), async move {
226            // Our siblings panic outright when there is no runtime to be found, but this one
227            // returns an io::Result, so say so properly instead of letting the caller trip over
228            // tokio's own "there is no reactor running" panic from inside connect. Asked from in
229            // here, the question is answered under whichever context InTokioContext just entered.
230            if !crate::util::inside_tokio() {
231                return Err(io::Error::other(NO_RUNTIME));
232            }
233            let stream = TcpStream::connect(addr).await?;
234            stream.set_nodelay(true)?;
235            Ok(stream.compat())
236        })
237    }
238}
239
240/// Drives a future inside a given tokio context, so it may be polled from a foreign executor.
241///
242/// The guard is taken and released within each `poll` rather than held across await points: an
243/// `EnterGuard` is not `Send`, and keeping one in the future's state would make the whole future
244/// `!Send`.
245///
246/// Only the handle is kept, deliberately: holding the kit's `Arc<TokioRT>` to keep the runtime
247/// alive would let this future become its last owner, and dropping a tokio `Runtime` from an
248/// async context panics outright, so a connect future dropped on a worker thread of that same
249/// runtime would blow up far from the cause. Entering a runtime which has since shut down merely
250/// fails the connect, which is the better of the two.
251struct InTokioContext<F: Future> {
252    handle: Option<Handle>,
253    // Boxed to get a stable address without hand-rolling a pin projection, as util::join does.
254    fut: Pin<Box<F>>,
255}
256
257impl<F: Future> InTokioContext<F> {
258    fn new(handle: Option<Handle>, fut: F) -> Self {
259        Self {
260            handle,
261            fut: Box::pin(fut),
262        }
263    }
264}
265
266impl<F: Future> Future for InTokioContext<F> {
267    type Output = F::Output;
268
269    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
270        let this = self.get_mut();
271        let _enter = this.handle.as_ref().map(Handle::enter);
272        this.fut.as_mut().poll(cx)
273    }
274}
275
276mod task {
277    use crate::util::TaskImpl;
278    use async_trait::async_trait;
279    use std::{
280        future::Future,
281        panic,
282        pin::Pin,
283        task::{Context, Poll},
284    };
285
286    /// A tokio task
287    #[derive(Debug)]
288    pub struct TTask<T: Send + 'static>(pub(super) Option<tokio::task::JoinHandle<T>>);
289
290    #[async_trait]
291    impl<T: Send + 'static> TaskImpl for TTask<T> {
292        async fn cancel(&mut self) -> Option<T> {
293            let task = self.0.take()?;
294            task.abort();
295            task.await.ok()
296        }
297    }
298
299    impl<T: Send + 'static> Future for TTask<T> {
300        type Output = T;
301
302        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
303            let task = self
304                .0
305                .as_mut()
306                .expect("Task polled after it was canceled or completed");
307            let res = match Pin::new(task).poll(cx) {
308                Poll::Pending => return Poll::Pending,
309                Poll::Ready(res) => res,
310            };
311
312            // Drop the handle now that it has completed: polling it again would trip tokio's own
313            // "JoinHandle polled after completion" assertion.
314            self.0 = None;
315
316            match res {
317                Ok(res) => Poll::Ready(res),
318                // Our Output is T, so a failed task has no value to yield. Report it the way
319                // async-task (and thus the smol and async-global-executor backends) already does
320                // rather than stalling forever on a Pending nobody will ever wake.
321                Err(err) if err.is_panic() => panic::resume_unwind(err.into_panic()),
322                Err(err) => panic!("Task did not complete: {err}"),
323            }
324        }
325    }
326}
327
328#[cfg(unix)]
329mod unix {
330    use super::*;
331    use futures_io::{AsyncRead, AsyncWrite};
332    use std::{
333        io::{IoSlice, IoSliceMut},
334        pin::Pin,
335        task::{Context, Poll},
336    };
337    use tokio::io::unix::AsyncFd;
338
339    pub(super) struct AsyncFdWrapper<H: Read + Write + AsSysFd>(pub(super) AsyncFd<H>);
340
341    impl<H: Read + Write + AsSysFd> AsyncFdWrapper<H> {
342        fn read<F: FnOnce(&mut AsyncFd<H>) -> io::Result<usize>>(
343            mut self: Pin<&mut Self>,
344            cx: &mut Context<'_>,
345            f: F,
346        ) -> Option<Poll<io::Result<usize>>> {
347            Some(match self.0.poll_read_ready_mut(cx) {
348                Poll::Pending => Poll::Pending,
349                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
350                Poll::Ready(Ok(mut guard)) => match guard.try_io(f) {
351                    Ok(res) => Poll::Ready(res),
352                    Err(_) => return None,
353                },
354            })
355        }
356
357        fn write<R, F: FnOnce(&mut AsyncFd<H>) -> io::Result<R>>(
358            mut self: Pin<&mut Self>,
359            cx: &mut Context<'_>,
360            f: F,
361        ) -> Option<Poll<io::Result<R>>> {
362            Some(match self.0.poll_write_ready_mut(cx) {
363                Poll::Pending => Poll::Pending,
364                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
365                Poll::Ready(Ok(mut guard)) => match guard.try_io(f) {
366                    Ok(res) => Poll::Ready(res),
367                    Err(_) => return None,
368                },
369            })
370        }
371    }
372
373    impl<H: Read + Write + AsSysFd> Unpin for AsyncFdWrapper<H> {}
374
375    impl<H: Read + Write + AsSysFd> AsyncRead for AsyncFdWrapper<H> {
376        fn poll_read(
377            mut self: Pin<&mut Self>,
378            cx: &mut Context<'_>,
379            buf: &mut [u8],
380        ) -> Poll<io::Result<usize>> {
381            loop {
382                if let Some(res) = self.as_mut().read(cx, |socket| socket.get_mut().read(buf)) {
383                    return res;
384                }
385            }
386        }
387
388        fn poll_read_vectored(
389            mut self: Pin<&mut Self>,
390            cx: &mut Context<'_>,
391            bufs: &mut [IoSliceMut<'_>],
392        ) -> Poll<io::Result<usize>> {
393            loop {
394                if let Some(res) = self
395                    .as_mut()
396                    .read(cx, |socket| socket.get_mut().read_vectored(bufs))
397                {
398                    return res;
399                }
400            }
401        }
402    }
403
404    impl<H: Read + Write + AsSysFd> AsyncWrite for AsyncFdWrapper<H> {
405        fn poll_write(
406            mut self: Pin<&mut Self>,
407            cx: &mut Context<'_>,
408            buf: &[u8],
409        ) -> Poll<io::Result<usize>> {
410            loop {
411                if let Some(res) = self
412                    .as_mut()
413                    .write(cx, |socket| socket.get_mut().write(buf))
414                {
415                    return res;
416                }
417            }
418        }
419
420        fn poll_write_vectored(
421            mut self: Pin<&mut Self>,
422            cx: &mut Context<'_>,
423            bufs: &[IoSlice<'_>],
424        ) -> Poll<io::Result<usize>> {
425            loop {
426                if let Some(res) = self
427                    .as_mut()
428                    .write(cx, |socket| socket.get_mut().write_vectored(bufs))
429                {
430                    return res;
431                }
432            }
433        }
434
435        fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
436            loop {
437                if let Some(res) = self.as_mut().write(cx, |socket| socket.get_mut().flush()) {
438                    return res;
439                }
440            }
441        }
442
443        fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<futures_io::Result<()>> {
444            self.poll_flush(cx)
445        }
446    }
447}
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452
453    #[test]
454    fn auto_traits() {
455        use crate::util::test::*;
456        let runtime = Runtime::tokio().unwrap();
457        assert_send(&runtime);
458        assert_sync(&runtime);
459        assert_clone(&runtime);
460    }
461
462    // A failed task used to resolve to a Pending nobody would ever wake, hanging the caller
463    // forever. Both of these must now come back, panicking, in bounded time.
464    #[test]
465    fn panicking_task_does_not_hang() {
466        let res = crate::util::test::with_timeout(|| {
467            let runtime = Runtime::tokio().unwrap();
468            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
469                runtime.block_on(runtime.spawn(async { panic!("boom") }))
470            }))
471        });
472        // Down to the payload: asserting only that something panicked would also pass if the
473        // panic were a fresh one of our own rather than the task's, resumed.
474        assert_eq!(
475            res.expect_err("task panic").downcast_ref::<&str>(),
476            Some(&"boom")
477        );
478    }
479
480    // The returned future must carry its tokio context with it: RuntimeParts pairs this reactor
481    // with a foreign executor, which polls it with no tokio runtime in scope.
482    #[test]
483    fn tcp_connect_addr_polled_off_runtime() {
484        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
485        let addr = listener.local_addr().unwrap();
486
487        // The runtime has to outlive the connect: the stream it hands back stays registered on it.
488        let (_runtime, mut stream) = crate::util::test::with_timeout(move || {
489            let runtime = Runtime::tokio().unwrap();
490            let connect = runtime.tcp_connect_addr(addr);
491            let stream = crate::util::simple_block_on(connect).expect("connect");
492            (runtime, stream)
493        });
494
495        // The listener never leaves this thread, so it is released however the test ends. Handing
496        // it to a helper to accept on would strand that helper on the very failure we guard here,
497        // holding its port for the rest of the binary.
498        let (mut socket, _) = listener.accept().expect("accept");
499        Write::write_all(&mut socket, b"hello").expect("write");
500
501        // Connecting is only half the property. The per-poll EnterGuard is long gone by now, and
502        // the stream we were handed still has to be usable off the runtime -- that is what makes
503        // entering per poll, rather than holding a guard across awaits, a safe design.
504        let read = crate::util::test::with_timeout(move || {
505            let mut buf = [0_u8; 5];
506            let mut read = 0;
507            crate::util::simple_block_on(std::future::poll_fn(|cx| {
508                while read < buf.len() {
509                    match Pin::new(&mut stream).poll_read(cx, &mut buf[read..]) {
510                        Poll::Ready(Ok(0)) => break,
511                        Poll::Ready(Ok(n)) => read += n,
512                        Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
513                        Poll::Pending => return Poll::Pending,
514                    }
515                }
516                Poll::Ready(Ok(buf))
517            }))
518            .expect("read")
519        });
520        assert_eq!(&read, b"hello");
521    }
522
523    // with_runtime records a handle too, so a kit handed each in turn used to resolve connect
524    // through one runtime and its timers and registrations through the other.
525    #[test]
526    fn one_kit_binds_everything_to_the_same_runtime() {
527        let other = TokioRT::new().unwrap();
528        let runtime = Runtime::new(
529            Tokio::default()
530                .with_runtime(TokioRT::new().unwrap())
531                .with_handle(other.handle().clone()),
532        );
533        // Nothing may be left pointing at `other` once it is gone.
534        drop(other);
535
536        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
537        let addr = listener.local_addr().unwrap();
538        let accepted = std::thread::spawn(move || listener.accept().map(|_| ()));
539        runtime.block_on(async { runtime.tcp_connect_addr(addr).await.expect("connect") });
540        accepted.join().expect("accept thread").expect("accept");
541    }
542
543    // The connect path returns an io::Result, so a missing runtime is reportable rather than a
544    // panic thrown from inside tokio once someone gets around to polling the future.
545    #[test]
546    fn tcp_connect_addr_without_a_runtime_reports_an_error() {
547        let runtime = Runtime::new(Tokio::default());
548        let addr = "127.0.0.1:1".parse().unwrap();
549        let Err(err) = crate::util::simple_block_on(runtime.tcp_connect_addr(addr)) else {
550            panic!("connect succeeded without a runtime");
551        };
552        assert!(err.to_string().contains("no tokio runtime"), "{err}");
553    }
554
555    // The mirror image of the test above: an unbound kit has no runtime of its own to carry, so
556    // the ambient one has to be looked for when the future is polled rather than when it is
557    // built. Resolving it eagerly condemns this future on the spot, on a thread which never had a
558    // runtime to offer, even though the one polling it does.
559    #[test]
560    fn tcp_connect_addr_built_off_runtime_uses_the_one_polling_it() {
561        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
562        let addr = listener.local_addr().unwrap();
563        let accepted = std::thread::spawn(move || listener.accept().map(|_| ()));
564
565        // Built here, with nothing in scope, and polled by a runtime it knows nothing about.
566        let connect = Runtime::new(Tokio::default()).tcp_connect_addr(addr);
567        TokioRT::new()
568            .unwrap()
569            .block_on(connect)
570            .expect("connect polled inside a runtime");
571        accepted.join().expect("accept thread").expect("accept");
572    }
573
574    // register hands back an io::Result too, so it owes the caller the same answer connect gives
575    // rather than tokio's panic from inside AsyncFd::new.
576    #[test]
577    #[cfg(unix)]
578    fn register_without_a_runtime_reports_an_error() {
579        let runtime = Runtime::new(Tokio::default());
580        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
581        let socket = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap();
582        let Err(err) = runtime.register(socket) else {
583            panic!("register succeeded without a runtime");
584        };
585        assert!(err.to_string().contains("no tokio runtime"), "{err}");
586    }
587
588    #[test]
589    fn panicking_blocking_task_does_not_hang() {
590        let res = crate::util::test::with_timeout(|| {
591            let runtime = Runtime::tokio().unwrap();
592            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
593                runtime.block_on(runtime.spawn_blocking(|| -> u32 { panic!("boom") }))
594            }))
595        });
596        assert_eq!(
597            res.expect_err("task panic").downcast_ref::<&str>(),
598            Some(&"boom")
599        );
600    }
601}
async_rs/implementors/tokio.rs

async_rs/implementors/
tokio.rs
