cache the timestamp on poll_start

Author: rcoh

dial9-tokio-telemetry/src/task_dumped.rs

@@ -35,6 +35,7 @@ use crate::telemetry::{Encodable, ThreadLocalEncoder};
 use pin_project_lite::pin_project;
 use smallvec::SmallVec;
 use std::future::Future;
+use std::num::NonZeroU64;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::atomic::Ordering;
@@ -54,7 +55,7 @@ pin_project! {
         frames: FrameBuf,
         // Monotonic nanoseconds when the frames in `frames` were captured.
         // Only meaningful when `frames.has_data()`.
-        pending_capture_ts: u64,
+        pending_capture_ts: Option<NonZeroU64>,
         // Snapshot of SharedState::task_dump_idle_threshold_ns at construction;
         // promotes the atomic load off the poll hot path.
         idle_threshold_ns: u64,
@@ -69,7 +70,7 @@ impl<F> TaskDumped<F> {
             shared,
             task_id,
             frames: FrameBuf::new(),
-            pending_capture_ts: 0,
+            pending_capture_ts: None,
             idle_threshold_ns,
         }
     }
@@ -91,38 +92,42 @@ impl<F: Future> Future for TaskDumped<F> {
             // after capture resumes.
             if this.frames.has_data() {
                 this.frames.clear();
-                *this.pending_capture_ts = 0;
+                *this.pending_capture_ts = None;
             }
             return this.inner.poll(cx);
         }
 
         // If we have captured frames from a previous idle, decide whether
         // that idle was long enough to emit.
-        let should_emit = if this.frames.has_data() {
-            let now = crate::telemetry::events::clock_monotonic_ns();
-            now.saturating_sub(*this.pending_capture_ts) > *this.idle_threshold_ns
-        } else {
-            false
+        let poll_start = crate::telemetry::recorder::poll_start_ts_or_now();
+        let should_emit = match *this.pending_capture_ts {
+            Some(ts) if this.frames.has_data() => {
+                poll_start.saturating_sub(ts.get()) > *this.idle_threshold_ns
+            }
+            _ => false,
         };
 
         let result = this.inner.as_mut().poll(cx);
 
         if should_emit {
-            this.frames
-                .emit(this.shared, *this.task_id, *this.pending_capture_ts);
+            this.frames.emit(
+                this.shared,
+                *this.task_id,
+                this.pending_capture_ts.unwrap().get(),
+            );
         }
 
         match &result {
             Poll::Ready(_) => {
                 // Terminal. Nothing more to capture; discard stale frames.
                 this.frames.clear();
-                *this.pending_capture_ts = 0;
+                *this.pending_capture_ts = None;
             }
             Poll::Pending => {
                 // Capture the yield point we just landed on, for the next
                 // poll's threshold check.
                 this.frames.capture(this.inner.as_mut());
-                *this.pending_capture_ts = crate::telemetry::events::clock_monotonic_ns();
+                *this.pending_capture_ts = NonZeroU64::new(poll_start);
             }
         }
 

dial9-tokio-telemetry/src/telemetry/recorder/mod.rs

@@ -4,6 +4,8 @@ mod shared_state;
 
 pub(crate) use runtime_context::RuntimeContext;
 pub use runtime_context::current_worker_id;
+#[cfg(feature = "taskdump")]
+pub(crate) use runtime_context::poll_start_ts_or_now;
 pub(crate) use shared_state::SharedState;
 
 use event_writer::EventWriter;
@@ -913,13 +915,18 @@ impl<P> TracedRuntimeBuilder<P> {
     }
 
     /// Capture async backtraces at yield points for tasks that stay idle
-    /// longer than the configured threshold. Requires the `taskdump` crate
-    /// feature to actually record events; the builder accepts the config
-    /// unconditionally so the API surface is stable.
+    /// longer than the configured threshold.
+    ///
+    /// Requires the `taskdump` crate feature to actually record events
     pub fn with_task_dumps(
         mut self,
         config: crate::telemetry::task_dump_config::TaskDumpConfig,
     ) -> Self {
+        if cfg!(not(feature = "taskdump")) {
+            tracing::warn!(
+                "taskdumps enabled but `taskdump` feature was not. No task dumps will be captured."
+            )
+        }
         self.task_dump_config = Some(config);
         self
     }

dial9-tokio-telemetry/src/telemetry/recorder/runtime_context.rs

@@ -7,6 +7,8 @@ use crate::telemetry::format::{
 use crate::telemetry::task_metadata::TaskId;
 use std::cell::Cell;
 use std::collections::HashMap;
+#[cfg(feature = "taskdump")]
+use std::num::NonZeroU64;
 use std::sync::OnceLock;
 use std::sync::RwLock;
 use tokio::runtime::RuntimeMetrics;
@@ -35,6 +37,21 @@ thread_local! {
     /// Whether we've registered this thread's OS tid for CPU profiling.
     #[cfg(feature = "cpu-profiling")]
     static TID_REGISTERED: Cell<bool> = const { Cell::new(false) };
+    /// Monotonic timestamp captured in `on_before_task_poll`, cleared in
+    /// `on_after_task_poll`. Allows code running inside a poll (e.g.
+    /// `TaskDumped`) to reuse the timestamp without an extra clock read.
+    #[cfg(feature = "taskdump")]
+    static POLL_START_TS: Cell<Option<NonZeroU64>> = const { Cell::new(None) };
+}
+
+/// Returns the poll-start timestamp if we're inside a poll, otherwise reads
+/// the clock.
+#[cfg(feature = "taskdump")]
+pub(crate) fn poll_start_ts_or_now() -> u64 {
+    POLL_START_TS.with(|c| c.get()).map_or_else(
+        crate::telemetry::events::clock_monotonic_ns,
+        NonZeroU64::get,
+    )
 }
 
 impl RuntimeContext {
@@ -209,8 +226,11 @@ pub(super) fn make_poll_start(
     let worker_local_queue_depth = resolved
         .map(|(_, idx)| ctx.local_queue_depth(idx))
         .unwrap_or(0);
+    let timestamp_ns = crate::telemetry::events::clock_monotonic_ns();
+    #[cfg(feature = "taskdump")]
+    POLL_START_TS.with(|c| c.set(NonZeroU64::new(timestamp_ns)));
     PollStart {
-        timestamp_ns: crate::telemetry::events::clock_monotonic_ns(),
+        timestamp_ns,
         worker_id: resolved.map(|(id, _)| id).unwrap_or(WorkerId::UNKNOWN),
         local_queue: worker_local_queue_depth as u8,
         task_id,
@@ -219,6 +239,8 @@ pub(super) fn make_poll_start(
 }
 
 pub(super) fn make_poll_end(ctx: &RuntimeContext, shared: &SharedState) -> PollEndEvent {
+    #[cfg(feature = "taskdump")]
+    POLL_START_TS.with(|c| c.set(None));
     let resolved = ctx.resolve_worker(shared);
     PollEndEvent {
         timestamp_ns: crate::telemetry::events::clock_monotonic_ns(),