orb_tokio/
lib.rs

1//! # Orb Tokio Runtime
2//!
3//! This crate provides a Tokio-based implementation of the Orb async runtime traits.
4//! It allows users to leverage Tokio's powerful async runtime with the unified Orb interface.
5//!
6//! The main type provided is [`TokioRT`], which implements the core runtime functionality.
7//!
8//! See the [main Orb documentation](https://github.com/NaturalIO/orb) for more information.
9//!
10//! ## Usage
11//!
12//! ```rust
13//! use orb_tokio::TokioRT;
14//! use orb::prelude::*;
15//!
16//! type RT = TokioRT;
17//!
18//! let rt = RT::multi(4);
19//! ```
20
21pub use orb::AsyncRuntime;
22use orb::io::{AsyncFd, AsyncIO};
23pub use orb::runtime::{AsyncExec, AsyncJoiner, ThreadJoiner};
24use orb::time::{AsyncTime, TimeInterval};
25use std::fmt;
26use std::future::Future;
27use std::io;
28use std::net::{SocketAddr, TcpStream};
29use std::ops::Deref;
30use std::os::fd::{AsFd, AsRawFd};
31use std::os::unix::net::UnixStream;
32use std::path::Path;
33use std::pin::Pin;
34use std::sync::Arc;
35use std::task::*;
36use std::time::{Duration, Instant};
37use tokio::runtime::{Builder, Handle, Runtime};
38
39pub struct TokioRT {}
40
41impl AsyncIO for TokioRT {
42    type AsyncFd<T: AsRawFd + AsFd + Send + Sync + 'static> = TokioFD<T>;
43
44    #[inline(always)]
45    async fn connect_tcp(addr: &SocketAddr) -> io::Result<Self::AsyncFd<TcpStream>> {
46        let stream = tokio::net::TcpStream::connect(addr).await?;
47        // into_std will not change back to blocking
48        Self::to_async_fd_rw(stream.into_std()?)
49    }
50
51    #[inline(always)]
52    async fn connect_unix(addr: &Path) -> io::Result<Self::AsyncFd<UnixStream>> {
53        let stream = tokio::net::UnixStream::connect(addr).await?;
54        // into_std will not change back to blocking
55        Self::to_async_fd_rw(stream.into_std()?)
56    }
57
58    #[inline(always)]
59    fn to_async_fd_rd<T: AsRawFd + AsFd + Send + Sync + 'static>(
60        fd: T,
61    ) -> io::Result<Self::AsyncFd<T>> {
62        use tokio::io;
63        Ok(TokioFD(io::unix::AsyncFd::with_interest(fd, io::Interest::READABLE)?))
64    }
65
66    #[inline(always)]
67    fn to_async_fd_rw<T: AsRawFd + AsFd + Send + Sync + 'static>(
68        fd: T,
69    ) -> io::Result<Self::AsyncFd<T>> {
70        use tokio::io;
71        use tokio::io::Interest;
72        Ok(TokioFD(io::unix::AsyncFd::with_interest(fd, Interest::READABLE | Interest::WRITABLE)?))
73    }
74}
75
76impl AsyncTime for TokioRT {
77    type Interval = TokioInterval;
78
79    #[inline(always)]
80    fn sleep(d: Duration) -> impl Future + Send {
81        tokio::time::sleep(d)
82    }
83
84    #[inline(always)]
85    fn interval(d: Duration) -> Self::Interval {
86        let later = tokio::time::Instant::now() + d;
87        TokioInterval(tokio::time::interval_at(later, d))
88    }
89}
90
91impl AsyncRuntime for TokioRT {
92    type Exec = TokioExec;
93
94    /// Initiate executor using current thread.
95    ///
96    /// # Safety
97    ///
98    /// You should run [AsyncExec::block_on()] with this executor.
99    ///
100    /// If spawn without a `block_on()` running, it's possible
101    /// the runtime just init future without scheduling.
102    fn current() -> Self::Exec {
103        TokioExec::new_current_thread()
104    }
105
106    /// Initiate executor with one background thread.
107    ///
108    /// # NOTE
109    ///
110    /// [AsyncExec::block_on()] is optional, you can directly call [AsyncExec::spawn] with it.
111    #[inline(always)]
112    fn one() -> Self::Exec {
113        TokioExec::new_multi_thread(1)
114    }
115
116    /// Initiate executor with multiple background threads.
117    ///
118    /// # NOTE
119    ///
120    /// When `num` == 0, start threads that match cpu number.
121    ///
122    /// [AsyncExec::block_on()] is optional, you can directly call [AsyncExec::spawn] with it.
123    #[inline(always)]
124    fn multi(num: usize) -> Self::Exec {
125        TokioExec::new_multi_thread(num)
126    }
127
128    /// Spawn a task in the background, returning a handle to await its result
129    #[inline]
130    fn spawn<F, R>(f: F) -> TokioJoinHandle<R>
131    where
132        F: Future<Output = R> + Send + 'static,
133        R: Send + 'static,
134    {
135        // Although AsyncJoiner don't need Send marker, but here in the spawn()
136        // need to restrict the requirements
137        return TokioJoinHandle(tokio::spawn(f));
138    }
139
140    /// Spawn a task and detach it (no handle returned)
141    #[inline]
142    fn spawn_detach<F, R>(f: F)
143    where
144        F: Future<Output = R> + Send + 'static,
145        R: Send + 'static,
146    {
147        tokio::spawn(f);
148    }
149
150    #[inline(always)]
151    fn spawn_blocking<F, R>(f: F) -> TokioThreadHandle<R>
152    where
153        F: FnOnce() -> R + Send + 'static,
154        R: Send + 'static,
155    {
156        TokioThreadHandle(tokio::task::spawn_blocking(f))
157    }
158}
159
160/// Associate type for TokioRT
161pub struct TokioInterval(tokio::time::Interval);
162
163impl TimeInterval for TokioInterval {
164    #[inline]
165    fn poll_tick(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Instant> {
166        let _self = self.get_mut();
167        if let Poll::Ready(i) = _self.0.poll_tick(ctx) {
168            Poll::Ready(i.into_std())
169        } else {
170            Poll::Pending
171        }
172    }
173}
174
175/// Associate type for TokioRT
176pub struct TokioFD<T: AsRawFd + AsFd + Send + Sync + 'static>(tokio::io::unix::AsyncFd<T>);
177
178impl<T: AsRawFd + AsFd + Send + Sync + 'static> AsyncFd<T> for TokioFD<T> {
179    #[inline(always)]
180    async fn async_read<R>(&self, f: impl FnMut(&T) -> io::Result<R> + Send) -> io::Result<R> {
181        self.0.async_io(tokio::io::Interest::READABLE, f).await
182    }
183
184    #[inline(always)]
185    async fn async_write<R>(&self, f: impl FnMut(&T) -> io::Result<R> + Send) -> io::Result<R> {
186        self.0.async_io(tokio::io::Interest::WRITABLE, f).await
187    }
188}
189
190impl<T: AsRawFd + AsFd + Send + Sync + 'static> Deref for TokioFD<T> {
191    type Target = T;
192
193    #[inline(always)]
194    fn deref(&self) -> &Self::Target {
195        self.0.get_ref()
196    }
197}
198
199/// A wrapper around tokio's JoinHandle that implements AsyncJoiner
200pub struct TokioJoinHandle<T>(tokio::task::JoinHandle<T>);
201
202impl<T: Send> AsyncJoiner<T> for TokioJoinHandle<T> {
203    #[inline]
204    fn is_finished(&self) -> bool {
205        self.0.is_finished()
206    }
207
208    #[inline]
209    fn detach(self) {
210        // Tokio's JoinHandle doesn't need explicit detach, it will run in background
211        // when the handle is dropped
212    }
213
214    #[inline]
215    fn abort(self) {
216        self.0.abort();
217    }
218
219    #[inline(always)]
220    fn abort_boxed(self: Box<Self>) {
221        self.0.abort();
222    }
223
224    #[inline(always)]
225    fn detach_boxed(self: Box<Self>) {
226        // Tokio's JoinHandle doesn't need explicit detach, it will run in background
227        // when the handle is dropped
228    }
229}
230
231impl<T> Future for TokioJoinHandle<T> {
232    type Output = Result<T, ()>;
233
234    #[inline]
235    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
236        let _self = unsafe { self.get_unchecked_mut() };
237        if let Poll::Ready(r) = Pin::new(&mut _self.0).poll(cx) {
238            return Poll::Ready(r.map_err(|_e| ()));
239        }
240        Poll::Pending
241    }
242}
243
244/// A wrapper around tokio's JoinHandle that implements ThreadJoiner
245pub struct TokioThreadHandle<T>(tokio::task::JoinHandle<T>);
246
247impl<T> ThreadJoiner<T> for TokioThreadHandle<T> {
248    #[inline]
249    fn is_finished(&self) -> bool {
250        self.0.is_finished()
251    }
252}
253
254impl<T> Future for TokioThreadHandle<T> {
255    type Output = Result<T, ()>;
256
257    #[inline]
258    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
259        let _self = unsafe { self.get_unchecked_mut() };
260        if let Poll::Ready(r) = Pin::new(&mut _self.0).poll(cx) {
261            return Poll::Ready(r.map_err(|_e| ()));
262        }
263        Poll::Pending
264    }
265}
266
267/// The main struct for tokio runtime IO, assign this type to AsyncIO trait when used.
268pub enum TokioExec {
269    // Runtime don't have clone, since we don't have thread context, we need to put runtime to Arc
270    // to impl Clone. (We usually need to clone self before calling block_on)
271    Runtime(Arc<Runtime>),
272    Handle(Handle),
273}
274
275impl Clone for TokioExec {
276    /// Clone a TokioRT::Handle out of runtime, for spawn
277    fn clone(&self) -> Self {
278        match self {
279            Self::Handle(h) => {
280                return Self::Handle(h.clone());
281            }
282            Self::Runtime(rt) => Self::Runtime(rt.clone()),
283        }
284    }
285}
286
287impl TokioExec {
288    /// Capture a runtime
289    #[inline]
290    pub fn new_with_runtime(rt: Runtime) -> Self {
291        Self::Runtime(Arc::new(rt))
292    }
293
294    #[inline]
295    pub fn new_multi_thread(workers: usize) -> Self {
296        let mut builder = Builder::new_multi_thread();
297        if workers > 0 {
298            builder.worker_threads(workers);
299        }
300        let rt = builder.enable_all().build().unwrap();
301        Self::Runtime(Arc::new(rt))
302    }
303
304    #[inline]
305    pub fn new_current_thread() -> Self {
306        let mut builder = Builder::new_current_thread();
307        let rt = builder.enable_all().build().unwrap();
308        Self::Runtime(Arc::new(rt))
309    }
310
311    /// Only capture a runtime handle. Should acquire with
312    /// `async { Handle::current() }`
313    #[inline]
314    pub fn new_with_handle(handle: Handle) -> Self {
315        Self::Handle(handle)
316    }
317}
318
319impl AsyncExec for TokioExec {
320    type AsyncJoiner<R: Send> = TokioJoinHandle<R>;
321
322    type ThreadJoiner<R: Send> = TokioThreadHandle<R>;
323
324    /// Spawn a task in the background, returning a handle to await its result
325    #[inline]
326    fn spawn<F, R>(&self, f: F) -> TokioJoinHandle<R>
327    where
328        F: Future<Output = R> + Send + 'static,
329        R: Send + 'static,
330    {
331        // Although AsyncJoiner don't need Send marker, but here in the spawn()
332        // need to restrict the requirements
333        match self {
334            Self::Runtime(s) => {
335                return TokioJoinHandle(s.spawn(f));
336            }
337            Self::Handle(s) => {
338                return TokioJoinHandle(s.spawn(f));
339            }
340        }
341    }
342
343    /// Spawn a task and detach it (no handle returned)
344    #[inline]
345    fn spawn_detach<F, R>(&self, f: F)
346    where
347        F: Future<Output = R> + Send + 'static,
348        R: Send + 'static,
349    {
350        match self {
351            Self::Runtime(s) => {
352                s.spawn(f);
353            }
354            Self::Handle(s) => {
355                s.spawn(f);
356            }
357        }
358    }
359
360    #[inline(always)]
361    fn spawn_blocking<F, R>(&self, f: F) -> TokioThreadHandle<R>
362    where
363        F: FnOnce() -> R + Send + 'static,
364        R: Send + 'static,
365    {
366        match self {
367            Self::Runtime(s) => TokioThreadHandle(s.spawn_blocking(f)),
368            Self::Handle(s) => TokioThreadHandle(s.spawn_blocking(f)),
369        }
370    }
371
372    /// Run a future to completion on the runtime
373    #[inline]
374    fn block_on<F, R>(&self, f: F) -> R
375    where
376        F: Future<Output = R>,
377        R: 'static,
378    {
379        match self {
380            Self::Runtime(s) => {
381                return s.block_on(f);
382            }
383            Self::Handle(_s) => {
384                // panic in order to prevent misbehaved code.
385                // refer to https://docs.rs/tokio/latest/tokio/runtime/struct.Handle.html#method.block_on
386                panic!("handle is not allowed to block_on");
387            }
388        }
389    }
390}
391
392impl fmt::Debug for TokioExec {
393    #[inline]
394    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
395        match self {
396            Self::Runtime(_) => write!(f, "tokio(rt)"),
397            Self::Handle(_) => write!(f, "tokio(handle)"),
398        }
399    }
400}
orb_tokio/lib.rs

orb_tokio/
lib.rs
