runtime: skip notified tasks during taskdumps

Fixes #6051.

Author: jswrenn

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

@@ -272,14 +272,19 @@ pub(in crate::runtime) fn trace_current_thread(
     injection: &Inject<Arc<current_thread::Handle>>,
 ) -> Vec<Trace> {
     // clear the local and injection queues
-    local.clear();
+
+    let mut dequeued = Vec::new();
+
+    while let Some(task) = local.pop_back() {
+        dequeued.push(task);
+    } 
 
     while let Some(task) = injection.pop() {
-        drop(task);
+        dequeued.push(task);
     }
 
     // precondition: We have drained the tasks from the injection queue.
-    trace_owned(owned)
+    trace_owned(owned, dequeued)
 }
 
 cfg_rt_multi_thread! {
@@ -299,22 +304,24 @@ cfg_rt_multi_thread! {
         synced: &Mutex<Synced>,
         injection: &Shared<Arc<multi_thread::Handle>>,
     ) -> Vec<Trace> {
+        let mut dequeued = Vec::new();
+
         // clear the local queue
         while let Some(notified) = local.pop() {
-            drop(notified);
+            dequeued.push(notified);
         }
 
         // clear the injection queue
         let mut synced = synced.lock();
         while let Some(notified) = injection.pop(&mut synced.inject) {
-            drop(notified);
+            dequeued.push(notified);
         }
 
         drop(synced);
 
         // precondition: we have drained the tasks from the local and injection
         // queues.
-        trace_owned(owned)
+        trace_owned(owned, dequeued)
     }
 }
 
@@ -324,12 +331,16 @@ cfg_rt_multi_thread! {
 ///
 /// This helper presumes exclusive access to each task. The tasks must not exist
 /// in any other queue.
-fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>) -> Vec<Trace> {
+fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>, dequeued: Vec<Notified<S>>) -> Vec<Trace> {
     // notify each task
-    let mut tasks = vec![];
+    let mut tasks = dequeued;
     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

@@ -147,10 +147,73 @@ 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 {
             let _ = tokio::join!(tokio::spawn(complete_during_trace()), dump());
         });
     }
 }
+
+/// 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;
+                // TODO: This tests hangs with this commented out. Why?
+                // 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]
+    // 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)
+    }
+}