taskdump: skip double wake on Trace::capture/Trace::trace_with (#8043)

This commit is contained in:
Russell Cohen
2026-05-12 19:56:50 +02:00
committed by GitHub
parent bdcea6b2cd
commit 0121120b6d
26 changed files with 113 additions and 91 deletions
+2 -1
View File
@@ -1,4 +1,4 @@
323
324
&
+
<
@@ -89,6 +89,7 @@ decrementing
demangled
dequeued
dereferenced
derefs
deregister
deregistered
deregistering
+7 -7
View File
@@ -599,7 +599,7 @@ impl AsyncRead for File {
cx: &mut Context<'_>,
dst: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let me = self.get_mut();
let inner = me.inner.get_mut();
@@ -694,7 +694,7 @@ impl AsyncSeek for File {
}
fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let inner = self.inner.get_mut();
loop {
@@ -730,7 +730,7 @@ impl AsyncWrite for File {
cx: &mut Context<'_>,
src: &[u8],
) -> Poll<io::Result<usize>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let me = self.get_mut();
let inner = me.inner.get_mut();
@@ -801,7 +801,7 @@ impl AsyncWrite for File {
cx: &mut Context<'_>,
bufs: &[io::IoSlice<'_>],
) -> Poll<Result<usize, io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let me = self.get_mut();
let inner = me.inner.get_mut();
@@ -872,13 +872,13 @@ impl AsyncWrite for File {
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let inner = self.inner.get_mut();
inner.poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
self.poll_flush(cx)
}
}
@@ -1077,7 +1077,7 @@ impl Inner {
}
fn poll_complete_inflight(&mut self, cx: &mut Context<'_>) -> Poll<()> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
match self.poll_flush(cx) {
Poll::Ready(Err(e)) => {
self.last_write_err = Some(e.kind());
+1 -1
View File
@@ -82,7 +82,7 @@ impl CopyBuffer {
R: AsyncRead + ?Sized,
W: AsyncWrite + ?Sized,
{
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
#[cfg(any(
feature = "fs",
feature = "io-std",
+7 -7
View File
@@ -71,7 +71,7 @@ impl AsyncRead for Empty {
cx: &mut Context<'_>,
_: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(()))
}
@@ -80,7 +80,7 @@ impl AsyncRead for Empty {
impl AsyncBufRead for Empty {
#[inline]
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(&[]))
}
@@ -96,21 +96,21 @@ impl AsyncWrite for Empty {
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(buf.len()))
}
#[inline]
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(()))
}
#[inline]
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(()))
}
@@ -126,7 +126,7 @@ impl AsyncWrite for Empty {
cx: &mut Context<'_>,
bufs: &[io::IoSlice<'_>],
) -> Poll<Result<usize, io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
let num_bytes = bufs.iter().map(|b| b.len()).sum();
Poll::Ready(Ok(num_bytes))
@@ -141,7 +141,7 @@ impl AsyncSeek for Empty {
#[inline]
fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(0))
}
+6 -6
View File
@@ -331,7 +331,7 @@ impl AsyncRead for SimplexStream {
cx: &mut task::Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let coop = ready!(crate::task::coop::poll_proceed(cx));
let ret = self.poll_read_internal(cx, buf);
@@ -348,7 +348,7 @@ impl AsyncRead for SimplexStream {
cx: &mut task::Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
self.poll_read_internal(cx, buf)
}
}
@@ -361,7 +361,7 @@ impl AsyncWrite for SimplexStream {
cx: &mut task::Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let coop = ready!(crate::task::coop::poll_proceed(cx));
let ret = self.poll_write_internal(cx, buf);
@@ -378,7 +378,7 @@ impl AsyncWrite for SimplexStream {
cx: &mut task::Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
self.poll_write_internal(cx, buf)
}
}
@@ -389,7 +389,7 @@ impl AsyncWrite for SimplexStream {
cx: &mut task::Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> Poll<Result<usize, std::io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let coop = ready!(crate::task::coop::poll_proceed(cx));
let ret = self.poll_write_vectored_internal(cx, bufs);
@@ -406,7 +406,7 @@ impl AsyncWrite for SimplexStream {
cx: &mut task::Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> Poll<Result<usize, std::io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
self.poll_write_vectored_internal(cx, bufs)
}
}
+1 -1
View File
@@ -56,7 +56,7 @@ impl AsyncRead for Repeat {
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
buf.put_bytes(self.byte, buf.remaining());
Poll::Ready(Ok(()))
+3 -3
View File
@@ -57,21 +57,21 @@ impl AsyncWrite for Sink {
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(buf.len()))
}
#[inline]
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(()))
}
#[inline]
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
ready!(poll_proceed_and_make_progress(cx));
Poll::Ready(Ok(()))
}
+2 -2
View File
@@ -574,14 +574,14 @@ mod trace {
cfg_not_taskdump! {
#[inline(always)]
#[allow(dead_code)]
pub(crate) fn trace_leaf(_: &mut std::task::Context<'_>) -> std::task::Poll<()> {
pub(crate) fn trace_leaf() -> std::task::Poll<()> {
std::task::Poll::Ready(())
}
}
#[cfg_attr(not(feature = "sync"), allow(dead_code))]
pub(crate) async fn async_trace_leaf() {
std::future::poll_fn(trace_leaf).await
std::future::poll_fn(|_cx| trace_leaf()).await
}
}
+1 -1
View File
@@ -1135,7 +1135,7 @@ where
type Output = Result<T, E>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));
+4 -3
View File
@@ -227,9 +227,10 @@ impl Trace {
/// should not be much slower than calling `f` directly.
///
/// Due to the way tracing is implemented, Tokio leaf futures will usually, instead of doing their
/// actual work, do the equivalent of a `yield_now` (returning a `Poll::Pending` and scheduling the
/// current context for execution), which means forward progress will probably not happen unless
/// you eventually call your future outside of `capture`.
/// actual work, return `Poll::Pending` without registering the task's waker with any driver.
/// This means forward progress will probably not happen unless you eventually call your future
/// outside of `capture`, or explicitly re-schedule the task (e.g. by calling
/// [`cx.waker().wake_by_ref()`][std::task::Waker::wake_by_ref]) after `capture` returns.
///
/// [`Handle::dump`]: crate::runtime::Handle::dump
///
+1 -1
View File
@@ -146,7 +146,7 @@ impl Registration {
cx: &mut Context<'_>,
direction: Direction,
) -> Poll<io::Result<ReadyEvent>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));
let ev = ready!(self.shared.poll_readiness(cx, direction));
+1 -1
View File
@@ -325,7 +325,7 @@ impl<T> Future for JoinHandle<T> {
type Output = super::Result<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let mut ret = Poll::Pending;
// Keep track of task budget
+14
View File
@@ -509,6 +509,20 @@ impl<S: Schedule> LocalNotified<S> {
mem::forget(self);
raw.poll();
}
/// Returns a `WakerRef` borrowing from this task.
///
/// `WakerRef` derefs to `Waker` without bumping the task's refcount.
#[cfg(all(
tokio_unstable,
feature = "taskdump",
feature = "rt",
target_os = "linux",
any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
))]
pub(crate) fn waker_ref(&self) -> waker::WakerRef<'_, S> {
waker::waker_ref::<S>(self.task.raw.header_ptr_ref())
}
}
impl<S: Schedule> UnownedTask<S> {
+11
View File
@@ -242,6 +242,17 @@ impl RawTask {
self.ptr
}
#[cfg(all(
tokio_unstable,
feature = "taskdump",
feature = "rt",
target_os = "linux",
any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
))]
pub(super) fn header_ptr_ref(&self) -> &NonNull<Header> {
&self.ptr
}
pub(super) fn trailer_ptr(&self) -> NonNull<Trailer> {
unsafe { Header::get_trailer(self.ptr) }
}
+28 -17
View File
@@ -257,7 +257,9 @@ impl Trace {
where
F: FnOnce() -> R,
{
trace_impl::capture(f)
let mut trace = Trace::empty();
let result = trace_with(f, |meta| trace_impl::trace_leaf(meta, &mut trace));
(result, trace)
}
pub(crate) fn empty() -> Self {
@@ -279,17 +281,17 @@ impl Trace {
}
}
/// If this is a sub-invocation of [`Trace::capture`], capture a backtrace.
/// If this is a sub-invocation of [`trace_with`], capture a backtrace.
///
/// The captured backtrace will be returned by [`Trace::capture`].
/// The captured backtrace will be returned by [`trace_with`].
///
/// Invoking this function does nothing when it is not a sub-invocation
/// [`Trace::capture`].
/// [`trace_with`].
// This function is marked `#[inline(never)]` to ensure that it gets a distinct `Frame` in the
// backtrace, below which frames should not be included in the backtrace (since they reflect the
// internal implementation details of this crate).
#[inline(never)]
pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
pub(crate) fn trace_leaf() -> Poll<()> {
let root_addr = Context::current_frame_addr();
let ret = Context::try_with_current_trace_leaf_fn(|leaf_fn| {
@@ -298,18 +300,6 @@ pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
trace_leaf_addr: trace_leaf as *const c_void,
};
leaf_fn(&meta);
// Use the same logic that `yield_now` uses to send out wakeups after
// the task yields.
context::with_scheduler(|scheduler| {
if let Some(scheduler) = scheduler {
match scheduler {
scheduler::Context::CurrentThread(s) => s.defer.defer(cx.waker()),
#[cfg(feature = "rt-multi-thread")]
scheduler::Context::MultiThread(s) => s.defer.defer(cx.waker()),
}
}
});
});
match ret {
@@ -459,6 +449,27 @@ fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>, dequeued: Vec<Notified<S>>) -
.map(|task| {
let local_notified = owned.assert_owner(task);
let id = local_notified.task.id();
// Re-enqueue the task's waker on the scheduler's defer queue so
// the task is polled again after the dump completes. This is the
// same mechanism `yield_now` uses; the defer queue is drained
// after `trace_current_thread` / `trace_multi_thread` returns.
//
// We do this before polling so the borrow of the task ends before
// the `LocalNotified` is consumed in `run()`. `defer` clones the
// waker into its own queue, so the deferred entry outlives the
// `WakerRef` here.
let waker_ref = local_notified.waker_ref();
context::with_scheduler(|scheduler| {
if let Some(scheduler) = scheduler {
match scheduler {
scheduler::Context::CurrentThread(s) => s.defer.defer(&waker_ref),
#[cfg(feature = "rt-multi-thread")]
scheduler::Context::MultiThread(s) => s.defer.defer(&waker_ref),
}
}
});
let ((), trace) = Trace::capture(|| local_notified.run());
(id, trace)
})
+1 -14
View File
@@ -4,20 +4,7 @@
use std::ptr;
use crate::runtime::task::trace::{trace_with, Trace, TraceMeta};
/// Capture using the default `backtrace::trace`-based implementation.
#[inline(never)]
pub(super) fn capture<F, R>(f: F) -> (R, Trace)
where
F: FnOnce() -> R,
{
let mut trace = Trace::empty();
let result = trace_with(f, |meta| trace_leaf(meta, &mut trace));
(result, trace)
}
use crate::runtime::task::trace::{Trace, TraceMeta};
/// Capture a backtrace via `backtrace::trace` and collect it into `trace`.
pub(crate) fn trace_leaf(meta: &TraceMeta, trace: &mut Trace) {
+1 -1
View File
@@ -6,7 +6,7 @@ use std::ops;
use std::ptr::NonNull;
use std::task::{RawWaker, RawWakerVTable, Waker};
pub(super) struct WakerRef<'a, S: 'static> {
pub(crate) struct WakerRef<'a, S: 'static> {
waker: ManuallyDrop<Waker>,
_p: PhantomData<(&'a Header, S)>,
}
+1 -1
View File
@@ -576,7 +576,7 @@ impl Future for Acquire<'_> {
type Output = Result<(), AcquireError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
#[cfg(all(tokio_unstable, feature = "tracing"))]
let _resource_span = self.node.ctx.resource_span.clone().entered();
+1 -1
View File
@@ -1605,7 +1605,7 @@ where
type Output = Result<T, RecvError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
let (receiver, waiter) = self.project();
+2 -2
View File
@@ -289,7 +289,7 @@ impl<T, S: Semaphore> Rx<T, S> {
pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
use super::block::Read;
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));
@@ -349,7 +349,7 @@ impl<T, S: Semaphore> Rx<T, S> {
) -> Poll<usize> {
use super::block::Read;
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));
+4 -6
View File
@@ -1227,9 +1227,8 @@ impl NotifiedProject<'_> {
}
State::Waiting => {
#[cfg(feature = "taskdump")]
if let Some(waker) = waker {
let mut ctx = Context::from_waker(waker);
std::task::ready!(crate::trace::trace_leaf(&mut ctx));
if let Some(_waker) = waker {
std::task::ready!(crate::trace::trace_leaf());
}
if waiter.notification.load(Acquire).is_some() {
@@ -1321,9 +1320,8 @@ impl NotifiedProject<'_> {
}
State::Done => {
#[cfg(feature = "taskdump")]
if let Some(waker) = waker {
let mut ctx = Context::from_waker(waker);
std::task::ready!(crate::trace::trace_leaf(&mut ctx));
if let Some(_waker) = waker {
std::task::ready!(crate::trace::trace_leaf());
}
return Poll::Ready(());
}
+2 -2
View File
@@ -818,7 +818,7 @@ impl<T> Sender<T> {
/// # }
/// ```
pub fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));
@@ -1313,7 +1313,7 @@ impl<T> Inner<T> {
}
fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));
+1 -1
View File
@@ -26,7 +26,7 @@ pub async fn consume_budget() {
let mut status = std::task::Poll::Pending;
std::future::poll_fn(move |cx| {
std::task::ready!(crate::trace::trace_leaf(cx));
std::task::ready!(crate::trace::trace_leaf());
if status.is_ready() {
return status;
}
+1 -1
View File
@@ -38,7 +38,7 @@ use std::task::{ready, Poll};
pub async fn yield_now() {
let mut yielded = false;
poll_fn(|cx| {
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
if yielded {
return Poll::Ready(());
+1 -1
View File
@@ -417,7 +417,7 @@ impl Sleep {
fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> {
let me = self.project();
ready!(crate::trace::trace_leaf(cx));
ready!(crate::trace::trace_leaf());
// Keep track of task budget
#[cfg(all(tokio_unstable, feature = "tracing"))]
+9 -10
View File
@@ -61,6 +61,10 @@ impl<F: Future> Future for PrettyFuture<F> {
let (res, trace) = tokio::runtime::dump::Trace::capture(|| this.f.as_mut().poll(cx));
this.logs.lock().unwrap().push(trace);
*this.t_last = State::Alerted;
// `Trace::capture` does not reschedule the task. Wake the task
// ourselves so the wrapped future gets polled again and can make
// progress now that tracing has captured its state.
cx.waker().wake_by_ref();
return res;
}
this.f.poll(cx)
@@ -184,21 +188,16 @@ impl<F: Future> Future for TaskDump<F> {
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
let mut this = self.project();
// Poll the future with a real waker. if it returns Ready, exit immediately
match this.f.as_mut().poll(cx) {
Poll::Ready(result) => return Poll::Ready(result),
Poll::Pending => {}
// if the future is ready, exit immediately
if let Poll::Ready(result) = this.f.as_mut().poll(cx) {
return Poll::Ready(result);
};
// if is pending, trace its location:
let mut logs = Vec::new();
// Tracing poll with a noop waker. If the future is at a yield
// point, trace_leaf fires our callback and returns Pending. We discard
// the result — this poll is purely for capturing the backtrace.
let noop = futures::task::noop_waker();
let mut noop_cx = Context::from_waker(&noop);
let trace_poll = trace_with(
|| this.f.as_mut().poll(&mut noop_cx),
|| this.f.as_mut().poll(cx),
|meta| trace_leaf_for_test(meta, &mut logs),
);
// trace should always produce poll pending