use ad-hoc runtime in tokio on last resort

It closes the gap between tokio and async-std, bringing similar functionality to tokio.

Author: ftelnov

src/spawner.rs

@@ -53,20 +53,69 @@ pub mod use_async_std {
 #[cfg(feature = "use-tokio")]
 pub mod use_tokio {
     use super::*;
-    use tokio::{runtime::Handle, task as tokio_task};
+    use tokio::{
+        runtime::{Handle, Runtime},
+        task::{self as tokio_task, block_in_place},
+    };
+
+    pub struct TokioSpawner(Option<TokioRuntime>);
+
+    impl Clone for TokioSpawner {
+        fn clone(&self) -> Self {
+            Self(self.0.as_ref().map(|rt| match rt {
+                TokioRuntime::ByHandle(handle) => TokioRuntime::ByHandle(handle.clone()),
+                TokioRuntime::Owned(runtime) => TokioRuntime::ByHandle(runtime.handle().clone()),
+            }))
+        }
+    }
+
+    const RUNTIME_INVARIANT_ERR: &str =
+        "invariant: runtime must be available during the spawner's lifetime";
 
-    pub struct TokioSpawner(Handle);
+    impl Drop for TokioSpawner {
+        /// Graceful shutdown owned runtime.
+        fn drop(&mut self) {
+            if let TokioRuntime::Owned(rt) = self.0.take().expect(RUNTIME_INVARIANT_ERR) {
+                rt.shutdown_background()
+            }
+        }
+    }
 
     impl TokioSpawner {
         pub fn new(rt_handle: Handle) -> Self {
-            Self(rt_handle)
+            Self(Some(TokioRuntime::ByHandle(rt_handle)))
         }
+
+        fn handle(&self) -> &Handle {
+            match &self.0.as_ref().expect(RUNTIME_INVARIANT_ERR) {
+                TokioRuntime::ByHandle(handle) => handle,
+                TokioRuntime::Owned(runtime) => runtime.handle(),
+            }
+        }
+    }
+
+    /// Variants of supplied tokio runtime.
+    /// Is needed because runtime can be either passed or created.
+    enum TokioRuntime {
+        /// User provides its own runtime, we'll refer to it by handle.
+        ByHandle(Handle),
+        /// We've created our own ad-hoc runtime, so we'll own it.
+        Owned(Runtime),
     }
 
-    // By default, `TokioSpawner` operates on global runtime.
+    // By default, `TokioSpawner` operates on globally available runtime.
+    // Ad-hoc runtime would only be created if it is not available globally.
+    // Newly created runtime would be destroyed when spawner is gone.
+    // That is needed because running runtime would prevent program from stopping.
     impl Default for TokioSpawner {
         fn default() -> Self {
-            Self(Handle::current())
+            if let Ok(handle) = Handle::try_current() {
+                return Self(Some(TokioRuntime::ByHandle(handle)));
+            }
+            let runtime = tokio::runtime::Builder::new_current_thread()
+                .build()
+                .unwrap();
+            Self(Some(TokioRuntime::Owned(runtime)))
         }
     }
 
@@ -75,7 +124,7 @@ pub mod use_tokio {
         type SpawnHandle = tokio_task::JoinHandle<T>;
 
         fn spawn<F: Future<Output = T> + Send + 'static>(&self, f: F) -> Self::SpawnHandle {
-            self.0.spawn(f)
+            self.handle().spawn(f)
         }
     }
 
@@ -84,13 +133,20 @@ pub mod use_tokio {
         type SpawnHandle = tokio_task::JoinHandle<T>;
 
         fn spawn_func<F: FnOnce() -> T + Send + 'static>(&self, f: F) -> Self::SpawnHandle {
-            self.0.spawn_blocking(f)
+            self.handle().spawn_blocking(f)
         }
     }
 
     unsafe impl Blocker for TokioSpawner {
         fn block_on<T, F: Future<Output = T>>(&self, f: F) -> T {
-            tokio_task::block_in_place(|| self.0.block_on(f))
+            let result = block_in_place(|| match self.0.as_ref().expect(RUNTIME_INVARIANT_ERR) {
+                TokioRuntime::ByHandle(handle) => handle.block_on(f),
+                // if runtime is owned, `block_on` must be called directly on it,
+                // not via it's handle. Otherwise, future won't be able to run IO-tasks.
+                // see `block_on` docs for more info.
+                TokioRuntime::Owned(runtime) => runtime.block_on(f),
+            });
+            result
         }
     }
 }