tokio: allow binding to a runtime directly

Author: Keruspe

README.md

@@ -25,29 +25,26 @@
 
 ```rust
 use async_rs::{Runtime, TokioRuntime, traits::*};
-use std::{io, sync::Arc, time::Duration};
+use std::{io, time::Duration};
 
-async fn get_a(rt: Arc<TokioRuntime>) -> io::Result<u32> {
-    rt.clone()
-        .spawn_blocking(move || rt.block_on(async { Ok(12) }))
-        .await
+async fn get_a(rt: &TokioRuntime) -> io::Result<u32> {
+    rt.spawn_blocking(|| Ok(12)).await
 }
 
-async fn get_b(rt: Arc<TokioRuntime>) -> io::Result<u32> {
+async fn get_b(rt: &TokioRuntime) -> io::Result<u32> {
     rt.spawn(async { Ok(30) }).await
 }
 
-async fn tokio_main() -> io::Result<()> {
-    let rt = Arc::new(Runtime::tokio());
-    let a = get_a(rt.clone()).await?;
-    let b = get_b(rt.clone()).await?;
+async fn tokio_main(rt: &TokioRuntime) -> io::Result<()> {
+    let a = get_a(&rt).await?;
+    let b = get_b(&rt).await?;
     rt.sleep(Duration::from_millis(500)).await;
     assert_eq!(a + b, 42);
     Ok(())
 }
 
-#[tokio::main]
-async fn main() -> io::Result<()> {
-    tokio_main().await
+fn main() -> io::Result<()> {
+    let rt = Runtime::tokio()?;
+    rt.block_on(tokio_main(&rt))
 }
 ```

examples/send-recv.rs

@@ -28,8 +28,7 @@ fn send(mut stream: impl AsyncRead + AsyncWrite + Unpin) -> io::Result<()> {
     }
 }
 
