From 613dcd6342a56d8f2dac4acc2e07bc4782f810fa Mon Sep 17 00:00:00 2001
From: noah <nomaxx117@gmail.com>
Date: Sat, 12 Oct 2024 09:51:40 -0500
Subject: [PATCH] update with tracing task size stuff

---
 tokio/src/runtime/handle.rs                |  4 ++--
 tokio/src/runtime/local_runtime/runtime.rs | 20 ++++++++++++++------
 tokio/src/task/local.rs                    |  4 ++--
 tokio/src/util/trace.rs                    |  1 +
 4 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs
index ae8f5969312..752640d75bd 100644
--- a/tokio/src/runtime/handle.rs
+++ b/tokio/src/runtime/handle.rs
@@ -353,7 +353,7 @@ impl Handle {
     pub(crate) unsafe fn spawn_local_named<F>(
         &self,
         future: F,
-        _name: Option<&str>,
+        _meta: SpawnMeta<'_>,
     ) -> JoinHandle<F::Output>
     where
         F: Future + 'static,
@@ -369,7 +369,7 @@ impl Handle {
         ))]
         let future = super::task::trace::Trace::root(future);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
-        let future = crate::util::trace::task(future, "task", _name, id.as_u64());
+        let future = crate::util::trace::task(future, "task", _meta, id.as_u64());
         self.inner.spawn_local(future, id)
     }
 
diff --git a/tokio/src/runtime/local_runtime/runtime.rs b/tokio/src/runtime/local_runtime/runtime.rs
index 9908c856709..0f2b944e4eb 100644
--- a/tokio/src/runtime/local_runtime/runtime.rs
+++ b/tokio/src/runtime/local_runtime/runtime.rs
@@ -5,8 +5,10 @@ use crate::runtime::scheduler::CurrentThread;
 use crate::runtime::{context, Builder, EnterGuard, Handle, BOX_FUTURE_THRESHOLD};
 use crate::task::JoinHandle;
 
+use crate::util::trace::SpawnMeta;
 use std::future::Future;
 use std::marker::PhantomData;
+use std::mem;
 use std::time::Duration;
 
 /// A local Tokio runtime.
@@ -147,12 +149,15 @@ impl LocalRuntime {
         F: Future + 'static,
         F::Output: 'static,
     {
+        let fut_size = std::mem::size_of::<F>();
+        let meta = SpawnMeta::new_unnamed(fut_size);
+
         // safety: spawn_local can only be called from `LocalRuntime`, which this is
         unsafe {
             if std::mem::size_of::<F>() > BOX_FUTURE_THRESHOLD {
-                self.handle.spawn_local_named(Box::pin(future), None)
+                self.handle.spawn_local_named(Box::pin(future), meta)
             } else {
-                self.handle.spawn_local_named(future, None)
+                self.handle.spawn_local_named(future, meta)
             }
         }
     }
@@ -211,15 +216,18 @@ impl LocalRuntime {
     /// ```
     #[track_caller]
     pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+        let fut_size = mem::size_of::<F>();
+        let meta = SpawnMeta::new_unnamed(fut_size);
+
         if std::mem::size_of::<F>() > BOX_FUTURE_THRESHOLD {
-            self.block_on_inner(Box::pin(future))
+            self.block_on_inner(Box::pin(future), meta)
         } else {
-            self.block_on_inner(future)
+            self.block_on_inner(future, meta)
         }
     }
 
     #[track_caller]
-    fn block_on_inner<F: Future>(&self, future: F) -> F::Output {
+    fn block_on_inner<F: Future>(&self, future: F, _meta: SpawnMeta<'_>) -> F::Output {
         #[cfg(all(
             tokio_unstable,
             tokio_taskdump,
@@ -233,7 +241,7 @@ impl LocalRuntime {
         let future = crate::util::trace::task(
             future,
             "block_on",
-            None,
+            _meta,
             crate::runtime::task::Id::next().as_u64(),
         );
 
diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs
index e30d21a88b4..edd02acbac0 100644
--- a/tokio/src/task/local.rs
+++ b/tokio/src/task/local.rs
@@ -409,7 +409,7 @@ cfg_rt! {
                 ))]
                 let future = task::trace::Trace::root(future);
                 let id = task::Id::next();
-                let task = crate::util::trace::task(future, "task", name, id.as_u64());
+                let task = crate::util::trace::task(future, "task", meta, id.as_u64());
 
                 // safety: we have verified that this is a `LocalRuntime` owned by the current thread
                 unsafe { handle.spawn_local(task, id) }
@@ -426,7 +426,7 @@ cfg_rt! {
             Ok(Some(join_handle)) => join_handle,
             Err(_) => match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) {
                 None => panic!("`spawn_local` called from outside of a `task::LocalSet` or LocalRuntime"),
-                Some(cx) => cx.spawn(future.unwrap(), name)
+                Some(cx) => cx.spawn(future.unwrap(), meta)
             }
         }
     }
diff --git a/tokio/src/util/trace.rs b/tokio/src/util/trace.rs
index 97006df474e..b6eadba2205 100644
--- a/tokio/src/util/trace.rs
+++ b/tokio/src/util/trace.rs
@@ -1,6 +1,7 @@
 cfg_rt! {
     use std::marker::PhantomData;
 
+    #[derive(Copy, Clone)]
     pub(crate) struct SpawnMeta<'a> {
         /// The name of the task
         #[cfg(all(tokio_unstable, feature = "tracing"))]