feat(subscriber): support multiple task callsites

console-subscriber/src/callsites.rs

@@ -0,0 +1,51 @@
+use std::{
+    ptr,
+    sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
+};
+use tracing_core::Metadata;
+
+#[derive(Debug, Default)]
+pub(crate) struct Callsites {
+    // In practice each of these will have like, 1-5 callsites in it, max, so
+    // 32 is probably fine...if it ever becomes not fine, we'll fix that.
+    ptrs: [AtomicPtr<Metadata<'static>>; 32],
+    len: AtomicUsize,
+}
+
+impl Callsites {
+    #[track_caller]
+    pub(crate) fn insert(&self, callsite: &'static Metadata<'static>) {
+        // The callsite may already have been inserted, if the callsite cache
+        // was invalidated and is being rebuilt. In that case, don't insert it
+        // again.'
+        if self.contains(callsite) {
+            return;
+        }
+
+        let idx = self.len.fetch_add(1, Ordering::AcqRel);
+        assert!(
+            idx < 64,
+            "you tried to store more than 64 callsites, \
+            time to make the callsite sets bigger i guess \
+            (please open an issue for this)"
+        );
+        self.ptrs[idx]
+            .compare_exchange(
+                ptr::null_mut(),
+                callsite as *const _ as *mut _,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            )
+            .expect("a callsite would have been clobbered by `insert` (this is a bug)");
+    }
+
+    pub(crate) fn contains(&self, callsite: &'static Metadata<'static>) -> bool {
+        let len = self.len.load(Ordering::Acquire);
+        for cs in &self.ptrs[..len + 1] {
+            if ptr::eq(cs.load(Ordering::Acquire), callsite) {
+                return true;
+            }
+        }
+        false
+    }
+}

console-subscriber/src/lib.rs

@@ -4,11 +4,7 @@ use tokio::sync::{mpsc, oneshot};
 
 use std::{
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    ptr,
-    sync::{
-        atomic::{AtomicPtr, Ordering::*},
-        Arc,
-    },
+    sync::Arc,
     time::{Duration, SystemTime},
 };
 use tracing_core::{
@@ -20,18 +16,21 @@ use tracing_core::{
 use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
 
 mod aggregator;
-use aggregator::Aggregator;
 mod builder;
+mod callsites;
+mod init;
+
+use aggregator::Aggregator;
 pub use builder::Builder;
+use callsites::Callsites;
 
-mod init;
 pub use init::{build, init};
 
 pub struct TasksLayer {
-    task_meta: AtomicPtr<Metadata<'static>>,
-    blocking_meta: AtomicPtr<Metadata<'static>>,
     tx: mpsc::Sender<Event>,
     flush: Arc<aggregator::Flush>,
+    spawn_callsites: Callsites,
+    waker_callsites: Callsites,
 }
 
 pub struct Server {
@@ -139,8 +138,8 @@ impl TasksLayer {
         let layer = Self {
             tx,
             flush,
-            task_meta: AtomicPtr::new(ptr::null_mut()),
-            blocking_meta: AtomicPtr::new(ptr::null_mut()),
+            spawn_callsites: Callsites::default(),
+            waker_callsites: Callsites::default(),
         };
         (layer, server)
     }
@@ -159,10 +158,8 @@ impl TasksLayer {
     // chosen by fair die roll, guaranteed to be random :)
     const FLUSH_AT_CAPACITY: usize = 100;
 
-    #[inline(always)]
     fn is_spawn(&self, meta: &'static Metadata<'static>) -> bool {
-        ptr::eq(self.task_meta.load(Relaxed), meta as *const _ as *mut _)
-        // || ptr::eq(self.blocking_meta.load(Relaxed), meta as *const _ as *mut _)
+        self.spawn_callsites.contains(meta)
     }
 
     fn is_id_spawned<S>(&self, id: &span::Id, cx: &Context<'_, S>) -> bool
@@ -216,24 +213,19 @@ where
     S: Subscriber + for<'a> LookupSpan<'a>,
 {
     fn register_callsite(&self, meta: &'static Metadata<'static>) -> subscriber::Interest {
-        if meta.target() == "tokio::task" && meta.name() == "task" {
-            if meta.fields().iter().any(|f| f.name() == "function") {
-                let _ = self.blocking_meta.compare_exchange(
-                    ptr::null_mut(),
-                    meta as *const _ as *mut _,
-                    AcqRel,
-                    Acquire,
-                );
-            } else {
-                let _ = self.task_meta.compare_exchange(
-                    ptr::null_mut(),
-                    meta as *const _ as *mut _,
-                    AcqRel,
-                    Acquire,
-                );
-            }
+        if meta.name() == "runtime.spawn"
+            // back compat until tokio is updated to use the standardized naming
+            // scheme
+            || (meta.name() == "task" && meta.target() == "tokio::task")
+        {
+            self.spawn_callsites.insert(meta);
+        } else if meta.target() == "runtime::waker"
+            // back compat until tokio is updated to use the standardized naming
+            // scheme
+            || meta.target() == "tokio::task::waker"
+        {
+            self.waker_callsites.insert(meta);
         }
-
         self.send(Event::Metadata(meta));
 
         subscriber::Interest::always()
@@ -259,9 +251,7 @@ where
     }
 
     fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
-        let meta = event.metadata();
-        // make faster like spawn metadata pointer check?
-        if meta.target() == "tokio::task::waker" {
+        if self.waker_callsites.contains(event.metadata()) {
             let at = SystemTime::now();
             let mut visitor = WakerVisitor { id: None, op: None };
             event.record(&mut visitor);

console/src/tasks.rs

@@ -44,6 +44,7 @@ pub(crate) struct Task {
     kind: &'static str,
     stats: Stats,
     completed_for: usize,
+    target: Arc<str>,
 }
 
 #[derive(Debug, Default)]
@@ -56,6 +57,7 @@ pub(crate) struct Details {
 #[derive(Debug)]
 pub(crate) struct Metadata {
     field_names: Vec<Arc<str>>,
+    target: Arc<str>,
     //TODO: add more metadata as needed
 }
 
@@ -141,6 +143,20 @@ impl State {
                 proto::tasks::task::Kind::Blocking => "B",
             };
 
+            let meta_id = match task.metadata.as_ref() {
+                Some(id) => id.id,
+                None => {
+                    tracing::warn!(?task, "task has no metadata ID, skipping");
+                    return None;
+                }
+            };
+            let meta = match metas.get(&meta_id) {
+                Some(meta) => meta,
+                None => {
+                    tracing::warn!(?task, meta_id, "no metadata for task, skipping");
+                    return None;
+                }
+            };
             let fields: Vec<Field> = task
                 .fields
                 .drain(..)
@@ -149,11 +165,12 @@ impl State {
                     let name: Option<Arc<str>> = match field_name {
                         proto::field::Name::StrName(n) => Some(n.clone().into()),
                         proto::field::Name::NameIdx(idx) => {
-                            let meta_id = f.metadata_id.as_ref()?;
-                            metas
-                                .get(&meta_id.id)
-                                .and_then(|meta| meta.field_names.get(*idx as usize))
-                                .cloned()
+                            debug_assert_eq!(
+                                f.metadata_id.map(|m| m.id),
+                                Some(meta_id),
+                                "malformed field name: metadata ID mismatch!"
+                            );
+                            meta.field_names.get(*idx as usize).cloned()
                         }
                     };
                     let value = f.value.as_ref().expect("no value").clone().into();
@@ -180,6 +197,7 @@ impl State {
                 kind,
                 stats,
                 completed_for: 0,
+                target: meta.target.clone(),
             };
             task.update();
             let task = Rc::new(RefCell::new(task));
@@ -246,6 +264,10 @@ impl Task {
         &self.id_hex
     }
 
+    pub(crate) fn target(&self) -> &str {
+        &self.target
+    }
+
     pub(crate) fn formatted_fields(&self) -> &[Vec<Span<'static>>] {
         &self.formatted_fields
     }
@@ -373,6 +395,7 @@ impl From<proto::Metadata> for Metadata {
     fn from(pb: proto::Metadata) -> Self {
         Self {
             field_names: pb.field_names.into_iter().map(|n| n.into()).collect(),
+            target: pb.target.into(),
         }
     }
 }

console/src/view/task.rs

@@ -56,7 +56,7 @@ impl TaskView {
             .constraints(
                 [
                     layout::Constraint::Length(1),
-                    layout::Constraint::Length(6),
+                    layout::Constraint::Length(7),
                     layout::Constraint::Length(9),
                     layout::Constraint::Percentage(60),
                 ]
@@ -102,6 +102,7 @@ impl TaskView {
         ]);
 
         let attrs = Spans::from(vec![bold("ID: "), Span::raw(task.id_hex())]);
+        let target = Spans::from(vec![bold("Target: "), Span::raw(task.target())]);
 
         let mut total = vec![
             bold("Total Time: "),
@@ -124,7 +125,7 @@ impl TaskView {
             Span::from(format!("{:.prec$?}", task.idle(now), prec = DUR_PRECISION,)),
         ]);
 
-        let metrics = vec![attrs, total, busy, idle];
+        let metrics = vec![attrs, target, total, busy, idle];
 
         let wakers = Spans::from(vec![
             bold("Current wakers: "),

console/src/view/tasks.rs

@@ -20,8 +20,9 @@ pub(crate) struct List {
 }
 
 impl List {
-    const HEADER: &'static [&'static str] =
-        &["TID", "KIND", "TOTAL", "BUSY", "IDLE", "POLLS", "FIELDS"];
+    const HEADER: &'static [&'static str] = &[
+        "TID", "KIND", "TOTAL", "BUSY", "IDLE", "POLLS", "TARGET", "FIELDS",
+    ];
 
     pub(crate) fn update_input(&mut self, event: input::Event) {
         // Clippy likes to remind us that we could use an `if let` here, since
@@ -83,6 +84,7 @@ impl List {
         // there's room for the unit!)
         const DUR_PRECISION: usize = 4;
         const POLLS_LEN: usize = 5;
+        const MIN_TARGET_LEN: usize = 15;
 
         self.sorted_tasks.extend(state.take_new_tasks());
         self.sort_by.sort(now, &mut self.sorted_tasks);
@@ -115,6 +117,7 @@ impl List {
                     prec = DUR_PRECISION,
                 )),
                 Cell::from(format!("{:>width$}", task.total_polls(), width = POLLS_LEN)),
+                Cell::from(task.target().to_owned()),
                 Cell::from(Spans::from(
                     task.formatted_fields()
                         .iter()
@@ -169,6 +172,7 @@ impl List {
                 layout::Constraint::Min(DUR_LEN as u16),
                 layout::Constraint::Min(DUR_LEN as u16),
                 layout::Constraint::Min(POLLS_LEN as u16),
+                layout::Constraint::Min(MIN_TARGET_LEN as u16),
                 layout::Constraint::Min(10),
             ])
             .highlight_symbol(">> ")