feat(telemetry): allow extra arguments with span end probes

examples/README.md

@@ -41,16 +41,22 @@ Hit Ctrl-C to end and print histograms.
 ^C
 
 @attr_task_ms:
-[64, 128)             23 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+[64, 128)             19 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+
+@cache_hit_ms:
+[128, 256)             9 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+
+@cache_miss_ms:
+[256, 512)             9 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
 
 @long_task_ms:
-[32, 64)              17 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
-[64, 128)              7 |@@@@@@@@@@@@@@@@@@@@@                               |
+[32, 64)              12 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+[64, 128)              7 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                      |
 
 @short_task_ms:
-[4, 8)                 5 |@@@@@@@@@@@@@@@@@@@@@@@@@@                          |
-[8, 16)                9 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@      |
-[16, 32)              10 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+[4, 8)                 2 |@@@@@@@@@@@                                         |
+[8, 16)                8 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@      |
+[16, 32)               9 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
 ```
 
 

examples/span_with_probe/main.rs

@@ -50,6 +50,29 @@ async fn attr_task(iter: u64) {
     tokio::time::sleep(Duration::from_millis(100 + iter % 20)).await;
 }
 
+/// Showcase `probe_args`: one span that serves alternating cache hits and
+/// misses, passing a hit flag as arg1 so the tracer can bucket the probe's
+/// durations by outcome. Hits are fast (200-220ms), misses slow (300-320ms).
+async fn cache_task(iter: u64) {
+    const HIT: u64 = 1;
+    const MISS: u64 = 0;
+
+    let (hit, base_ms) = if iter % 2 == 0 {
+        (HIT, 200)
+    } else {
+        (MISS, 300)
+    };
+
+    let work_fut = async move {
+        tokio::time::sleep(Duration::from_millis(base_ms + iter % 20)).await;
+    };
+
+    span_with_probe!("example::cache_task", probe_args = [hit])
+        .into_context()
+        .apply(work_fut)
+        .await;
+}
+
 #[tokio::main(flavor = "current_thread")]
 async fn main() {
     println!("pid {}", std::process::id());
@@ -59,6 +82,7 @@ async fn main() {
         short_task(iter).await;
         long_task(iter).await;
         attr_task(iter).await;
+        cache_task(iter).await;
 
         print!("\riteration {iter}");
         std::io::stdout().flush().unwrap();

examples/span_with_probe/span_durations.bt

@@ -1,15 +1,16 @@
 /*
  * Duration histograms for the span_with_probe example's spans.
  * `attr_task` is instrumented with `span_fn`'s `end_probe = true`
- * option, the other two with the `span_with_probe!` macro.
+ * option, the others with the `span_with_probe!` macro.
  *
  * Run from the repository root (the probe path is relative to the cwd):
  *
  *   cargo run --example span_with_probe
  *   sudo bpftrace examples/span_with_probe/span_durations.bt -p <pid>
  *
  * Probes receive the span duration in nanoseconds as arg0; recorded here in
- * milliseconds.
+ * milliseconds. `cache_task` additionally passes a hit flag as arg1
+ * (probe_args), used to bucket its durations by outcome.
  */
 BEGIN {
     printf("Attaching to span end probes, durations in milliseconds...\n");
@@ -27,3 +28,12 @@ usdt:target/debug/examples/span_with_probe:foundations:span_end__example__long_t
 usdt:target/debug/examples/span_with_probe:foundations:span_end__example__attr_task {
     @attr_task_ms = hist((uint64)arg0 / 1000000);
 }
+
+/* One probe, two latency populations told apart by arg1 (the hit flag). */
+usdt:target/debug/examples/span_with_probe:foundations:span_end__example__cache_task /arg1 == 1/ {
+    @cache_hit_ms = hist((uint64)arg0 / 1000000);
+}
+
+usdt:target/debug/examples/span_with_probe:foundations:span_end__example__cache_task /arg1 == 0/ {
+    @cache_miss_ms = hist((uint64)arg0 / 1000000);
+}

foundations-macros/src/lib.rs

@@ -44,16 +44,17 @@ pub fn span_fn(args: TokenStream, item: TokenStream) -> TokenStream {
 /// semaphore is non-zero, the span start timestamp is recorded in the span
 /// state (regardless of span sampling), and the per-span `probe_end` function's
 /// address is stored alongside it. When the last clone of the span drops,
-/// `probe_end` is called with the span duration in nanoseconds, executing the
-/// NOP whose address the `stapsdt` ELF note publishes as the
-/// `span_end__<sanitized span name>` probe location.
+/// `probe_end` is called with the span duration in nanoseconds followed by the
+/// `probe_args` values, executing the NOP whose address the `stapsdt` ELF note
+/// publishes as the `span_end__<sanitized span name>` probe location. Tracers
+/// see the duration as arg0 and each `probe_args` value after it.
 ///
 /// # Example
 ///
 /// ```rust,ignore
 /// use foundations::telemetry::tracing::span_with_probe;
 ///
