Add track_caller to tokio core APIs (tokio-rs#4413)

Functions that may panic can be annotated with `#[track_caller]` so that
in the event of a panic, the function where the user called the
panicking function is shown instead of the file and line within Tokio
source.

This change adds track caller to all the non-unstable public APIs in
Tokio core where the documentation describes how the function may panic
due to incorrect context or inputs.  Since each internal function needs
to be annotated down to the actual panic, it makes sense to start in
Tokio core functionality.

Tests are needed to ensure that all the annotations remain in place in
case internal refactoring occurs.

The test installs a panic hook to extract the file location from the
`PanicInfo` struct and clone it up to the outer scope to check that the
panic was indeed reported from within the test file.  The downside to
this approach is that the panic hook is global while set and so we need
a lot of extra functionality to effectively serialize the tests so that
only a single panic can occur at a time.

The annotation of `block_on` was removed as it did not work. It appears
to be impossible to correctly chain track caller when the call stack to
the panic passes through clojures, as the track caller annotation can
only be applied to functions. Also, the panic message itself is very
descriptive.

Author: hds

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()
+}