Initial implementation of tracking task wakes (#4)

* Add integration test & fix worker id issue

* Initial implementation of tracking task wakes

* Fix merge issues: add WakeEvent handling and restore traced.rs dependencies

* enable waker instrumentation for example service

* Rewrite waker with ArcWaker

* Add better tracing & add raw wake event

* improvements to trace viewer for wakes

* refactor code that produces wake events

* Improvements to trace viewer

* cleanup clippy

Author: rcoh

Cargo.lock

@@ -8,6 +8,7 @@ publish = true
 tokio = { version = "1", features = ["rt", "macros", "time", "taskdump", "rt-multi-thread", "sync", "net", "io-util"] }
 arc-swap = "1"
 libc.workspace = true
+futures-util = { version = "0.3", default-features = false, features = ["alloc"] }
 pin-project-lite = "0.2"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"

dial9-tokio-telemetry/Cargo.toml

@@ -14,7 +14,7 @@
 //! Duration defaults to 30 seconds. A 3-second warmup precedes measurement.
 
 use dial9_tokio_telemetry::telemetry::{
-    NullWriter, SimpleBinaryWriter, TelemetryGuard, TracedRuntime,
+    NullWriter, SimpleBinaryWriter, TelemetryGuard, TelemetryHandle, TracedRuntime,
 };
 use hdrhistogram::Histogram;
 use std::sync::Arc;
@@ -29,13 +29,13 @@ const WARMUP_SECS: u64 = 3;
 
 // ── Echo server (runs on the traced runtime) ─────────────────────────────────
 
-async fn echo_server(listener: TcpListener) {
+async fn echo_server(listener: TcpListener, handle: Option<TelemetryHandle>) {
     loop {
         let (mut sock, _) = match listener.accept().await {
             Ok(c) => c,
             Err(_) => break,
         };
-        tokio::spawn(async move {
+        let conn = async move {
             let mut buf = [0u8; 256];
             loop {
                 let n = match sock.read(&mut buf).await {
@@ -46,7 +46,12 @@ async fn echo_server(listener: TcpListener) {
                     break;
                 }
             }
-        });
+        };
+        if let Some(h) = &handle {
+            h.spawn(conn);
+        } else {
+            tokio::spawn(conn);
+        }
     }
 }
 
@@ -190,7 +195,8 @@ fn main() {
     let port = server_rt.block_on(async {
         let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
         let port = listener.local_addr().unwrap().port();
-        tokio::spawn(echo_server(listener));
+        let handle = _guard.as_ref().map(|g| g.handle());
+        tokio::spawn(echo_server(listener, handle));
         port
     });
 

dial9-tokio-telemetry/benches/overhead_bench.rs

@@ -1,6 +1,8 @@
 use dial9_tokio_telemetry::telemetry::{
-    TraceReader, analyze_trace, detect_idle_workers, print_analysis,
+    TelemetryEvent, TraceReader, analyze_trace, compute_wake_to_poll_delays, detect_idle_workers,
+    print_analysis,
 };
+use std::collections::HashMap;
 use std::env;
 
 fn main() {
@@ -29,6 +31,127 @@ fn main() {
 
     println!("\n=== Idle Worker Detection ===");
     let idle_periods = detect_idle_workers(&events);
+
+    let delays = compute_wake_to_poll_delays(&events);
+    if !delays.is_empty() {
+        let p50 = delays[delays.len() * 50 / 100];
+        let p99 = delays[delays.len() * 99 / 100];
+        let p999 = delays[delays.len() * 999 / 1000];
+        let max = *delays.last().unwrap();
+        println!("\n=== Wake→Poll Delays ({} samples) ===", delays.len());
+        println!(
+            "  p50: {:.1}µs, p99: {:.1}µs, p99.9: {:.1}µs, max: {:.1}µs",
+            p50 as f64 / 1000.0,
+            p99 as f64 / 1000.0,
+            p999 as f64 / 1000.0,
+            max as f64 / 1000.0,
+        );
+    }
+
+    // Build task_id → spawn_loc_id from PollStart events (more complete than TaskSpawn alone)
+    let mut task_locs: HashMap<u32, u16> = HashMap::new();
+    for e in &events {
+        if let TelemetryEvent::PollStart {
+            task_id,
+            spawn_loc_id,
+            ..
+        } = e
+        {
+            task_locs
+                .entry(task_id.to_u32())
+                .or_insert(spawn_loc_id.as_u16());
+        }
+    }
+    // Also include TaskSpawn mappings from the reader
+    for (task_id, spawn_loc_id) in &reader.task_spawn_locs {
+        task_locs
+            .entry(task_id.to_u32())
+            .or_insert(spawn_loc_id.as_u16());
+    }
+
+    // Count wakes by waker spawn location
+    let mut wakes_by_loc: HashMap<Option<&str>, usize> = HashMap::new();
+    let mut resolved = 0usize;
+    let mut unresolved = 0usize;
+    for e in &events {
+        if let TelemetryEvent::WakeEvent { waker_task_id, .. } = e {
+            let id = waker_task_id.to_u32();
+            if id == 0 {
+                *wakes_by_loc.entry(Some("<non-task context>")).or_default() += 1;
+                resolved += 1;
+            } else if let Some(loc_id) = task_locs.get(&id) {
+                let loc = reader
+                    .spawn_locations
+                    .get(&dial9_tokio_telemetry::telemetry::SpawnLocationId::from_u16(*loc_id));
+                *wakes_by_loc.entry(loc.map(|s| s.as_str())).or_default() += 1;
+                resolved += 1;
+            } else {
+                unresolved += 1;
+            }
+        }
+    }
+    if resolved + unresolved > 0 {
+        println!(
+            "\n=== Waker Identity ({} resolved, {} unresolved of {} total tasks in trace) ===",
+            resolved,
+            unresolved,
+            task_locs.len()
+        );
+
+        // Debug: show some unresolved waker task IDs vs known task IDs
+        let mut unresolved_ids: HashMap<u32, usize> = HashMap::new();
+        for e in &events {
+            if let TelemetryEvent::WakeEvent { waker_task_id, .. } = e {
+                let id = waker_task_id.to_u32();
+                if id != 0 && !task_locs.contains_key(&id) {
+                    *unresolved_ids.entry(id).or_default() += 1;
+                }
+            }
+        }
+        let mut top_unresolved: Vec<_> = unresolved_ids.into_iter().collect();
+        top_unresolved.sort_by(|a, b| b.1.cmp(&a.1));
+        println!("  Top unresolved waker IDs:");
+        for (id, count) in top_unresolved.iter().take(5) {
+            println!("    0x{:08x} ({}) — {} wakes", id, id, count);
+        }
+        println!("  Known task IDs (sample):");
+        for (id, loc_id) in task_locs.iter().take(5) {
+            let loc = reader
+                .spawn_locations
+                .get(&dial9_tokio_telemetry::telemetry::SpawnLocationId::from_u16(*loc_id));
+            println!(
+                "    0x{:08x} ({}) — {}",
+                id,
+                id,
+                loc.map(|s| s.as_str()).unwrap_or("?")
+            );
+        }
+        println!(
+            "  task_spawn_locs from reader: {} entries",
+            reader.task_spawn_locs.len()
+        );
+        println!("  task_locs from PollStart: {} entries", task_locs.len());
+        // Check: are the unresolved IDs in task_spawn_locs?
+        for (id, count) in top_unresolved.iter().take(5) {
+            let in_spawn = reader
+                .task_spawn_locs
+                .contains_key(&dial9_tokio_telemetry::telemetry::TaskId::from_u32(*id));
+            println!(
+                "    0x{:08x}: in task_spawn_locs={}, in task_locs={}, wakes={}",
+                id,
+                in_spawn,
+                task_locs.contains_key(id),
+                count
+            );
+        }
+        let mut by_count: Vec<_> = wakes_by_loc.into_iter().collect();
+        by_count.sort_by(|a, b| b.1.cmp(&a.1));
+        for (loc, count) in &by_count {
+            let name = loc.unwrap_or("<unknown>");
+            println!("  {:>8} wakes from {}", count, name);
+        }
+    }
+
     if idle_periods.is_empty() {
         println!("No significant idle periods detected with work in queue");
     } else {

dial9-tokio-telemetry/examples/analyze_trace.rs

@@ -1,4 +1,5 @@
 pub mod telemetry;
+pub mod traced;
 
 use arc_swap::ArcSwap;
 use pin_project_lite::pin_project;

dial9-tokio-telemetry/src/lib.rs

@@ -192,6 +192,7 @@ pub fn analyze_trace(events: &[TelemetryEvent]) -> TraceAnalysis {
                 stats.max_sched_wait_ns = stats.max_sched_wait_ns.max(*sched_wait_delta_nanos);
             }
             TelemetryEvent::SpawnLocationDef { .. } | TelemetryEvent::TaskSpawn { .. } => {}
+            TelemetryEvent::WakeEvent { .. } => {}
         }
     }
 
@@ -209,6 +210,51 @@ pub fn analyze_trace(events: &[TelemetryEvent]) -> TraceAnalysis {
     }
 }
 
+/// Compute wake→poll scheduling delays from wake events and poll starts.
+/// Returns a sorted vec of delay durations in nanoseconds.
+pub fn compute_wake_to_poll_delays(events: &[TelemetryEvent]) -> Vec<u64> {
+    // Index: task_id → sorted vec of wake timestamps
+    let mut wakes_by_task: HashMap<TaskId, Vec<u64>> = HashMap::new();
+    for e in events {
+        if let TelemetryEvent::WakeEvent {
+            timestamp_nanos,
+            woken_task_id,
+            ..
+        } = e
+        {
+            wakes_by_task
+                .entry(*woken_task_id)
+                .or_default()
+                .push(*timestamp_nanos);
+        }
+    }
+    for v in wakes_by_task.values_mut() {
+        v.sort_unstable();
+    }
+
+    let mut delays = Vec::new();
+    for e in events {
+        if let TelemetryEvent::PollStart {
+            timestamp_nanos,
+            task_id,
+            ..
+        } = e
+            && let Some(wakes) = wakes_by_task.get(task_id)
+        {
+            // Binary search for last wake <= poll start
+            let idx = wakes.partition_point(|&t| t <= *timestamp_nanos);
+            if idx > 0 {
+                let delay = timestamp_nanos - wakes[idx - 1];
+                if delay > 0 && delay < 1_000_000_000 {
+                    delays.push(delay);
+                }
+            }
+        }
+    }
+    delays.sort_unstable();
+    delays
+}
+
 /// An active period between WorkerUnpark and WorkerPark, with scheduling ratio.
 #[derive(Debug)]
 pub struct ActivePeriod {

dial9-tokio-telemetry/src/telemetry/analysis.rs

@@ -68,6 +68,13 @@ pub enum TelemetryEvent {
         task_id: TaskId,
         spawn_loc_id: SpawnLocationId,
     },
+    WakeEvent {
+        #[serde(rename = "timestamp_ns")]
+        timestamp_nanos: u64,
+        waker_task_id: TaskId,
+        woken_task_id: TaskId,
+        target_worker: u8,
+    },
 }
 
 impl TelemetryEvent {
@@ -90,6 +97,9 @@ impl TelemetryEvent {
             }
             | TelemetryEvent::QueueSample {
                 timestamp_nanos, ..
+            }
+            | TelemetryEvent::WakeEvent {
+                timestamp_nanos, ..
             } => Some(*timestamp_nanos),
             TelemetryEvent::SpawnLocationDef { .. } | TelemetryEvent::TaskSpawn { .. } => None,
         }
@@ -104,7 +114,8 @@ impl TelemetryEvent {
             | TelemetryEvent::WorkerUnpark { worker_id, .. } => Some(*worker_id),
             TelemetryEvent::QueueSample { .. }
             | TelemetryEvent::SpawnLocationDef { .. }
-            | TelemetryEvent::TaskSpawn { .. } => None,
+            | TelemetryEvent::TaskSpawn { .. }
+            | TelemetryEvent::WakeEvent { .. } => None,
         }
     }
 
@@ -152,6 +163,12 @@ pub enum RawEvent {
         task_id: crate::telemetry::task_metadata::TaskId,
         location: &'static std::panic::Location<'static>,
     },
+    WakeEvent {
+        timestamp_nanos: u64,
+        waker_task_id: crate::telemetry::task_metadata::TaskId,
+        woken_task_id: crate::telemetry::task_metadata::TaskId,
+        target_worker: u8,
+    },
 }
 
 /// Read the calling thread's CPU time via `CLOCK_THREAD_CPUTIME_ID`.