-/// span_with_probe!("http::client::send_request", usdt_provider = "myapp")
+/// span_with_probe!("http::client::send_request", usdt_provider = "myapp", probe_args = [opcode as u64])
 ///      .into_context()
 ///      .apply(do_exchange())
 ///      .await
@@ -65,6 +66,9 @@ pub fn span_fn(args: TokenStream, item: TokenStream) -> TokenStream {
 ///   environment variable at compile time — settable per project via `[env]`
 ///   in `.cargo/config.toml` — or `"foundations"` when unset); must be
 ///   non-empty and must not contain `:`
+/// - `probe_args = [expr, ...]` (defaults to `[]`); up to 3 opaque `u64`
+///   values passed to the tracer as the probe's arguments after the duration
+///   (arg0), each evaluated only when the probe is armed
 #[proc_macro]
 pub fn span_with_probe(input: TokenStream) -> TokenStream {
     span_with_probe::expand(input)

foundations-macros/src/span_fn.rs

@@ -296,7 +296,7 @@ fn end_probe_setup(args: &Args) -> Option<TokenStream2> {
         .as_ref()
         .expect("provider defaulted and validated during parse");
 
-    let probe_setup = span_with_probe::probe_setup(span_name, usdt_provider);
+    let probe_setup = span_with_probe::probe_setup(span_name, usdt_provider, None);
     let track_env = span_with_probe::track_provider_env();
 
     Some(quote!(
@@ -791,7 +791,9 @@ mod tests {
         assert!(actual.contains(
             "let mut __span = :: foundations :: telemetry :: tracing :: span (\"sync_span\") ;"
         ));
-        assert!(actual.contains("if enabled { __span . __arm_probe (span_end_probe) ; }"));
+        assert!(
+            actual.contains("__span . __arm_probe (span_end_probe , [0u64 , 0u64 , 0u64 , 0u64]")
+        );
         assert!(actual.contains(".asciz \\\"span_end__sync_span\\\""));
         assert!(actual.contains(".asciz \\\"foundations\\\""));
         // Probe arming is linux/x86_64-only.
@@ -817,7 +819,9 @@ mod tests {
         assert!(actual.contains(
             "let mut __span = :: foundations :: telemetry :: tracing :: span (\"async_span\") ;"
         ));
-        assert!(actual.contains("if enabled { __span . __arm_probe (span_end_probe) ; }"));
+        assert!(
+            actual.contains("__span . __arm_probe (span_end_probe , [0u64 , 0u64 , 0u64 , 0u64]")
+        );
         assert!(actual.contains("__span . into_context () . apply (async move"));
     }
 
@@ -838,7 +842,9 @@ mod tests {
         assert!(actual.contains(
             "let mut __span = :: foundations :: telemetry :: tracing :: dual_span (\"user_span\") ;"
         ));
-        assert!(actual.contains("if enabled { __span . __arm_probe (span_end_probe) ; }"));
+        assert!(
+            actual.contains("__span . __arm_probe (span_end_probe , [0u64 , 0u64 , 0u64 , 0u64]")
+        );
     }
 
     #[test]
@@ -858,7 +864,9 @@ mod tests {
         assert!(actual.contains(
             "let mut __span = :: foo :: bar :: telemetry :: tracing :: span (\"sync_span\") ;"
         ));
-        assert!(actual.contains("if enabled { __span . __arm_probe (span_end_probe) ; }"));
+        assert!(
+            actual.contains("__span . __arm_probe (span_end_probe , [0u64 , 0u64 , 0u64 , 0u64]")
+        );
     }
 
     #[test]

foundations-macros/src/span_with_probe.rs

@@ -4,7 +4,7 @@ use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
 use syn::parse::{Parse, ParseStream};
-use syn::{LitStr, Path, parse_quote};
+use syn::{Expr, LitStr, Path, parse_quote};
 
 struct Args {
     span_name: LitStr,
@@ -18,6 +18,12 @@ pub(crate) struct Options {
 
     #[darling(default = "Options::default_usdt_provider")]
     usdt_provider: LitStr,
+
+    /// Opaque `u64` values passed to the probe as extra arguments (after the
+    /// span duration, which is always arg0); evaluated only when the probe is
+    /// armed.
+    #[darling(default)]
+    probe_args: Option<Expr>,
 }
 
 impl Options {
@@ -85,7 +91,11 @@ fn expand_from_parsed(args: Args) -> TokenStream2 {
     let span_name = &args.span_name;
     let crate_path = &args.options.crate_path;
 
-    let probe_setup = probe_setup(span_name, &args.options.usdt_provider);
+    let probe_setup = probe_setup(
+        span_name,
+        &args.options.usdt_provider,
+        args.options.probe_args.as_ref(),
+    );
     let track_env = track_provider_env();
 
     // The USDT machinery is linux/x86_64-only; elsewhere the macro degrades
@@ -107,10 +117,63 @@ fn expand_from_parsed(args: Args) -> TokenStream2 {
     })
 }
 
+/// Maximum number of u64 probe args (arg0 = duration); mirrors
+/// `MAX_PROBE_ARGS` in foundations. Keep in sync.
+const MAX_PROBE_ARGS: usize = 4;
+
 /// The probe scaffolding shared by `span_with_probe!` and `span_fn`.
-pub(crate) fn probe_setup(span_name: &LitStr, usdt_provider: &LitStr) -> TokenStream2 {
+/// `probe_args` is an optional `[u64; N]` array expression (0..=3 extra
+/// values) exposed to the tracer as the probe's arguments after the span
+/// duration (arg0); it is evaluated only when the probe is armed.
+pub(crate) fn probe_setup(
+    span_name: &LitStr,
+    usdt_provider: &LitStr,
+    probe_args: Option<&Expr>,
+) -> TokenStream2 {
     let probe_name = probe_name(&span_name.value());
-    let template = asm_template(&usdt_provider.value(), &probe_name);
+
+    // Split `probe_args` into its element expressions. An absent option or an
+    // empty array means duration-only; any other arity (up to MAX) is
+    // supported. A non-array expression is rejected.
+    let extra: Vec<Expr> = match probe_args {
+        None => Vec::new(),
+        Some(Expr::Array(arr)) => arr.elems.iter().cloned().collect(),
+        Some(other) => {
+            return syn::Error::new_spanned(
+                other,
+                "probe_args must be an array expression, e.g. `probe_args = [a, b]`",
+            )
+            .to_compile_error();
+        }
+    };
+
+    let nargs = 1 + extra.len();
+    if nargs > MAX_PROBE_ARGS {
+        return syn::Error::new_spanned(
+            probe_args.expect("extra is non-empty only when probe_args is set"),
+            format!(
+                "probe_args supports at most {} extra values ({} total args incl. duration)",
+                MAX_PROBE_ARGS - 1,
+                MAX_PROBE_ARGS
+            ),
+        )
+        .to_compile_error();
+    }
+
+    let template = asm_template(&usdt_provider.value(), &probe_name, nargs);
+
+    // The probe reads its operands from the `args` slice by constant index
+    // (arg0 = duration, then the extras). Each becomes one `-8@<reg>` in the
+    // note; constant indices let `asm!` place each in its own register.
+    let asm_operands = (0..nargs).map(|i| quote!(in(reg) args[#i] as isize,));
+
+    // The args array holds the duration placeholder (index 0, overwritten by
+    // the runtime) followed by the extra expressions, each evaluated once when
+    // the array is built. It is always exactly `MAX_PROBE_ARGS` elements, so
+    // its type is inferred from `__arm_probe`'s signature.
+    let arg_count = nargs as u8;
+    let pad = (0..MAX_PROBE_ARGS - nargs).map(|_| quote!(0u64));
+    let args_array = quote!([ 0u64 #(, #extra )* #(, #pad)* ]);
 
     quote!(
         #[unsafe(link_section = ".probes")]
@@ -119,11 +182,11 @@ pub(crate) fn probe_setup(span_name: &LitStr, usdt_provider: &LitStr) -> TokenSt
         // `#[inline(never)]` keeps the NOP inside this function so the
         // note's address is hit exactly when the span ends.
         #[inline(never)]
-        fn span_end_probe(duration_ns: u64) {
+        fn span_end_probe(args: &[u64]) {
             unsafe {
                 ::core::arch::asm!(#template,
                     sym SEMAPHORE,
-                    in(reg) duration_ns as isize,
+                    #( #asm_operands )*
                     options(readonly, nostack, preserves_flags, att_syntax),
                 )
             }
@@ -132,16 +195,22 @@ pub(crate) fn probe_setup(span_name: &LitStr, usdt_provider: &LitStr) -> TokenSt
         let enabled = unsafe { ::core::ptr::read_volatile(&raw const SEMAPHORE) } != 0;
 
         if enabled {
-            __span.__arm_probe(span_end_probe);
+            __span.__arm_probe(
+                span_end_probe,
+                #args_array,
+                #arg_count,
+            );
         }
     )
 }
 
 /// `stapsdt` note + NOP, adapted from probe-rs' `sdt!` (x86_64, SystemTap
-/// semaphore in `.probes`). The two `{}` operands are the semaphore symbol
-/// and the duration argument.
-fn asm_template(usdt_provider: &str, probe_name: &str) -> String {
+/// semaphore in `.probes`). The `{}` operands are the semaphore symbol and one
+/// `-8@{}` arg descriptor per probe argument (arg0 = duration, then the
+/// caller-chosen extras).
+fn asm_template(usdt_provider: &str, probe_name: &str, nargs: usize) -> String {
     let usdt_provider = sanitize(usdt_provider);
+    let args_desc = vec!["-8@{}"; nargs].join(" ");
 
     format!(
         r#"
@@ -156,7 +225,7 @@ fn asm_template(usdt_provider: &str, probe_name: &str) -> String {
         .8byte {{}}
         .asciz "{usdt_provider}"
         .asciz "{probe_name}"
-        .asciz "-8@{{}}"
+        .asciz "{args_desc}"
 994:    .balign 4
         .popsection
 .ifndef _.stapsdt.base
@@ -203,13 +272,51 @@ mod tests {
         assert!(actual.contains(
             "let mut __span = :: foundations :: telemetry :: tracing :: span (\"http::client::send_request\") ;"
         ));
-        assert!(actual.contains("if enabled { __span . __arm_probe (span_end_probe) ; }"));
+        // Duration-only probe: single arg, args array of just the duration slot.
+        assert!(actual.contains("fn span_end_probe (args : & [u64])"));
+        assert!(actual.contains("in (reg) args [0usize] as isize"));
+        assert!(
+            actual.contains("__span . __arm_probe (span_end_probe , [0u64 , 0u64 , 0u64 , 0u64]")
+        );
+        assert!(actual.contains(", 1u8 ,)"));
         assert!(actual.contains(".asciz \\\"span_end__http__client__send_request\\\""));
+        assert!(actual.contains(".asciz \\\"-8@{}\\\""));
         assert!(actual.contains(".asciz \\\"foundations\\\""));
         // Probe arming is linux/x86_64-only.
         assert!(actual.contains("cfg (all (target_os = \"linux\" , target_arch = \"x86_64\"))"));
     }
 
+    #[test]
+    fn expand_span_with_probe_with_probe_args() {
+        let args = parse_attr! {
+            #[span_with_probe("some::span", probe_args = [u8::from(msg.opcode) as u64, flags])]
+        };
+
+        let actual = expand_from_parsed(args).to_string();
+
+        // Two extra values → three operands and a three-element note.
+        assert!(actual.contains("fn span_end_probe (args : & [u64])"));
+        assert!(actual.contains("in (reg) args [0usize] as isize"));
+        assert!(actual.contains("in (reg) args [1usize] as isize"));
+        assert!(actual.contains("in (reg) args [2usize] as isize"));
+        assert!(actual.contains(
+            "__span . __arm_probe (span_end_probe , [0u64 , u8 :: from (msg . opcode) as u64 , flags , 0u64]"
+        ));
+        assert!(actual.contains(", 3u8 ,)"));
+        assert!(actual.contains(".asciz \\\"-8@{} -8@{} -8@{}\\\""));
+    }
+
+    #[test]
+    fn rejects_too_many_probe_args() {
+        let args = parse_attr! {
+            #[span_with_probe("some::span", probe_args = [1, 2, 3, 4, 5])]
+        };
+
+        let actual = expand_from_parsed(args).to_string();
+
+        assert!(actual.contains("probe_args supports at most 3 extra values"));
+    }
+
     #[test]
     fn expand_span_with_probe_with_crate_path() {
         let args = parse_attr! {