diff --git a/spellcheck.dic b/spellcheck.dic index ea6ab93ae..aeb1e9db4 100644 --- a/spellcheck.dic +++ b/spellcheck.dic @@ -1,4 +1,4 @@ -323 +324 & + < @@ -89,6 +89,7 @@ decrementing demangled dequeued dereferenced +derefs deregister deregistered deregistering diff --git a/tokio/src/fs/file.rs b/tokio/src/fs/file.rs index 12c1f2dbe..c23164356 100644 --- a/tokio/src/fs/file.rs +++ b/tokio/src/fs/file.rs @@ -599,7 +599,7 @@ impl AsyncRead for File { cx: &mut Context<'_>, dst: &mut ReadBuf<'_>, ) -> Poll> { - 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> { - 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> { - 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> { - 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> { - 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> { - 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()); diff --git a/tokio/src/io/util/copy.rs b/tokio/src/io/util/copy.rs index ac7eb891d..75f81dc8e 100644 --- a/tokio/src/io/util/copy.rs +++ b/tokio/src/io/util/copy.rs @@ -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", diff --git a/tokio/src/io/util/empty.rs b/tokio/src/io/util/empty.rs index 0cc6211cb..f875dafec 100644 --- a/tokio/src/io/util/empty.rs +++ b/tokio/src/io/util/empty.rs @@ -71,7 +71,7 @@ impl AsyncRead for Empty { cx: &mut Context<'_>, _: &mut ReadBuf<'_>, ) -> Poll> { - 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> { - 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> { - 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> { - 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> { - 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> { - 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> { - ready!(crate::trace::trace_leaf(cx)); + ready!(crate::trace::trace_leaf()); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(0)) } diff --git a/tokio/src/io/util/mem.rs b/tokio/src/io/util/mem.rs index 6c3910662..d4cd08eec 100644 --- a/tokio/src/io/util/mem.rs +++ b/tokio/src/io/util/mem.rs @@ -331,7 +331,7 @@ impl AsyncRead for SimplexStream { cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - 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> { - 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> { - 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> { - 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> { - 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> { - ready!(crate::trace::trace_leaf(cx)); + ready!(crate::trace::trace_leaf()); self.poll_write_vectored_internal(cx, bufs) } } diff --git a/tokio/src/io/util/repeat.rs b/tokio/src/io/util/repeat.rs index e5bce46ab..90edb284f 100644 --- a/tokio/src/io/util/repeat.rs +++ b/tokio/src/io/util/repeat.rs @@ -56,7 +56,7 @@ impl AsyncRead for Repeat { cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - 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(())) diff --git a/tokio/src/io/util/sink.rs b/tokio/src/io/util/sink.rs index 28adc5408..4bdecba1f 100644 --- a/tokio/src/io/util/sink.rs +++ b/tokio/src/io/util/sink.rs @@ -57,21 +57,21 @@ impl AsyncWrite for Sink { cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { - 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> { - 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> { - ready!(crate::trace::trace_leaf(cx)); + ready!(crate::trace::trace_leaf()); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs index 7da7e21fd..51f1cfc34 100644 --- a/tokio/src/lib.rs +++ b/tokio/src/lib.rs @@ -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 } } diff --git a/tokio/src/process/mod.rs b/tokio/src/process/mod.rs index aa57a9d40..fb233cb4e 100644 --- a/tokio/src/process/mod.rs +++ b/tokio/src/process/mod.rs @@ -1135,7 +1135,7 @@ where type Output = Result; fn poll(mut self: Pin<&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)); diff --git a/tokio/src/runtime/dump.rs b/tokio/src/runtime/dump.rs index 40f364f91..0a00bca54 100644 --- a/tokio/src/runtime/dump.rs +++ b/tokio/src/runtime/dump.rs @@ -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 /// diff --git a/tokio/src/runtime/io/registration.rs b/tokio/src/runtime/io/registration.rs index 63b6a6d78..bc5c54d4a 100644 --- a/tokio/src/runtime/io/registration.rs +++ b/tokio/src/runtime/io/registration.rs @@ -146,7 +146,7 @@ impl Registration { cx: &mut Context<'_>, direction: Direction, ) -> 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)); let ev = ready!(self.shared.poll_readiness(cx, direction)); diff --git a/tokio/src/runtime/task/join.rs b/tokio/src/runtime/task/join.rs index b79f80af3..a22c22d9f 100644 --- a/tokio/src/runtime/task/join.rs +++ b/tokio/src/runtime/task/join.rs @@ -325,7 +325,7 @@ impl Future for JoinHandle { type Output = super::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - ready!(crate::trace::trace_leaf(cx)); + ready!(crate::trace::trace_leaf()); let mut ret = Poll::Pending; // Keep track of task budget diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 53c477d52..6f964c1de 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -509,6 +509,20 @@ impl LocalNotified { 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::(self.task.raw.header_ptr_ref()) + } } impl UnownedTask { diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index 96453e640..a9d143f74 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/tokio/src/runtime/task/raw.rs @@ -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
{ + &self.ptr + } + pub(super) fn trailer_ptr(&self) -> NonNull { unsafe { Header::get_trailer(self.ptr) } } diff --git a/tokio/src/runtime/task/trace/mod.rs b/tokio/src/runtime/task/trace/mod.rs index 2f98670fa..8e088d5d4 100644 --- a/tokio/src/runtime/task/trace/mod.rs +++ b/tokio/src/runtime/task/trace/mod.rs @@ -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(owned: &OwnedTasks, dequeued: Vec>) - .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) }) diff --git a/tokio/src/runtime/task/trace/trace_impl.rs b/tokio/src/runtime/task/trace/trace_impl.rs index 3197cca8f..14ec8ba1b 100644 --- a/tokio/src/runtime/task/trace/trace_impl.rs +++ b/tokio/src/runtime/task/trace/trace_impl.rs @@ -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: 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) { diff --git a/tokio/src/runtime/task/waker.rs b/tokio/src/runtime/task/waker.rs index 712260d22..bf8e2a6b1 100644 --- a/tokio/src/runtime/task/waker.rs +++ b/tokio/src/runtime/task/waker.rs @@ -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, _p: PhantomData<(&'a Header, S)>, } diff --git a/tokio/src/sync/batch_semaphore.rs b/tokio/src/sync/batch_semaphore.rs index 6447cefb4..405cc01fe 100644 --- a/tokio/src/sync/batch_semaphore.rs +++ b/tokio/src/sync/batch_semaphore.rs @@ -576,7 +576,7 @@ impl Future for Acquire<'_> { type Output = Result<(), AcquireError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - 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(); diff --git a/tokio/src/sync/broadcast.rs b/tokio/src/sync/broadcast.rs index 49e879568..50f637e15 100644 --- a/tokio/src/sync/broadcast.rs +++ b/tokio/src/sync/broadcast.rs @@ -1605,7 +1605,7 @@ where type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - ready!(crate::trace::trace_leaf(cx)); + ready!(crate::trace::trace_leaf()); let (receiver, waiter) = self.project(); diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs index f9a72de60..0945d8805 100644 --- a/tokio/src/sync/mpsc/chan.rs +++ b/tokio/src/sync/mpsc/chan.rs @@ -289,7 +289,7 @@ impl Rx { pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll> { 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 Rx { ) -> Poll { 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)); diff --git a/tokio/src/sync/notify.rs b/tokio/src/sync/notify.rs index 7fb856c98..1af800e79 100644 --- a/tokio/src/sync/notify.rs +++ b/tokio/src/sync/notify.rs @@ -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(()); } diff --git a/tokio/src/sync/oneshot.rs b/tokio/src/sync/oneshot.rs index 9603081b9..3b8055b68 100644 --- a/tokio/src/sync/oneshot.rs +++ b/tokio/src/sync/oneshot.rs @@ -818,7 +818,7 @@ impl Sender { /// # } /// ``` 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 Inner { } fn poll_recv(&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)); diff --git a/tokio/src/task/coop/consume_budget.rs b/tokio/src/task/coop/consume_budget.rs index 4133facb6..61050b238 100644 --- a/tokio/src/task/coop/consume_budget.rs +++ b/tokio/src/task/coop/consume_budget.rs @@ -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; } diff --git a/tokio/src/task/yield_now.rs b/tokio/src/task/yield_now.rs index a4832e40c..a58e85ee7 100644 --- a/tokio/src/task/yield_now.rs +++ b/tokio/src/task/yield_now.rs @@ -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(()); diff --git a/tokio/src/time/sleep.rs b/tokio/src/time/sleep.rs index b2dcf5c59..2aa19b666 100644 --- a/tokio/src/time/sleep.rs +++ b/tokio/src/time/sleep.rs @@ -417,7 +417,7 @@ impl Sleep { fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll> { 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"))] diff --git a/tokio/tests/task_trace_self.rs b/tokio/tests/task_trace_self.rs index e3e0c4791..6553a0a1f 100644 --- a/tokio/tests/task_trace_self.rs +++ b/tokio/tests/task_trace_self.rs @@ -61,6 +61,10 @@ impl Future for PrettyFuture { 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 Future for TaskDump { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 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