add track_caller to public APIs (#4413)

tokio/Cargo.toml

@@ -140,6 +140,7 @@ version = "0.3.6"
 tokio-test = { version = "0.4.0", path = "../tokio-test" }
 tokio-stream = { version = "0.1", path = "../tokio-stream" }
 futures = { version = "0.3.0", features = ["async-await"] }
+lazy_static = "1.4"
 mockall = "0.11.1"
 tempfile = "3.1.0"
 async-stream = "0.3"

tokio/src/runtime/builder.rs

@@ -200,7 +200,7 @@ impl Builder {
     ///
     /// The default value is the number of cores available to the system.
     ///
-    /// # Panic
+    /// # Panics
     ///
     /// When using the `current_thread` runtime this method will panic, since
     /// those variants do not allow setting worker thread counts.
@@ -237,9 +237,10 @@ impl Builder {
     /// rt.block_on(async move {});
     /// ```
     ///
-    /// # Panic
+    /// # Panics
     ///
     /// This will panic if `val` is not larger than `0`.
+    #[track_caller]
     pub fn worker_threads(&mut self, val: usize) -> &mut Self {
         assert!(val > 0, "Worker threads cannot be set to 0");
         self.worker_threads = Some(val);
@@ -255,7 +256,7 @@ impl Builder {
     ///
     /// The default value is 512.
     ///
-    /// # Panic
+    /// # Panics
     ///
     /// This will panic if `val` is not larger than `0`.
     ///
@@ -267,6 +268,7 @@ impl Builder {
     /// [`spawn_blocking`]: fn@crate::task::spawn_blocking
     /// [`worker_threads`]: Self::worker_threads
     /// [`thread_keep_alive`]: Self::thread_keep_alive
+    #[track_caller]
     #[cfg_attr(docsrs, doc(alias = "max_threads"))]
     pub fn max_blocking_threads(&mut self, val: usize) -> &mut Self {
         assert!(val > 0, "Max blocking threads cannot be set to 0");

tokio/src/runtime/context.rs

@@ -15,6 +15,7 @@ pub(crate) fn try_current() -> Result<Handle, crate::runtime::TryCurrentError> {
     }
 }
 
+#[track_caller]
 pub(crate) fn current() -> Handle {
     match try_current() {
         Ok(handle) => handle,

tokio/src/runtime/handle.rs

@@ -87,7 +87,7 @@ impl Handle {
 
     /// Returns a `Handle` view over the currently running `Runtime`.
     ///
-    /// # Panic
+    /// # Panics
     ///
     /// This will panic if called outside the context of a Tokio runtime. That means that you must
     /// call this on one of the threads **being run by the runtime**, or from a thread with an active
@@ -129,6 +129,7 @@ impl Handle {
     /// # });
     /// # }
     /// ```
+    #[track_caller]
     pub fn current() -> Self {
         context::current()
     }

tokio/src/runtime/mod.rs

@@ -469,7 +469,6 @@ cfg_rt! {
         /// ```
         ///
         /// [handle]: fn@Handle::block_on
-        #[track_caller]
         pub fn block_on<F: Future>(&self, future: F) -> F::Output {
             #[cfg(all(tokio_unstable, feature = "tracing"))]
             let future = crate::util::trace::task(future, "block_on", None, task::Id::next().as_u64());

tokio/src/runtime/task/error.rs

@@ -82,6 +82,7 @@ impl JoinError {
     ///     }
     /// }
     /// ```
+    #[track_caller]
     pub fn into_panic(self) -> Box<dyn Any + Send + 'static> {
         self.try_into_panic()
             .expect("`JoinError` reason is not a panic.")

tokio/tests/rt_panic.rs

@@ -0,0 +1,105 @@
+#![warn(rust_2018_idioms)]
+#![cfg(feature = "full")]
+
+use lazy_static::lazy_static;
+use std::error::Error;
+use std::panic;
+use std::sync::{Arc, Mutex, Once};
+use tokio::runtime::{Builder, Handle, Runtime};
+
+fn test_panic<Func: FnOnce() + panic::UnwindSafe>(func: Func) -> Option<String> {
+    lazy_static! {
+        static ref SET_UP_PANIC_HOOK: Once = Once::new();
+        static ref PANIC_MUTEX: Arc<Mutex<()>> = Arc::new(Mutex::new(()));
+        static ref PANIC_FILE: Mutex<String> = Mutex::new(String::new());
+    }
+
+    SET_UP_PANIC_HOOK.call_once(|| {
+        panic::set_hook(Box::new(|panic_info| {
+            let panic_location = panic_info.location().unwrap();
+            PANIC_FILE
+                .lock()
+                .unwrap()
+                .clone_from(&panic_location.file().to_string());
+        }));
+    });
+
+    {
+        let _guard = PANIC_MUTEX.lock();
+
+        let result = panic::catch_unwind(func);
+
+        if result.is_err() {
+            Some(PANIC_FILE.lock().ok()?.clone())
+        } else {
+            None
+        }
+    }
+}
+
+#[test]
+fn test_current_handle_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(|| {
+        let _ = Handle::current();
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}
+
+#[test]
+fn test_into_panic_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(move || {
+        let rt = basic();
+        rt.block_on(async {
+            let handle = tokio::spawn(async {
+                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+            });
+
+            handle.abort();
+
+            let err = handle.await.unwrap_err();
+            assert!(!&err.is_panic());
+
+            let _ = err.into_panic();
+        });
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}
+
+#[test]
+fn test_builder_worker_threads_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(|| {
+        let _ = Builder::new_multi_thread().worker_threads(0).build();
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}
+
+#[test]
+fn test_builder_max_blocking_threads_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(|| {
+        let _ = Builder::new_multi_thread().max_blocking_threads(0).build();
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}
+
+fn basic() -> Runtime {
+    tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap()
+}