Skip to content

Commit

Permalink
update with tracing task size stuff
Browse files Browse the repository at this point in the history
  • Loading branch information
Noah-Kennedy committed Oct 12, 2024
1 parent 017adc9 commit 613dcd6
Show file tree
Hide file tree
Showing 4 changed files with 19 additions and 10 deletions.
4 changes: 2 additions & 2 deletions tokio/src/runtime/handle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
}

Expand Down
20 changes: 14 additions & 6 deletions tokio/src/runtime/local_runtime/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
}
}
}
Expand Down Expand Up @@ -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,
Expand All @@ -233,7 +241,7 @@ impl LocalRuntime {
let future = crate::util::trace::task(
future,
"block_on",
None,
_meta,
crate::runtime::task::Id::next().as_u64(),
);

Expand Down
4 changes: 2 additions & 2 deletions tokio/src/task/local.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) }
Expand All @@ -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)
}
}
}
Expand Down
1 change: 1 addition & 0 deletions tokio/src/util/trace.rs
Original file line number Diff line number Diff line change
@@ -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"))]
Expand Down

0 comments on commit 613dcd6

Please sign in to comment.