From ef5873186756584ac828ddfe912f8e552281429f Mon Sep 17 00:00:00 2001
From: Zahari Dichev <zaharidichev@gmail.com>
Date: Tue, 14 Sep 2021 20:13:44 +0000
Subject: [PATCH] move track caller into func

Signed-off-by: Zahari Dichev <zaharidichev@gmail.com>
---
 tokio/src/time/driver/sleep.rs | 16 ++++------------
 tokio/src/time/timeout.rs      | 13 +++++--------
 tokio/src/util/trace.rs        |  8 ++++++++
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/tokio/src/time/driver/sleep.rs b/tokio/src/time/driver/sleep.rs
index 015a802c0fe..db1d444bf9a 100644
--- a/tokio/src/time/driver/sleep.rs
+++ b/tokio/src/time/driver/sleep.rs
@@ -1,5 +1,6 @@
 use crate::time::driver::{Handle, TimerEntry};
 use crate::time::{error::Error, Duration, Instant};
+use crate::util::trace;
 
 use pin_project_lite::pin_project;
 use std::future::Future;
@@ -46,13 +47,7 @@ cfg_trace! {
 #[cfg_attr(docsrs, doc(alias = "delay_until"))]
 #[cfg_attr(tokio_track_caller, track_caller)]
 pub fn sleep_until(deadline: Instant) -> Sleep {
-    #[cfg(tokio_track_caller)]
-    let location = std::panic::Location::caller();
-    #[cfg(tokio_track_caller)]
-    return Sleep::new_timeout(deadline, Some(location));
-
-    #[cfg(not(tokio_track_caller))]
-    Sleep::new_timeout(deadline, None)
+    return Sleep::new_timeout(deadline, trace::caller_location());
 }
 
 /// Waits until `duration` has elapsed.
@@ -96,10 +91,7 @@ pub fn sleep_until(deadline: Instant) -> Sleep {
 #[cfg_attr(docsrs, doc(alias = "wait"))]
 #[cfg_attr(tokio_track_caller, track_caller)]
 pub fn sleep(duration: Duration) -> Sleep {
-    #[cfg(tokio_track_caller)]
-    let location = Some(std::panic::Location::caller());
-    #[cfg(not(tokio_track_caller))]
-    let location = None;
+    let location = trace::caller_location();
 
     match Instant::now().checked_add(duration) {
         Some(deadline) => Sleep::new_timeout(deadline, location),
@@ -226,7 +218,7 @@ cfg_not_trace! {
 }
 
 impl Sleep {
-    #[allow(unused_variables)]
+    #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_variables))]
     pub(crate) fn new_timeout(
         deadline: Instant,
         location: Option<&'static Location<'static>>,
diff --git a/tokio/src/time/timeout.rs b/tokio/src/time/timeout.rs
index 04cd7f1e6b1..007fde86cb2 100644
--- a/tokio/src/time/timeout.rs
+++ b/tokio/src/time/timeout.rs
@@ -4,16 +4,16 @@
 //!
 //! [`Timeout`]: struct@Timeout
 
-use crate::time::{error::Elapsed, sleep_until, Duration, Instant, Sleep};
+use crate::{
+    time::{error::Elapsed, sleep_until, Duration, Instant, Sleep},
+    util::trace,
+};
 
 use pin_project_lite::pin_project;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{self, Poll};
 
-#[cfg(all(tokio_track_caller, tokio_unstable, feature = "tracing"))]
-use std::panic::Location;
-
 /// Require a `Future` to complete before the specified duration has elapsed.
 ///
 /// If the future completes before the duration has elapsed, then the completed
@@ -53,10 +53,7 @@ pub fn timeout<T>(duration: Duration, future: T) -> Timeout<T>
 where
     T: Future,
 {
-    #[cfg(all(tokio_track_caller, tokio_unstable, feature = "tracing"))]
-    let location = Some(Location::caller());
-    #[cfg(not(all(tokio_track_caller, tokio_unstable, feature = "tracing")))]
-    let location = None;
+    let location = trace::caller_location();
 
     let deadline = Instant::now().checked_add(duration);
     let delay = match deadline {
diff --git a/tokio/src/util/trace.rs b/tokio/src/util/trace.rs
index 61c155c5b03..8dae5090014 100644
--- a/tokio/src/util/trace.rs
+++ b/tokio/src/util/trace.rs
@@ -27,6 +27,14 @@ cfg_trace! {
         }
     }
 }
+cfg_time! {
+    pub(crate) fn caller_location() -> Option<&'static std::panic::Location<'static>> {
+        #[cfg(all(tokio_track_caller, tokio_unstable, feature = "tracing"))]
+        return Some(std::panic::Location::caller());
+        #[cfg(not(all(tokio_track_caller, tokio_unstable, feature = "tracing")))]
+        None
+    }
+}
 
 cfg_not_trace! {
     cfg_rt! {