Tokio current thread runtime (#3)

* tokio current thread

Signed-off-by: Yoshua Wuyts <yoshuawuyts@gmail.com>

* tokio current thread benches

Signed-off-by: Yoshua Wuyts <yoshuawuyts@gmail.com>

Author: yoshuawuyts

benches/tokio.rs

@@ -9,3 +9,7 @@ mod common;
 mod tokio {
     benchmark_suite!(runtime_tokio::Tokio);
 }
+
+mod tokio_current_thread {
+    benchmark_suite!(runtime_tokio::TokioCurrentThread);
+}

runtime-tokio/src/lib.rs

@@ -20,30 +20,32 @@ use lazy_static::lazy_static;
 use std::io;
 use std::net::SocketAddr;
 use std::pin::Pin;
+use std::sync::{mpsc, Mutex};
+use std::thread;
 
 mod tcp;
 mod udp;
 
 use tcp::{TcpListener, TcpStream};
 use udp::UdpSocket;
 
-lazy_static! {
-    static ref TOKIO_RUNTIME: tokio::runtime::Runtime = {
-        tokio::runtime::Builder::new()
-            .after_start(|| {
-                runtime_raw::set_runtime(&Tokio);
-            })
-            .build()
-            .unwrap()
-    };
-}
-
-/// The Tokio runtime.
+/// The default Tokio runtime.
 #[derive(Debug)]
 pub struct Tokio;
 
 impl runtime_raw::Runtime for Tokio {
     fn spawn_obj(&self, fut: FutureObj<'static, ()>) -> Result<(), SpawnError> {
+        lazy_static! {
+            static ref TOKIO_RUNTIME: tokio::runtime::Runtime = {
+                tokio::runtime::Builder::new()
+                    .after_start(|| {
+                        runtime_raw::set_runtime(&Tokio);
+                    })
+                    .build()
+                    .unwrap()
+            };
+        }
+
         TOKIO_RUNTIME
             .executor()
             .spawn(Compat::new(fut.map(|_| Ok(()))));
@@ -81,3 +83,70 @@ impl runtime_raw::Runtime for Tokio {
         Ok(Box::pin(UdpSocket { tokio_socket }))
     }
 }
+
+/// The single-threaded Tokio runtime based on `tokio-current-thread`.
+#[derive(Debug)]
+pub struct TokioCurrentThread;
+
+impl runtime_raw::Runtime for TokioCurrentThread {
+    fn spawn_obj(&self, fut: FutureObj<'static, ()>) -> Result<(), SpawnError> {
+        lazy_static! {
+            static ref TOKIO_RUNTIME: Mutex<tokio::runtime::current_thread::Handle> = {
+                let (tx, rx) = mpsc::channel();
+
+                thread::spawn(move || {
+                    let mut rt = tokio::runtime::current_thread::Runtime::new().unwrap();
+                    let handle = rt.handle();
+                    tx.send(handle).unwrap();
+
+                    runtime_raw::set_runtime(&TokioCurrentThread);
+                    let forever = futures01::future::poll_fn(|| {
+                        Ok::<futures01::Async<()>, ()>(futures01::Async::NotReady)
+                    });
+                    rt.block_on(forever).unwrap();
+                });
+
+                let handle = rx.recv().unwrap();
+                Mutex::new(handle)
+            };
+        }
+
+        TOKIO_RUNTIME
+            .lock()
+            .unwrap()
+            .spawn(Compat::new(fut.map(|_| Ok(()))))
+            .unwrap();
+        Ok(())
+    }
+
+    fn connect_tcp_stream(
+        &self,
+        addr: &SocketAddr,
+    ) -> Pin<Box<dyn Future<Output = io::Result<Pin<Box<dyn runtime_raw::TcpStream>>>> + Send>>
+    {
+        use futures::compat::Compat01As03;
+        use futures01::Future;
+
+        let tokio_connect = tokio::net::TcpStream::connect(addr);
+        let connect = tokio_connect.map(|tokio_stream| {
+            Box::pin(TcpStream { tokio_stream }) as Pin<Box<dyn runtime_raw::TcpStream>>
+        });
+        Box::pin(Compat01As03::new(connect))
+    }
+
+    fn bind_tcp_listener(
+        &self,
+        addr: &SocketAddr,
+    ) -> io::Result<Pin<Box<dyn runtime_raw::TcpListener>>> {
+        let tokio_listener = tokio::net::TcpListener::bind(&addr)?;
+        Ok(Box::pin(TcpListener { tokio_listener }))
+    }
+
+    fn bind_udp_socket(
+        &self,
+        addr: &SocketAddr,
+    ) -> io::Result<Pin<Box<dyn runtime_raw::UdpSocket>>> {
+        let tokio_socket = tokio::net::UdpSocket::bind(&addr)?;
+        Ok(Box::pin(UdpSocket { tokio_socket }))
+    }
+}

tests/tokio-current-thread.rs

@@ -0,0 +1,10 @@
+#![feature(async_await, await_macro, futures_api)]
+
+#[runtime::test(runtime_tokio::TokioCurrentThread)]
+async fn spawn() {
+    let handle = runtime::spawn(async {
+        println!("hello planet from Tokio current-thread");
+        42
+    });
+    assert_eq!(await!(handle), 42);
+}