From 11d646eaeba720e848c39e80b06ceb9b2d598bb1 Mon Sep 17 00:00:00 2001
From: "Stainsby, Hayden" <hayden.stainsby@here.com>
Date: Mon, 20 Jun 2022 10:57:51 +0200
Subject: [PATCH] stream: add track_caller to public APIs (#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-stream (only `chunks_timeout` in `StreamExt`) where the
documentation describes how this function may panic due to incorrect
input.

A test has been included to cover the panic.
---
 tokio-stream/Cargo.toml            |  1 +
 tokio-stream/src/stream_ext.rs     |  1 +
 tokio-stream/tests/stream_panic.rs | 55 ++++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 tokio-stream/tests/stream_panic.rs

diff --git a/tokio-stream/Cargo.toml b/tokio-stream/Cargo.toml
index 50899c26c39..a84c924bfd9 100644
--- a/tokio-stream/Cargo.toml
+++ b/tokio-stream/Cargo.toml
@@ -34,6 +34,7 @@ tokio-util = { version = "0.7.0", path = "../tokio-util", optional = true }
 [dev-dependencies]
 tokio = { version = "1.2.0", path = "../tokio", features = ["full", "test-util"] }
 async-stream = "0.3"
+parking_lot = "0.12.0"
 tokio-test = { path = "../tokio-test" }
 futures = { version = "0.3", default-features = false }
 
diff --git a/tokio-stream/src/stream_ext.rs b/tokio-stream/src/stream_ext.rs
index 223b9e102dd..6cea7b55d5f 100644
--- a/tokio-stream/src/stream_ext.rs
+++ b/tokio-stream/src/stream_ext.rs
@@ -1056,6 +1056,7 @@ pub trait StreamExt: Stream {
     /// ```
     #[cfg(feature = "time")]
     #[cfg_attr(docsrs, doc(cfg(feature = "time")))]
+    #[track_caller]
     fn chunks_timeout(self, max_size: usize, duration: Duration) -> ChunksTimeout<Self>
     where
         Self: Sized,
diff --git a/tokio-stream/tests/stream_panic.rs b/tokio-stream/tests/stream_panic.rs
new file mode 100644
index 00000000000..c0a3a45a476
--- /dev/null
+++ b/tokio-stream/tests/stream_panic.rs
@@ -0,0 +1,55 @@
+#![warn(rust_2018_idioms)]
+#![cfg(feature = "time")]
+
+use parking_lot::{const_mutex, Mutex};
+use std::error::Error;
+use std::panic;
+use std::sync::Arc;
+use tokio::time::Duration;
+use tokio_stream::{self as stream, StreamExt};
+
+fn test_panic<Func: FnOnce() + panic::UnwindSafe>(func: Func) -> Option<String> {
+    static PANIC_MUTEX: Mutex<()> = const_mutex(());
+
+    {
+        let _guard = PANIC_MUTEX.lock();
+        let panic_file: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
+
+        let prev_hook = panic::take_hook();
+        {
+            let panic_file = panic_file.clone();
+            panic::set_hook(Box::new(move |panic_info| {
+                let panic_location = panic_info.location().unwrap();
+                panic_file
+                    .lock()
+                    .clone_from(&Some(panic_location.file().to_string()));
+            }));
+        }
+
+        let result = panic::catch_unwind(func);
+        // Return to the previously set panic hook (maybe default) so that we get nice error
+        // messages in the tests.
+        panic::set_hook(prev_hook);
+
+        if result.is_err() {
+            panic_file.lock().clone()
+        } else {
+            None
+        }
+    }
+}
+
+#[test]
+fn stream_chunks_timeout_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(|| {
+        let iter = vec![1, 2, 3].into_iter();
+        let stream0 = stream::iter(iter);
+
+        let _chunk_stream = stream0.chunks_timeout(0, Duration::from_secs(2));
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}