rt: avoid cloning runtime::Handle in spawn (#5724)

This commit updates `tokio::spawn` to avoid having to clone
`runtime::Handle`.

tokio/src/runtime/context.rs

@@ -182,9 +182,13 @@ cfg_rt! {
         CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None)
     }
 
-    pub(crate) fn try_current() -> Result<scheduler::Handle, TryCurrentError> {
-        match CONTEXT.try_with(|ctx| ctx.handle.borrow().clone()) {
-            Ok(Some(handle)) => Ok(handle),
+    pub(crate) fn with_current<F, R>(f: F) -> Result<R, TryCurrentError>
+    where
+        F: FnOnce(&scheduler::Handle) -> R,
+    {
+
+        match CONTEXT.try_with(|ctx| ctx.handle.borrow().as_ref().map(f)) {
+            Ok(Some(ret)) => Ok(ret),
             Ok(None) => Err(TryCurrentError::new_no_context()),
             Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()),
         }

tokio/src/runtime/coop.rs

@@ -200,9 +200,9 @@ cfg_coop! {
         cfg_metrics! {
             #[inline(always)]
             fn inc_budget_forced_yield_count() {
-                if let Ok(handle) = context::try_current() {
+                let _ = context::with_current(|handle| {
                     handle.scheduler_metrics().inc_budget_forced_yield_count();
-                }
+                });
             }
         }
 

tokio/src/runtime/handle.rs

@@ -109,7 +109,9 @@ impl Handle {
     ///
     /// Contrary to `current`, this never panics
     pub fn try_current() -> Result<Self, TryCurrentError> {
-        context::try_current().map(|inner| Handle { inner })
+        context::with_current(|inner| Handle {
+            inner: inner.clone(),
+        })
     }
 
     /// Spawns a future onto the Tokio runtime.

tokio/src/runtime/scheduler/mod.rs

@@ -52,7 +52,7 @@ cfg_rt! {
     impl Handle {
         #[track_caller]
         pub(crate) fn current() -> Handle {
-            match context::try_current() {
+            match context::with_current(Clone::clone) {
                 Ok(handle) => handle,
                 Err(e) => panic!("{}", e),
             }

tokio/src/task/spawn.rs

@@ -1,4 +1,3 @@
-use crate::runtime::Handle;
 use crate::task::JoinHandle;
 
 use std::future::Future;
@@ -178,7 +177,8 @@ cfg_rt! {
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        use crate::runtime::task;
+        use crate::runtime::{context, task};
+
         #[cfg(all(
             tokio_unstable,
             tokio_taskdump,
@@ -193,7 +193,10 @@ cfg_rt! {
         let future = task::trace::Trace::root(future);
         let id = task::Id::next();
         let task = crate::util::trace::task(future, "task", name, id.as_u64());
-        let handle = Handle::current();
-        handle.inner.spawn(task, id)
+
+        match context::with_current(|handle| handle.spawn(task, id)) {
+            Ok(join_handle) => join_handle,
+            Err(e) => panic!("{}", e),
+        }
     }
 }