From cbb4c38bd932dfaae492b39d1af701231b7b7627 Mon Sep 17 00:00:00 2001
From: Jason Orendorff <jorendorff@github.com>
Date: Tue, 15 Nov 2022 21:50:17 -0600
Subject: [PATCH] BlockingSchedule::new: Avoid calling Handle::current().

It does not always work; and is it happens, the caller has a handle already and
can just pass it in.
---
 tokio/src/runtime/blocking/pool.rs     |  2 +-
 tokio/src/runtime/blocking/schedule.rs | 30 ++++++++++++++------------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs
index 3c2c15e830d..e9f6b66e0fc 100644
--- a/tokio/src/runtime/blocking/pool.rs
+++ b/tokio/src/runtime/blocking/pool.rs
@@ -379,7 +379,7 @@ impl Spawner {
         #[cfg(not(all(tokio_unstable, feature = "tracing")))]
         let _ = name;
 
-        let (task, handle) = task::unowned(fut, BlockingSchedule::new(), id);
+        let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id);
 
         let spawned = self.spawn_task(Task::new(task, is_mandatory), rt);
         (handle, spawned)
diff --git a/tokio/src/runtime/blocking/schedule.rs b/tokio/src/runtime/blocking/schedule.rs
index d6b27a8a256..edf775be8be 100644
--- a/tokio/src/runtime/blocking/schedule.rs
+++ b/tokio/src/runtime/blocking/schedule.rs
@@ -1,6 +1,7 @@
-use crate::runtime::task::{self, Task};
 #[cfg(feature = "test-util")]
-use crate::runtime::{scheduler, Handle};
+use crate::runtime::scheduler;
+use crate::runtime::task::{self, Task};
+use crate::runtime::Handle;
 
 /// `task::Schedule` implementation that does nothing (except some bookkeeping
 /// in test-util builds). This is unique to the blocking scheduler as tasks
@@ -14,20 +15,21 @@ pub(crate) struct BlockingSchedule {
 }
 
 impl BlockingSchedule {
-    pub(crate) fn new() -> Self {
+    #[cfg_attr(not(feature = "test-util"), allow(unused_variables))]
+    pub(crate) fn new(handle: &Handle) -> Self {
+        #[cfg(feature = "test-util")]
+        {
+            match &handle.inner {
+                scheduler::Handle::CurrentThread(handle) => {
+                    handle.driver.clock.inhibit_auto_advance();
+                }
+                #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
+                scheduler::Handle::MultiThread(_) => {}
+            }
+        }
         BlockingSchedule {
             #[cfg(feature = "test-util")]
-            handle: {
-                let handle = Handle::current();
-                match &handle.inner {
-                    scheduler::Handle::CurrentThread(handle) => {
-                        handle.driver.clock.inhibit_auto_advance();
-                    }
-                    #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
-                    scheduler::Handle::MultiThread(_) => {}
-                }
-                handle
-            },
+            handle: handle.clone(),
         }
     }
 }