runtime: skip notified tasks during taskdumps

Fixes tokio-rs#6051.

tokio/src/runtime/task/mod.rs

@@ -363,10 +363,14 @@ impl<S: 'static> Task<S> {
     }
 
     cfg_taskdump! {
-        pub(super) fn notify_for_tracing(&self) -> Notified<S> {
-            self.as_raw().state().transition_to_notified_for_tracing();
-            // SAFETY: `transition_to_notified_for_tracing` increments the refcount.
-            unsafe { Notified(Task::new(self.raw)) }
+        pub(super) fn notify_for_tracing(&self) -> Option<Notified<S>> {
+            if self.as_raw().state().transition_to_notified_for_tracing() {
+                // SAFETY: `transition_to_notified_for_tracing` increments the
+                // refcount.
+                Some(unsafe { Notified(Task::new(self.raw)) })
+            } else {
+                None
+            }
         }
     }
 }

tokio/src/runtime/task/state.rs

@@ -278,12 +278,16 @@ impl State {
         target_os = "linux",
         any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
     ))]
-    pub(super) fn transition_to_notified_for_tracing(&self) {
+    pub(super) fn transition_to_notified_for_tracing(&self) -> bool {
         self.fetch_update_action(|mut snapshot| {
-            snapshot.set_notified();
-            snapshot.ref_inc();
-            ((), Some(snapshot))
-        });
+            if snapshot.is_notified() {
+                (false, None)
+            } else {
+                snapshot.set_notified();
+                snapshot.ref_inc();
+                (true, Some(snapshot))
+            }
+        })
     }
 
     /// Sets the cancelled bit and transitions the state to `NOTIFIED` if idle.

tokio/src/runtime/task/trace/mod.rs

@@ -328,8 +328,12 @@ fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>) -> Vec<Trace> {
     // notify each task
     let mut tasks = vec![];
     owned.for_each(|task| {
-        // notify the task (and thus make it poll-able) and stash it
-        tasks.push(task.notify_for_tracing());
+        // Notify the task (and thus make it poll-able) and stash it. This fails
+        // if the task is already notified. In these cases, we skip tracing the
+        // task.
+        if let Some(notified) = task.notify_for_tracing() {
+            tasks.push(notified);
+        }
         // we do not poll it here since we hold a lock on `owned` and the task
         // may complete and need to remove itself from `owned`.
     });

tokio/tests/dump.rs

@@ -154,3 +154,65 @@ 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;
+                //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]
+    #[ignore]
+    // TODO: This currently hangs, with or without the regression fix. Adding a
+    // `yield_now` to the end of `fn dump()` above fixes the issue, but why?
+    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)
+    }
+}