-async fn tokio_main() -> io::Result<()> {
-    let rt = Runtime::tokio();
+async fn tokio_main(rt: &TokioRuntime) -> io::Result<()> {
     let listener = listener(&rt).await?;
     let sender = sender(&rt).await?;
     let stream = rt
@@ -47,7 +46,7 @@ async fn tokio_main() -> io::Result<()> {
     Ok(())
 }
 
-#[tokio::main]
-async fn main() -> io::Result<()> {
-    tokio_main().await
+fn main() -> io::Result<()> {
+    let rt = Runtime::tokio()?;
+    rt.block_on(tokio_main(&rt))
 }

examples/tokio.rs

@@ -1,31 +1,29 @@
 use async_rs::{Runtime, TokioRuntime, traits::*};
-use std::{io, sync::Arc, time::Duration};
+use std::{io, time::Duration};
 
-async fn get_a(rt: Arc<TokioRuntime>) -> io::Result<u32> {
-    rt.clone()
-        .spawn_blocking(move || rt.block_on(async { Ok(12) }))
-        .await
+async fn get_a(rt: &TokioRuntime) -> io::Result<u32> {
+    rt.spawn_blocking(|| Ok(12)).await
 }
 
-async fn get_b(rt: Arc<TokioRuntime>) -> io::Result<u32> {
+async fn get_b(rt: &TokioRuntime) -> io::Result<u32> {
     rt.spawn(async { Ok(30) }).await
 }
 
-async fn tokio_main() -> io::Result<()> {
-    let rt = Arc::new(Runtime::tokio());
-    let a = get_a(rt.clone()).await?;
-    let b = get_b(rt.clone()).await?;
+async fn tokio_main(rt: &TokioRuntime) -> io::Result<()> {
+    let a = get_a(&rt).await?;
+    let b = get_b(&rt).await?;
     rt.sleep(Duration::from_millis(500)).await;
     assert_eq!(a + b, 42);
     Ok(())
 }
 
-#[tokio::main]
-async fn main() -> io::Result<()> {
-    tokio_main().await
+fn main() -> io::Result<()> {
+    let rt = Runtime::tokio()?;
+    rt.block_on(tokio_main(&rt))
 }
 
-#[tokio::test]
-async fn tokio() -> io::Result<()> {
-    tokio_main().await
+#[test]
+fn tokio() -> io::Result<()> {
+    let rt = Runtime::tokio()?;
+    rt.block_on(tokio_main(&rt))
 }

src/implementors/tokio.rs

@@ -17,29 +17,43 @@ use std::{
     task::{Context, Poll},
     time::{Duration, Instant},
 };
-use tokio::{net::TcpStream, runtime::Handle};
+use tokio::{
+    net::TcpStream,
+    runtime::{EnterGuard, Handle, Runtime as TokioRT},
+};
 use tokio_stream::{StreamExt, wrappers::IntervalStream};
 use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
 
 /// Type alias for the tokio runtime
 pub type TokioRuntime = Runtime<Tokio>;
 
 impl TokioRuntime {
+    /// Create a new TokioRuntime and bind it to this tokio runtime.
+    pub fn tokio() -> io::Result<Self> {
+        Ok(Self::tokio_with_runtime(TokioRT::new()?))
+    }
+
     /// Create a new TokioRuntime and bind it to the current tokio runtime by default.
-    pub fn tokio() -> Self {
+    pub fn tokio_current() -> Self {
         Self::new(Tokio::current())
     }
 
     /// Create a new TokioRuntime and bind it to the tokio runtime associated to this handle by default.
     pub fn tokio_with_handle(handle: Handle) -> Self {
         Self::new(Tokio::default().with_handle(handle))
     }
+
+    /// Create a new TokioRuntime and bind it to this tokio runtime.
+    pub fn tokio_with_runtime(runtime: TokioRT) -> Self {
+        Self::new(Tokio::default().with_runtime(runtime))
+    }
 }
 
 /// Dummy object implementing async common interfaces on top of tokio
-#[derive(Default, Debug, Clone)]
+#[derive(Default, Debug)]
 pub struct Tokio {
     handle: Option<Handle>,
+    runtime: Option<TokioRT>,
 }
 
 impl Tokio {
@@ -49,13 +63,32 @@ impl Tokio {
         self
     }
 
+    /// Bind to the tokio Runtime associated to this handle by default.
+    pub fn with_runtime(mut self, runtime: TokioRT) -> Self {
+        let handle = runtime.handle().clone();
+        self.runtime = Some(runtime);
+        self.with_handle(handle)
+    }
+
     /// Bind to the current tokio Runtime by default.
     pub fn current() -> Self {
         Self::default().with_handle(Handle::current())
     }
 
-    pub(crate) fn handle(&self) -> Option<Handle> {
-        Handle::try_current().ok().or_else(|| self.handle.clone())
+    fn handle(&self) -> Option<Handle> {
+        self.runtime
+            .as_ref()
+            .map(|r| r.handle().clone())
+            .or_else(|| Handle::try_current().ok())
+            .or_else(|| self.handle.clone())
+    }
+
+    fn enter(&self) -> Option<EnterGuard<'_>> {
+        self.runtime
+            .as_ref()
+            .map(TokioRT::handle)
+            .or(self.handle.as_ref())
+            .map(Handle::enter)
     }
 }
 
@@ -65,7 +98,9 @@ impl RuntimeKit for Tokio {}
 
 impl Executor for Tokio {
     fn block_on<T, F: Future<Output = T>>(&self, f: F) -> T {
-        if let Some(handle) = self.handle() {
+        if let Some(runtime) = self.runtime.as_ref() {
+            runtime.block_on(f)
+        } else if let Some(handle) = self.handle() {
             handle.block_on(f)
         } else {
             Handle::current().block_on(f)
@@ -123,7 +158,7 @@ impl Reactor for Tokio {
         &self,
         socket: H,
     ) -> io::Result<impl AsyncRead + AsyncWrite + Send + Unpin + 'static> {
-        let _enter = self.handle().as_ref().map(|handle| handle.enter());
+        let _enter = self.enter();
         cfg_if! {
             if #[cfg(unix)] {
                 Ok(unix::AsyncFdWrapper(
@@ -138,19 +173,20 @@ impl Reactor for Tokio {
     }
 
     fn sleep(&self, dur: Duration) -> impl Future<Output = ()> + Send + 'static {
+        let _enter = self.enter();
         tokio::time::sleep(dur)
     }
 
     fn interval(&self, dur: Duration) -> impl Stream<Item = Instant> + Send + 'static {
-        let _enter = self.handle().as_ref().map(|handle| handle.enter());
+        let _enter = self.enter();
         IntervalStream::new(tokio::time::interval(dur)).map(tokio::time::Instant::into_std)
     }
 
     fn tcp_connect_addr(
         &self,
         addr: SocketAddr,
     ) -> impl Future<Output = io::Result<Self::TcpStream>> + Send + 'static {
-        let _enter = self.handle().as_ref().map(|handle| handle.enter());
+        let _enter = self.enter();
         async move { Ok(TcpStream::connect(addr).await?.compat()) }
     }
 }

src/lib.rs

@@ -14,30 +14,27 @@
 //!
 //! ```rust
 //! use async_rs::{Runtime, TokioRuntime, traits::*};
-//! use std::{io, sync::Arc, time::Duration};
+//! use std::{io, time::Duration};
 //!
-//! async fn get_a(rt: Arc<TokioRuntime>) -> io::Result<u32> {
-//!     rt.clone()
-//!         .spawn_blocking(move || rt.block_on(async { Ok(12) }))
-//!         .await
+//! async fn get_a(rt: &TokioRuntime) -> io::Result<u32> {
+//!     rt.spawn_blocking(|| Ok(12)).await
 //! }
 //!
-//! async fn get_b(rt: Arc<TokioRuntime>) -> io::Result<u32> {
+//! async fn get_b(rt: &TokioRuntime) -> io::Result<u32> {
 //!     rt.spawn(async { Ok(30) }).await
 //! }
 //!
-//! async fn tokio_main() -> io::Result<()> {
-//!     let rt = Arc::new(Runtime::tokio());
-//!     let a = get_a(rt.clone()).await?;
-//!     let b = get_b(rt.clone()).await?;
+//! async fn tokio_main(rt: &TokioRuntime) -> io::Result<()> {
+//!     let a = get_a(&rt).await?;
+//!     let b = get_b(&rt).await?;
 //!     rt.sleep(Duration::from_millis(500)).await;
 //!     assert_eq!(a + b, 42);
 //!     Ok(())
 //! }
 //!
-//! #[tokio::main]
-//! async fn main() -> io::Result<()> {
-//!     tokio_main().await
+//! fn main() -> io::Result<()> {
+//!     let rt = Runtime::tokio()?;
+//!     rt.block_on(tokio_main(&rt))
 //! }
 //! ```