Skip to content

Commit

Permalink
taskdump: implement task dumps for multi-thread runtime
Browse files Browse the repository at this point in the history
This PR implements task dumps on the multi-thread runtime. It
complements #5608, which implemented task dumps on the current-thread
runtime.
  • Loading branch information
jswrenn committed Jun 5, 2023
1 parent 076d77c commit 09cebae
Show file tree
Hide file tree
Showing 8 changed files with 399 additions and 77 deletions.
32 changes: 20 additions & 12 deletions examples/dump.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
target_os = "linux",
any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
))]
#[tokio::main(flavor = "current_thread")]
#[tokio::main]
async fn main() {
use std::hint::black_box;

Expand All @@ -22,21 +22,29 @@ async fn main() {

#[inline(never)]
async fn c() {
black_box(tokio::task::yield_now()).await
loop {
tokio::task::yield_now().await;
}
}

tokio::spawn(a());
tokio::spawn(b());
tokio::spawn(c());
async fn dump() {
let handle = tokio::runtime::Handle::current();
let dump = handle.dump().await;

let handle = tokio::runtime::Handle::current();
let dump = handle.dump();

for (i, task) in dump.tasks().iter().enumerate() {
let trace = task.trace();
println!("task {i} trace:");
println!("{trace}");
for (i, task) in dump.tasks().iter().enumerate() {
let trace = task.trace();
println!("task {i} trace:");
println!("{trace}\n");
}
}

tokio::select!(
biased;
_ = tokio::spawn(a()) => {},
_ = tokio::spawn(b()) => {},
_ = tokio::spawn(c()) => {},
_ = dump() => {},
);
}

#[cfg(not(all(
Expand Down
31 changes: 28 additions & 3 deletions tokio/src/runtime/handle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -343,15 +343,40 @@ cfg_metrics! {
cfg_taskdump! {
impl Handle {
/// Capture a snapshot of this runtime's state.
pub fn dump(&self) -> crate::runtime::Dump {
pub async fn dump(&self) -> crate::runtime::Dump {
match &self.inner {
scheduler::Handle::CurrentThread(handle) => handle.dump(),
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
scheduler::Handle::MultiThread(_) =>
unimplemented!("taskdumps are unsupported on the multi-thread runtime"),
scheduler::Handle::MultiThread(handle) => {
// perform the trace in a separate thread so that the
// trace itself does not appear in the taskdump.
let handle = handle.clone();
spawn_thread(async {
let handle = handle;
handle.dump().await
}).await
},
}
}
}

cfg_rt_multi_thread! {
/// Spawn a new thread and asynchronously await on its result.
async fn spawn_thread<F>(f: F) -> <F as Future>::Output
where
F: Future + Send + 'static,
<F as Future>::Output: Send + 'static
{
let (tx, rx) = crate::sync::oneshot::channel();
crate::loom::thread::spawn(|| {
let rt = crate::runtime::Builder::new_current_thread().build().unwrap();
rt.block_on(async {
let _ = tx.send(f.await);
});
});
rx.await.unwrap()
}
}
}

/// Error returned by `try_current` when no Runtime has been started
Expand Down
25 changes: 25 additions & 0 deletions tokio/src/runtime/scheduler/multi_thread/handle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,31 @@ cfg_metrics! {
}
}

cfg_taskdump! {
impl Handle {
pub(crate) async fn dump(&self) -> crate::runtime::Dump {
let trace_status = &self.shared.trace_status;

// If a dump is in progress, block.
trace_status.start_trace_request(&self).await;

let result = loop {
if let Some(result) = trace_status.take_result() {
break result;
} else {
self.notify_all();
trace_status.result_ready.notified().await;
}
};

// Allow other queued dumps to proceed.
trace_status.end_trace_request(&self).await;

result
}
}
}

impl fmt::Debug for Handle {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("multi_thread::Handle { ... }").finish()
Expand Down
4 changes: 4 additions & 0 deletions tokio/src/runtime/scheduler/multi_thread/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,10 @@ pub(crate) mod queue;
mod worker;
pub(crate) use worker::{Context, Launch, Shared};

cfg_taskdump! {
pub(crate) use worker::Synced;
}

pub(crate) use worker::block_in_place;

use crate::loom::sync::Arc;
Expand Down
Loading

0 comments on commit 09cebae

Please sign in to comment.