From 6ba8e7621d6b7058dda8619dd2831046cc8c250e Mon Sep 17 00:00:00 2001
From: Jon Gjengset <jon@thesquareplanet.com>
Date: Wed, 11 Jul 2018 18:32:58 -0400
Subject: [PATCH] Add free block_on_all in current thread Runtime (#477)

---
 src/runtime/current_thread/mod.rs | 17 ++++++++++++++++
 tests/runtime.rs                  | 34 +++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/src/runtime/current_thread/mod.rs b/src/runtime/current_thread/mod.rs
index 30bd8f6dd..f0a927b9b 100644
--- a/src/runtime/current_thread/mod.rs
+++ b/src/runtime/current_thread/mod.rs
@@ -68,3 +68,20 @@ mod runtime;
 
 pub use self::builder::Builder;
 pub use self::runtime::{Runtime, Handle};
+
+use futures::Future;
+
+/// Run the provided future to completion using a runtime running on the current thread.
+///
+/// This first creates a new [`Runtime`], and calls [`Runtime::block_on`] with the provided future,
+/// which blocks the current thread until the provided future completes. It then calls
+/// [`Runtime::run`] to wait for any other spawned futures to resolve.
+pub fn block_on_all<F>(future: F) -> Result<F::Item, F::Error>
+where
+    F: Future,
+{
+    let mut r = Runtime::new().expect("failed to start runtime on current thread");
+    let v = r.block_on(future)?;
+    r.run().expect("failed to resolve remaining futures");
+    Ok(v)
+}
diff --git a/tests/runtime.rs b/tests/runtime.rs
index 06ada9c80..e61707c7d 100644
--- a/tests/runtime.rs
+++ b/tests/runtime.rs
@@ -65,6 +65,40 @@ fn runtime_single_threaded() {
     runtime.run().unwrap();
 }
 
+#[test]
+fn runtime_single_threaded_block_on() {
+    let _ = env_logger::init();
+
+    tokio::runtime::current_thread::block_on_all(create_client_server_future()).unwrap();
+}
+
+#[test]
+fn runtime_single_threaded_block_on_all() {
+    let cnt = Arc::new(Mutex::new(0));
+    let c = cnt.clone();
+
+    let msg = tokio::runtime::current_thread::block_on_all(lazy(move || {
+        {
+            let mut x = c.lock().unwrap();
+            *x = 1 + *x;
+        }
+
+        // Spawn!
+        tokio::spawn(lazy(move || {
+            {
+                let mut x = c.lock().unwrap();
+                *x = 1 + *x;
+            }
+            Ok::<(), ()>(())
+        }));
+
+        Ok::<_, ()>("hello")
+    })).unwrap();
+
+    assert_eq!(2, *cnt.lock().unwrap());
+    assert_eq!(msg, "hello");
+}
+
 #[test]
 fn runtime_multi_threaded() {
     let _ = env_logger::init();