feat: emit and show self-wake counts for tasks (console-rs#112)

This counts all wakes of a task that happen while the task is "currently
polled". This could be from `task::yield_now()`, from `tokio::coop`, or
any other pattern that triggers the task to wake even though it's
currently being polled.

This implementation builds upon @seanmonstar's branch console-rs#55. However,
rather than detecting self-wakes by simply looking at _whether_ the task
is being polled at all, we instead detect them by checking whether the
task's span is on the current thread's stack of entered spans. This
approach should _not_ produce self-positives when a task is woken by
another thread while it is currently entered (as discussed in
tokio-rs/console#55 (comment)).

Testing this with the "burn" subcommand for the example app, which
spawns one task that does a lot of self-wakes via `task::yield`,
indicates that it is, in fact, working properly. Here are the stats for the
`burn` task:
![image](https://user-images.githubusercontent.com/2796466/132136922-2d0d9fbd-e8eb-475c-8ab8-855f8fce5127.png)
Note the very large number of self-wakes.

In comparison, here are the task stats for a "normal" task, which does
not wake itself (and is instead woken by timers):
![image](https://user-images.githubusercontent.com/2796466/132136964-514d7560-074a-41b1-9d6b-621c0dfac935.png)
Note that no self-wakes stat is displayed.

Closes console-rs#55.

Signed-off-by: Eliza Weisman <[email protected]>
Co-authored-by: Sean McArthur <[email protected]>

Author: hawkw

console-subscriber/src/aggregator/mod.rs

@@ -184,6 +184,7 @@ struct TaskStats {
     wakes: u64,
     waker_clones: u64,
     waker_drops: u64,
+    self_wakes: u64,
     last_wake: Option<SystemTime>,
 
     poll_times_histogram: Histogram<u64>,
@@ -273,6 +274,7 @@ impl Default for TaskStats {
             wakes: 0,
             waker_clones: 0,
             waker_drops: 0,
+            self_wakes: 0,
             last_wake: None,
             // significant figures should be in the [0-5] range and memory usage
             // grows exponentially with higher a sigfig
@@ -663,10 +665,16 @@ impl Aggregator {
                 // "wasted" waker ops, but we'll leave that for another time.
                 if let Some(mut task_stats) = self.task_stats.update(&id) {
                     match op {
-                        WakeOp::Wake | WakeOp::WakeByRef => {
+                        WakeOp::Wake { self_wake } | WakeOp::WakeByRef { self_wake } => {
                             task_stats.wakes += 1;
                             task_stats.last_wake = Some(at);
 
+                            // If the  task has woken itself, increment the
+                            // self-wake count.
+                            if self_wake {
+                                task_stats.self_wakes += 1;
+                            }
+
                             // Note: `Waker::wake` does *not* call the `drop`
                             // implementation, so waking by value doesn't
                             // trigger a drop event. so, count this as a `drop`
@@ -675,7 +683,7 @@ impl Aggregator {
                             //
                             // see
                             // https://github.com/rust-lang/rust/blob/673d0db5e393e9c64897005b470bfeb6d5aec61b/library/core/src/task/wake.rs#L211-L212
-                            if let WakeOp::Wake = op {
+                            if let WakeOp::Wake { .. } = op {
                                 task_stats.waker_drops += 1;
                             }
                         }
@@ -866,6 +874,7 @@ impl ToProto for TaskStats {
             total_time: total_time(self.created_at, self.closed_at).map(Into::into),
             wakes: self.wakes,
             waker_clones: self.waker_clones,
+            self_wakes: self.self_wakes,
             waker_drops: self.waker_drops,
             last_wake: self.last_wake.map(Into::into),
         }

console-subscriber/src/lib.rs

@@ -174,8 +174,8 @@ enum AttributeUpdateOp {
 
 #[derive(Clone, Debug, Copy, Serialize)]
 enum WakeOp {
-    Wake,
-    WakeByRef,
+    Wake { self_wake: bool },
+    WakeByRef { self_wake: bool },
     Clone,
     Drop,
 }
@@ -398,7 +398,17 @@ where
             let at = SystemTime::now();
             let mut visitor = WakerVisitor::default();
             event.record(&mut visitor);
-            if let Some((id, op)) = visitor.result() {
+            if let Some((id, mut op)) = visitor.result() {
+                if op.is_wake() {
+                    // Are we currently inside the task's span? If so, the task
+                    // has woken itself.
+                    let self_wake = self
+                        .current_spans
+                        .get()
+                        .map(|spans| spans.borrow().iter().any(|span| span == &id))
+                        .unwrap_or(false);
+                    op = op.self_wake(self_wake);
+                }
                 self.send(Event::Waker { id, op, at });
             }
             // else unknown waker event... what to do? can't trace it from here...
@@ -618,3 +628,18 @@ impl proto::instrument::instrument_server::Instrument for Server {
         Ok(tonic::Response::new(proto::instrument::ResumeResponse {}))
     }
 }
+
+impl WakeOp {
+    /// Returns `true` if `self` is a `Wake` or `WakeByRef` event.
+    fn is_wake(self) -> bool {
+        matches!(self, Self::Wake { .. } | Self::WakeByRef { .. })
+    }
+
+    fn self_wake(self, self_wake: bool) -> Self {
+        match self {
+            Self::Wake { .. } => Self::Wake { self_wake },
+            Self::WakeByRef { .. } => Self::WakeByRef { self_wake },
+            x => x,
+        }
+    }
+}

console-subscriber/src/stack.rs

@@ -49,6 +49,12 @@ impl SpanStack {
         false
     }
 
+    pub(crate) fn iter(&self) -> impl Iterator<Item = &Id> {
+        self.stack
+            .iter()
+            .filter_map(|ContextId { id, duplicate }| if *duplicate { None } else { Some(id) })
+    }
+
     pub(crate) fn stack(&self) -> &Vec<ContextId> {
         &self.stack
     }

console-subscriber/src/visitors.rs

@@ -243,8 +243,8 @@ impl Visit for WakerVisitor {
     fn record_str(&mut self, field: &tracing_core::Field, value: &str) {
         if field.name() == "op" {
             self.op = Some(match value {
-                Self::WAKE => WakeOp::Wake,
-                Self::WAKE_BY_REF => WakeOp::WakeByRef,
+                Self::WAKE => WakeOp::Wake { self_wake: false },
+                Self::WAKE_BY_REF => WakeOp::WakeByRef { self_wake: false },
                 Self::CLONE => WakeOp::Clone,
                 Self::DROP => WakeOp::Drop,
                 _ => return,

console/src/tasks.rs

@@ -101,6 +101,8 @@ struct Stats {
 
     /// The timestamp of when the task was last woken.
     last_wake: Option<SystemTime>,
+    /// Total number of times the task has woken itself.
+    self_wakes: u64,
 }
 
 #[derive(Debug)]
@@ -372,6 +374,11 @@ impl Task {
         self.stats.wakes
     }
 
+    /// Returns the total number of times this task has woken itself.
+    pub(crate) fn self_wakes(&self) -> u64 {
+        self.stats.self_wakes
+    }
+
     fn update(&mut self) {
         let completed = self.stats.total.is_some() && self.completed_for == 0;
         if completed {
@@ -420,6 +427,7 @@ impl From<proto::tasks::Stats> for Stats {
             waker_clones: pb.waker_clones,
             waker_drops: pb.waker_drops,
             last_wake: pb.last_wake.map(|v| v.try_into().unwrap()),
+            self_wakes: pb.self_wakes,
         }
     }
 }

console/src/view/task.rs

@@ -131,6 +131,12 @@ impl TaskView {
             Span::from(format!("{} times", task.wakes())),
         ];
 
+        if task.self_wakes() > 0 {
+            wakeups.push(Span::raw(", "));
+            wakeups.push(bold("Self Wakes: "));
+            wakeups.push(Span::from(format!("{} times", task.self_wakes())));
+        }
+
         // If the task has been woken, add the time since wake to its stats as well.
         if let Some(since) = task.since_wake(now) {
             wakeups.push(Span::raw(", "));

proto/tasks.proto

@@ -104,4 +104,6 @@ message Stats {
     optional google.protobuf.Timestamp last_wake = 6;
     // Contains task poll statistics.
     common.PollStats poll_stats = 7;
+    // The total number of times this task has woken itself.
+    uint64 self_wakes = 8;
 }