From 7fbfa9b649b16de6096eb673f8debfb900618987 Mon Sep 17 00:00:00 2001
From: Marc-Antoine Perennou <Marc-Antoine@Perennou.com>
Date: Wed, 21 Oct 2020 20:00:48 +0200
Subject: tokio: deduplicate spawn_blocking (#3017)

Move common code and tracing integration into Handle

Fixes #2998
Closes #3004

Signed-off-by: Marc-Antoine Perennou <Marc-Antoine@Perennou.com>
---
 tokio/src/runtime/blocking/pool.rs |  5 +----
 tokio/src/runtime/handle.rs        | 25 +++++++++++++++++++++++++
 tokio/src/runtime/mod.rs           |  5 +----
 tokio/src/task/blocking.rs         | 13 -------------
 4 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs
index d0f2c1c8..2967a109 100644
--- a/tokio/src/runtime/blocking/pool.rs
+++ b/tokio/src/runtime/blocking/pool.rs
@@ -72,10 +72,7 @@ where
     F: FnOnce() -> R + Send + 'static,
 {
     let rt = context::current().expect("not currently running on the Tokio runtime.");
-
-    let (task, handle) = task::joinable(BlockingTask::new(func));
-    let _ = rt.blocking_spawner.spawn(task, &rt);
-    handle
+    rt.spawn_blocking(func)
 }
 
 #[allow(dead_code)]
diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs
index 9c2cfa5f..b1e8d8f1 100644
--- a/tokio/src/runtime/handle.rs
+++ b/tokio/src/runtime/handle.rs
@@ -1,3 +1,5 @@
+use crate::runtime::blocking::task::BlockingTask;
+use crate::runtime::task::{self, JoinHandle};
 use crate::runtime::{blocking, driver, Spawner};
 
 /// Handle to the runtime.
@@ -36,4 +38,27 @@ impl Handle {
     // {
     //     context::enter(self.clone(), f)
     // }
+
+    /// Run the provided function on an executor dedicated to blocking operations.
+    pub(crate) fn spawn_blocking<F, R>(&self, func: F) -> JoinHandle<R>
+    where
+        F: FnOnce() -> R + Send + 'static,
+    {
+        #[cfg(feature = "tracing")]
+        let func = {
+            let span = tracing::trace_span!(
+                target: "tokio::task",
+                "task",
+                kind = %"blocking",
+                function = %std::any::type_name::<F>(),
+            );
+            move || {
+                let _g = span.enter();
+                func()
+            }
+        };
+        let (task, handle) = task::joinable(BlockingTask::new(func));
+        let _ = self.blocking_spawner.spawn(task, &self);
+        handle
+    }
 }
diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs
index 4eba7923..be4aa38b 100644
--- a/tokio/src/runtime/mod.rs
+++ b/tokio/src/runtime/mod.rs
@@ -187,7 +187,6 @@ cfg_rt! {
 
     mod blocking;
     use blocking::BlockingPool;
-    use blocking::task::BlockingTask;
     pub(crate) use blocking::spawn_blocking;
 
     mod builder;
@@ -390,9 +389,7 @@ cfg_rt! {
         where
             F: FnOnce() -> R + Send + 'static,
         {
-            let (task, handle) = task::joinable(BlockingTask::new(func));
-            let _ = self.handle.blocking_spawner.spawn(task, &self.handle);
-            handle
+            self.handle.spawn_blocking(func)
         }
 
         /// Run a future to completion on the Tokio runtime. This is the
diff --git a/tokio/src/task/blocking.rs b/tokio/src/task/blocking.rs
index 88ba678b..fc6632be 100644
--- a/tokio/src/task/blocking.rs
+++ b/tokio/src/task/blocking.rs
@@ -109,18 +109,5 @@ where
     F: FnOnce() -> R + Send + 'static,
     R: Send + 'static,
 {
-    #[cfg(feature = "tracing")]
-    let f = {
-        let span = tracing::trace_span!(
-            target: "tokio::task",
-            "task",
-            kind = %"blocking",
-            function = %std::any::type_name::<F>(),
-        );
-        move || {
-            let _g = span.enter();
-            f()
-        }
-    };
     crate::runtime::spawn_blocking(f)
 }
-- 
cgit v1.2.3