feat: get runtime handle in the creation of timeout

tokio/src/time/sleep.rs

@@ -256,6 +256,16 @@ impl Sleep {
         use crate::runtime::scheduler;
 
         let handle = scheduler::Handle::current();
+        Self::new_timeout_with_handle(deadline, location, handle)
+    }
+
+    #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_variables))]
+    #[track_caller]
+    pub(crate) fn new_timeout_with_handle(
+        deadline: Instant,
+        location: Option<&'static Location<'static>>,
+        handle: crate::runtime::scheduler::Handle,
+    ) -> Sleep {
         let entry = TimerEntry::new(&handle, deadline);
 
         #[cfg(all(tokio_unstable, feature = "tracing"))]

tokio/src/time/timeout.rs

@@ -139,10 +139,14 @@ pub fn timeout_at<F>(deadline: Instant, future: F) -> Timeout<F>
 where
     F: Future,
 {
+    use crate::runtime::scheduler;
+    let handle = scheduler::Handle::current();
+
     Timeout {
         value: future,
         deadline: Some(deadline),
         delay: None,
+        handle,
     }
 }
 
@@ -153,18 +157,23 @@ pin_project! {
     pub struct Timeout<T> {
         #[pin]
         value: T,
-        deadline : Option<Instant>,
         #[pin]
         delay: Option<Sleep>,
+        deadline : Option<Instant>,
+        handle: crate::runtime::scheduler::Handle,
     }
 }
 
 impl<T> Timeout<T> {
     pub(crate) fn new_with_delay(value: T, deadline: Option<Instant>) -> Timeout<T> {
+        use crate::runtime::scheduler;
+        let handle = scheduler::Handle::current();
+
         Timeout {
             value,
             deadline,
             delay: None,
+            handle,
         }
     }
 
@@ -208,7 +217,9 @@ where
         if me.delay.is_none() {
             let location = trace::caller_location();
             let delay = match me.deadline {
-                Some(deadline) => Sleep::new_timeout(*deadline, location),
+                Some(deadline) => {
+                    Sleep::new_timeout_with_handle(*deadline, location, me.handle.clone())
+                }
                 None => Sleep::far_future(location),
             };
             me.delay.as_mut().set(Some(delay));