From 71c8f561e36035bba4938cc5b199a93beb3d71e4 Mon Sep 17 00:00:00 2001
From: Marc-Antoine Perennou <Marc-Antoine@Perennou.com>
Date: Fri, 15 Jun 2018 05:13:39 +0000
Subject: [PATCH] runtime: add block_on_all (#398)

Signed-off-by: Marc-Antoine Perennou <Marc-Antoine@Perennou.com>
---
 src/runtime/mod.rs | 25 +++++++++++++++++++++++++
 tests/runtime.rs   | 30 ++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs
index ed416fd19..38feb53f7 100644
--- a/src/runtime/mod.rs
+++ b/src/runtime/mod.rs
@@ -389,6 +389,31 @@ impl Runtime {
         rx.wait().unwrap()
     }
 
+    /// Run a future to completion on the Tokio runtime, then wait for all
+    /// background futures to complete too.
+    ///
+    /// This runs the given future on the runtime, blocking until it is
+    /// complete, waiting for background futures to complete, and yielding
+    /// its resolved result. Any tasks or timers which the future spawns
+    /// internally will be executed on the runtime and waited for completion.
+    ///
+    /// This method should not be called from an asynchrounous context.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if the executor is at capacity, if the provided
+    /// future panics, or if called within an asynchronous execution context.
+    pub fn block_on_all<F, R, E>(mut self, future: F) -> Result<R, E>
+    where
+        F: Send + 'static + Future<Item = R, Error = E>,
+        R: Send + 'static,
+        E: Send + 'static,
+    {
+        let res = self.block_on(future);
+        self.shutdown_on_idle().wait().unwrap();
+        res
+    }
+
     /// Signals the runtime to shutdown once it becomes idle.
     ///
     /// Returns a future that completes once the shutdown operation has
diff --git a/tests/runtime.rs b/tests/runtime.rs
index 7012a1783..738e4a9df 100644
--- a/tests/runtime.rs
+++ b/tests/runtime.rs
@@ -173,3 +173,33 @@ fn spawn_many() {
     runtime.shutdown_on_idle().wait().unwrap();
     assert_eq!(ITER, *cnt.lock().unwrap());
 }
+
+#[test]
+fn spawn_from_block_on_all() {
+    let cnt = Arc::new(Mutex::new(0));
+    let c = cnt.clone();
+
+    let mut runtime = Runtime::new().unwrap();
+    let msg = runtime
+        .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");
+}