Use in-memory buffer and set debug to true

Author: connoryy

Cargo.toml

@@ -45,7 +45,7 @@ path = "tests/e2e/mod.rs"
 # This results in roughly a 5% reduction in performance when compiling locally vs when
 # compiled via the CI pipeline.
 [profile.release]
-debug = false # Do not include debug symbols in the executable.
+debug = true # Do not include debug symbols in the executable.
 
 [profile.bench]
 debug = true
@@ -419,6 +419,7 @@ windows-service = "0.8.0"
 nix = { version = "0.26.2", default-features = false, features = ["socket", "signal"] }
 
 [target.'cfg(target_os = "linux")'.dependencies]
+libc = "0.2.174"
 netlink-packet-utils = "0.5.2"
 netlink-packet-sock-diag = "0.4.2"
 netlink-packet-core = "0.7.0"

src/internal_telemetry/allocations/allocator/token.rs

@@ -9,6 +9,12 @@ use tracing::Span;
 use super::stack::GroupStack;
 use super::tracing::WithAllocationGroup;
 
+#[cfg(all(target_os = "linux", feature = "component-probes"))]
+fn thread_id() -> usize {
+    thread_local! { static TID: usize = unsafe { libc::gettid() } as usize; }
+    TID.with(|t| *t)
+}
+
 thread_local! {
     /// The currently executing allocation token.
     ///
@@ -81,17 +87,18 @@ pub struct AllocationGroupToken {
 impl AllocationGroupToken {
     pub fn enter(&self) {
         _ = LOCAL_ALLOCATION_GROUP_STACK.try_with(|stack| stack.borrow_mut().push(self.id));
-        #[cfg(feature = "component-probes")]
-        crate::internal_telemetry::component_probes::vector_component_enter(self.id.as_raw());
+        #[cfg(all(target_os = "linux", feature = "component-probes"))]
+        crate::internal_telemetry::component_probes::VECTOR_COMPONENT_LABELS
+            [thread_id() % crate::internal_telemetry::component_probes::LABELS_LEN]
+            .store(self.id.as_raw(), std::sync::atomic::Ordering::Relaxed);
     }
 
     pub fn exit(&self) {
         _ = LOCAL_ALLOCATION_GROUP_STACK.try_with(|stack| stack.borrow_mut().pop());
-    }
-
-    pub fn close(&self) {
-        #[cfg(feature = "component-probes")]
-        crate::internal_telemetry::component_probes::vector_component_exit();
+        #[cfg(all(target_os = "linux", feature = "component-probes"))]
+        crate::internal_telemetry::component_probes::VECTOR_COMPONENT_LABELS
+            [thread_id() % crate::internal_telemetry::component_probes::LABELS_LEN]
+            .store(0, std::sync::atomic::Ordering::Relaxed);
     }
 }
 

src/internal_telemetry/allocations/allocator/tracing.rs

@@ -69,15 +69,6 @@ where
         }
     }
 
-    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
-        #[cfg(feature = "component-probes")]
-        if let Some(span_ref) = ctx.span(&id) {
-            if let Some(token) = span_ref.extensions().get::<AllocationGroupToken>() {
-                token.close();
-            }
-        }
-    }
-
     unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
         match id {
             id if id == TypeId::of::<Self>() => Some(addr_of!(self).cast::<()>()),

src/internal_telemetry/allocations/mod.rs

@@ -206,6 +206,9 @@ pub fn acquire_allocation_group_id(
                     group_id.as_raw(),
                     id_bytes.as_ptr(),
                     id_bytes.len(),
+                    crate::internal_telemetry::component_probes::VECTOR_COMPONENT_LABELS.as_ptr()
+                        as *const u8,
+                    crate::internal_telemetry::component_probes::LABELS_LEN,
                 );
             }
 

src/internal_telemetry/component_probes.rs

@@ -1,38 +1,45 @@
+use std::sync::atomic::AtomicU8;
+
+pub const LABELS_LEN: usize = 4096;
+
+/// Shared-memory array indexed by `tid % LABELS_LEN`.
+///
+/// Rust writes the component's group ID on enter and 0 on exit (~1ns, no kernel
+/// involvement).  bpftrace reads this array on a fixed-rate profile timer
+/// (`profile:hz:997`) to attribute stack samples to components.
+///
+/// 4096 entries = 4KB.  TID collision requires two Vector threads whose TIDs
+/// differ by exactly 4096 — effectively impossible for a process with ~8-32
+/// threads.
+#[no_mangle]
+pub static VECTOR_COMPONENT_LABELS: [AtomicU8; LABELS_LEN] = {
+    const ZERO: AtomicU8 = AtomicU8::new(0);
+    [ZERO; LABELS_LEN]
+};
+
 /// Uprobe attachment point: called once per component at startup to register
 /// the mapping from allocation group ID to component name.
 ///
 /// bpftrace attaches `uprobe:BINARY:vector_register_component` here at probe
-/// startup to build a `group_id → component_id` lookup table. At sample time,
-/// bpftrace reads the thread-local allocation group stack to get the current
-/// group ID and looks it up in that table.
+/// startup to build a `group_id -> component_id` lookup table.  It also
+/// captures `labels_ptr` / `labels_len` on the first call so the profile
+/// handler can read the shared-memory array at runtime (ASLR makes the
+/// compile-time address unusable).
 ///
 /// Arguments follow the C ABI so bpftrace can read them reliably:
 ///   arg0 = group_id (u8)
 ///   arg1/arg2 = component_id (ptr, len)
+///   arg3/arg4 = labels array (ptr, len)
 ///
 /// `black_box` prevents LTO from eliding the call site.
 #[no_mangle]
 #[inline(never)]
-pub extern "C" fn vector_register_component(id: u8, name_ptr: *const u8, name_len: usize) {
-    std::hint::black_box((id, name_ptr, name_len));
-}
-
-/// Uprobe attachment point: called on every allocation group enter (i.e. every Tokio task poll
-/// that enters a component's group). bpftrace attaches here to record the active component ID
-/// on the current thread without needing span machinery.
-///
-/// Arguments follow the C ABI so bpftrace can read them reliably:
-///   arg0 = group_id (u8)
-#[no_mangle]
-#[inline(never)]
-pub extern "C" fn vector_component_enter(id: u8) {
-    std::hint::black_box(id);
-}
-
-/// Uprobe attachment point: called on every allocation group exit (i.e. every Tokio task poll
-/// that exits a component's group). bpftrace attaches here to clear the active component ID.
-#[no_mangle]
-#[inline(never)]
-pub extern "C" fn vector_component_exit() {
-    std::hint::black_box(0u8);
+pub extern "C" fn vector_register_component(
+    id: u8,
+    name_ptr: *const u8,
+    name_len: usize,
+    labels_ptr: *const u8,
+    labels_len: usize,
+) {
+    std::hint::black_box((id, name_ptr, name_len, labels_ptr, labels_len));
 }