From 2849fb3b374210cae446dabe71e24fb861cc8618 Mon Sep 17 00:00:00 2001
From: Jack Wrenn <jack@wrenn.fyi>
Date: Mon, 2 Oct 2023 17:27:54 +0000
Subject: [PATCH] rt: don't persist raw tasks

---
 tokio/src/runtime/task/mod.rs       | 10 +++++
 tokio/src/runtime/task/trace/mod.rs | 69 +++++++++++++----------------
 2 files changed, 42 insertions(+), 37 deletions(-)

diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs
index abf7cc266e7..75157625252 100644
--- a/tokio/src/runtime/task/mod.rs
+++ b/tokio/src/runtime/task/mod.rs
@@ -444,6 +444,16 @@ impl<S: Schedule> UnownedTask<S> {
     }
 }
 
+impl<S: 'static> Clone for Task<S> {
+    fn clone(&self) -> Task<S> {
+        // SAFETY: We increment the ref count.
+        unsafe {
+            self.raw.ref_inc();
+            Task::new(self.raw)
+        }
+    }
+}
+
 impl<S: 'static> Drop for Task<S> {
     fn drop(&mut self) {
         // Decrement the ref count
diff --git a/tokio/src/runtime/task/trace/mod.rs b/tokio/src/runtime/task/trace/mod.rs
index 5fb9e71474e..f6f1d9ef72f 100644
--- a/tokio/src/runtime/task/trace/mod.rs
+++ b/tokio/src/runtime/task/trace/mod.rs
@@ -278,24 +278,8 @@ pub(in crate::runtime) fn trace_current_thread(
         drop(task);
     }
 
-    // notify each task
-    let mut tasks = vec![];
-    owned.for_each(|task| {
-        // set the notified bit
-        task.as_raw().state().transition_to_notified_for_tracing();
-        // store the raw tasks into a vec
-        tasks.push(task.as_raw());
-        // do NOT poll `task` here, since we hold a lock on `owned` and the task
-        // may complete and need to remove itself from `owned`.
-    });
-
-    tasks
-        .into_iter()
-        .map(|task| {
-            let ((), trace) = Trace::capture(|| task.poll());
-            trace
-        })
-        .collect()
+    // precondition: We have drained the tasks from the injection queue.
+    trace_owned(owned)
 }
 
 cfg_rt_multi_thread! {
@@ -328,24 +312,35 @@ cfg_rt_multi_thread! {
 
         drop(synced);
 
-        // notify each task
-        let mut tasks = vec![];
-        owned.for_each(|task| {
-            // set the notified bit
-            task.as_raw().state().transition_to_notified_for_tracing();
-            // store the raw tasks into a vec
-            tasks.push(task.as_raw());
-            // do NOT poll `task` here, since we hold a lock on `owned` and the task
-            // may complete and need to remove itself from `owned`.
-        });
-
-        tasks
-            .into_iter()
-            .map(|task| {
-                // trace the task
-                let ((), trace) = Trace::capture(|| task.poll());
-                trace
-            })
-            .collect()
+        // precondition: we have drained the tasks from the local and injection
+        // queues.
+        trace_owned(owned)
     }
 }
+
+/// Trace the `OwnedTasks`.
+///
+/// # Preconditions
+///
+/// This helper presumes exclusive access to each task. The tasks must not exist
+/// in any other queue.
+fn trace_owned<S>(owned: &OwnedTasks<Arc<S>>) -> Vec<Trace> {
+    // notify each task
+    let mut tasks = vec![];
+    owned.for_each(|task| {
+        // set the notified bit
+        task.as_raw().state().transition_to_notified_for_tracing();
+        // store the tasks into a vec
+        tasks.push(task.clone());
+        // do NOT poll `task` here, since we hold a lock on `owned` and the task
+        // may complete and need to remove itself from `owned`.
+    });
+
+    tasks
+        .into_iter()
+        .map(|task| {
+            let ((), trace) = Trace::capture(|| task.as_raw().poll());
+            trace
+        })
+        .collect()
+}