From 7ae7347400c17eeb795f8c7a65b594f3912fcf48 Mon Sep 17 00:00:00 2001
From: Jack Wrenn <jack@wrenn.fyi>
Date: Wed, 6 Dec 2023 15:54:29 +0000
Subject: [PATCH] tests: add regression test for #6051

---
 tokio/tests/dump.rs | 61 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/tokio/tests/dump.rs b/tokio/tests/dump.rs
index 4da0c9e8e18..811bd7856af 100644
--- a/tokio/tests/dump.rs
+++ b/tokio/tests/dump.rs
@@ -147,6 +147,7 @@ mod future_completes_during_trace {
         async fn dump() {
             let handle = Handle::current();
             let _dump = handle.dump().await;
+            tokio::task::yield_now().await;
         }
 
         rt.block_on(async {
@@ -154,3 +155,63 @@ mod future_completes_during_trace {
         });
     }
 }
+
+/// Regression tests for #6051.
+///
+/// These tests ensure that tasks notified outside of a worker will not be
+/// traced, since doing so will un-set their notified bit prior to them being
+/// run and panic.
+mod notified_during_tracing {
+    use super::*;
+
+    fn test(rt: tokio::runtime::Runtime) {
+        async fn dump() {
+            loop {
+                let handle = Handle::current();
+                let _dump = handle.dump().await;
+                // Without this yield, the `current_runtime` test hangs.
+                tokio::task::yield_now().await;
+            }
+        }
+
+        rt.block_on(async {
+            let timer = tokio::spawn(async {
+                loop {
+                    tokio::time::sleep(tokio::time::Duration::from_nanos(1)).await;
+                }
+            });
+
+            let timeout = async {
+                tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+            };
+
+            tokio::select!(
+                biased;
+                _ = timeout => {},
+                _ = timer => {},
+                _ = dump() => {},
+            );
+        });
+    }
+
+    #[test]
+    fn current_thread() {
+        let rt = runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+
+        test(rt)
+    }
+
+    #[test]
+    fn multi_thread() {
+        let rt = runtime::Builder::new_multi_thread()
+            .enable_all()
+            .worker_threads(3)
+            .build()
+            .unwrap();
+
+        test(rt)
+    }
+}