From 7a6a83790d139362a4ef90ccb587dcdceb2cbb83 Mon Sep 17 00:00:00 2001 From: Jakub Sitnicki Date: Fri, 28 Aug 2026 10:22:33 +0200 Subject: [PATCH 1/2] feat(telemetry): add span_with_probe! macro and span_fn's end_probe option Our code bases are already instrumented with tracing spans, but there is no easy way to attach a BPF tracer to a specific span's start/end and measure its duration: one would have to craft custom uprobes at specific code offsets, which shift with every build, and then match span start and end events to each other across await points and thread hops, since an async span can be suspended and resumed on a different thread. span_with_probe! (and its span_fn attribute counterpart, end_probe = true) emit a per-span USDT probe named span_end__ which fires at span end with the span duration as an argument, so a tracer can attach to a single span by name while all other spans pay no cost. The probe is armed via a per-span semaphore only while a tracer is attached and works regardless of span sampling. The USDT provider defaults to the FOUNDATIONS_USDT_PROVIDER environment variable at compile time (settable per project via [env] in .cargo/config.toml), or "foundations" when unset, and can be overridden per call site with usdt_provider = "...". --- foundations-macros/src/lib.rs | 43 +++ foundations-macros/src/span_fn.rs | 305 ++++++++++++++++- foundations-macros/src/span_with_probe.rs | 317 ++++++++++++++++++ foundations/src/telemetry/tracing/internal.rs | 36 +- foundations/src/telemetry/tracing/mod.rs | 105 +++++- 5 files changed, 789 insertions(+), 17 deletions(-) create mode 100644 foundations-macros/src/span_with_probe.rs diff --git a/foundations-macros/src/lib.rs b/foundations-macros/src/lib.rs index 3b280dff..09a5ff14 100644 --- a/foundations-macros/src/lib.rs +++ b/foundations-macros/src/lib.rs @@ -3,6 +3,7 @@ mod info_metric; mod metrics; mod settings; mod span_fn; +mod span_with_probe; mod with_test_telemetry; use proc_macro::TokenStream; @@ -27,6 +28,48 @@ pub fn span_fn(args: TokenStream, item: TokenStream) -> TokenStream { span_fn::expand(args, item) } +/// Like `foundations::telemetry::tracing::span`, plus a per-span USDT probe +/// fired at span end. +/// +/// Probes are only emitted on linux/x86_64; on any other platform the macro +/// degrades to a plain `foundations::telemetry::tracing::span` call. +/// +/// The probe shows up to tracers as +/// `::span_end__`, where +/// sanitization replaces `::` with `__` and any other non-alphanumeric +/// character with `_`. +/// +/// Expands to a dedicated probe semaphore: a `static` in the `.probes` ELF +/// section that the tracer (like bpftrace) increments on attach. When the +/// 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__` probe location. +/// +/// # Example +/// +/// ```rust,ignore +/// use foundations::telemetry::tracing::span_with_probe; +/// +/// span_with_probe!("http::client::send_request", usdt_provider = "myapp") +/// .into_context() +/// .apply(do_exchange()) +/// .await +/// ``` +/// +/// Options: +/// - `crate_path = "..."` (defaults to `::foundations`) +/// - `usdt_provider = "..."` (defaults to the `FOUNDATIONS_USDT_PROVIDER` +/// 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 `:` +#[proc_macro] +pub fn span_with_probe(input: TokenStream) -> TokenStream { + span_with_probe::expand(input) +} + #[proc_macro_attribute] pub fn with_test_telemetry(args: TokenStream, item: TokenStream) -> TokenStream { with_test_telemetry::expand(args, item) diff --git a/foundations-macros/src/span_fn.rs b/foundations-macros/src/span_fn.rs index e9712d81..1d9606bd 100644 --- a/foundations-macros/src/span_fn.rs +++ b/foundations-macros/src/span_fn.rs @@ -1,9 +1,11 @@ use crate::common::parse_optional_trailing_meta_list; +use crate::span_with_probe; use darling::FromMeta; use proc_macro::TokenStream; use proc_macro2::{Span, TokenStream as TokenStream2}; use quote::{ToTokens, quote}; use syn::parse::{Parse, ParseStream}; +use syn::spanned::Spanned as _; use syn::{Block, Expr, ExprCall, ItemFn, LitStr, Path, Signature, Stmt, parse_quote}; const ERR_APPLIED_TO_NON_FN: &str = "`span_fn` macro can only be used on functions"; @@ -50,6 +52,16 @@ struct Options { #[darling(default = "Options::default_user")] user: bool, + + /// Add a per-span USDT probe fired at span end (`span_with_probe!` + /// semantics). + #[darling(default)] + end_probe: bool, + + /// USDT provider for the probe (defaults to `$FOUNDATIONS_USDT_PROVIDER` + /// or `"foundations"`); only valid together with `end_probe = true`. + #[darling(default)] + usdt_provider: Option, } impl Options { @@ -79,7 +91,39 @@ impl Parse for Args { fn parse(input: ParseStream) -> syn::Result { let span_name = input.parse::()?; let meta_list = parse_optional_trailing_meta_list(&input)?; - let options = Options::from_list(&meta_list)?; + let mut options = Options::from_list(&meta_list)?; + + if !options.end_probe + && let Some(provider) = &options.usdt_provider + { + return Err(syn::Error::new( + provider.span(), + "usdt_provider requires end_probe = true", + )); + } + + if options.end_probe { + let provider = options + .usdt_provider + .get_or_insert_with(span_with_probe::Options::default_usdt_provider); + + if !span_with_probe::is_valid_usdt_provider(&provider.value()) { + return Err(syn::Error::new( + provider.span(), + "usdt_provider must be non-empty and must not contain `:`", + )); + } + } + + // The probe name is derived from the span name at compile time. + if options.end_probe + && let SpanName::Const(path) = &span_name + { + return Err(syn::Error::new( + path.span(), + "end_probe spans require a string literal span name", + )); + } Ok(Self { span_name, options }) } @@ -124,14 +168,20 @@ fn expand_from_parsed(args: Args, item_fn: ItemFn) -> TokenStream2 { let body = match asyncness { Some(_) => wrap_with_span(&args, quote!(async move { #block })), None => try_async_trait_fn_rewrite(&args, &block).unwrap_or_else(|| { - let span_name = args.span_name.as_tokens(); - let crate_path = &args.options.crate_path; - let span_ctor = span_ctor(&args.options); - - quote!( - let __span = #crate_path::telemetry::tracing::#span_ctor(#span_name); - #block - ) + let span_expr = span_expr(&args); + + match end_probe_setup(&args) { + Some(probe_setup) => quote!( + #[allow(unused_mut)] + let mut __span = #span_expr; + #probe_setup + #block + ), + None => quote!( + let __span = #span_expr; + #block + ), + } }), }; @@ -197,16 +247,66 @@ fn wrap_with_span(args: &Args, block: TokenStream2) -> TokenStream2 { quote!(apply) }; + let span_expr = span_expr(args); + + match end_probe_setup(args) { + Some(probe_setup) => quote!( + { + #[allow(unused_mut)] + let mut __span = #span_expr; + #probe_setup + __span + .into_context() + .#apply_fn(#block) + .await + } + ), + None => quote!( + #span_expr + .into_context() + .#apply_fn(#block) + .await + ), + } +} + +/// The span-construction expression: a plain `span`/`dual_span` call. +fn span_expr(args: &Args) -> TokenStream2 { let span_name = args.span_name.as_tokens(); let crate_path = &args.options.crate_path; let span_ctor = span_ctor(&args.options); - quote!( - #crate_path::telemetry::tracing::#span_ctor(#span_name) - .into_context() - .#apply_fn(#block) - .await - ) + quote!(#crate_path::telemetry::tracing::#span_ctor(#span_name)) +} + +/// When `end_probe` is enabled, the linux/x86_64-only block that sets up the +/// USDT span-end probe on the just-created `__span`. +fn end_probe_setup(args: &Args) -> Option { + if !args.options.end_probe { + return None; + } + + let SpanName::Str(span_name) = &args.span_name else { + unreachable!("end_probe spans require a string literal span name"); + }; + + let usdt_provider = args + .options + .usdt_provider + .as_ref() + .expect("provider defaulted and validated during parse"); + + let probe_setup = span_with_probe::probe_setup(span_name, usdt_provider); + let track_env = span_with_probe::track_provider_env(); + + Some(quote!( + #track_env + + #[cfg(all(target_os = "linux", target_arch = "x86_64"))] + { + #probe_setup + } + )) } /// The span constructor to call: `dual_span` when `user = true` (internal + parallel user span), @@ -670,4 +770,179 @@ mod tests { assert_eq!(actual, expected); } + + #[test] + fn expand_sync_fn_probe() { + let args = parse_attr! { + #[span_fn("sync_span", end_probe = true)] + }; + + let item_fn = parse_quote! { + fn do_sync() -> io::Result { + do_something_else(); + + Ok("foo".into()) + } + }; + + let actual = expand_from_parsed(args, item_fn).to_string(); + + // Plain span construction, with the probe armed on the scope after. + assert!(actual.contains( + "let mut __span = :: foundations :: telemetry :: tracing :: span (\"sync_span\") ;" + )); + assert!(actual.contains("__span . __arm_probe (enabled , span_end_probe) ;")); + assert!(actual.contains(".asciz \\\"span_end__sync_span\\\"")); + 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_async_fn_probe() { + let args = parse_attr! { + #[span_fn("async_span", end_probe = true)] + }; + + let item_fn = parse_quote! { + async fn do_async() -> io::Result { + do_something_else().await; + + Ok("foo".into()) + } + }; + + let actual = expand_from_parsed(args, item_fn).to_string(); + + assert!(actual.contains( + "let mut __span = :: foundations :: telemetry :: tracing :: span (\"async_span\") ;" + )); + assert!(actual.contains("__span . __arm_probe (enabled , span_end_probe) ;")); + assert!(actual.contains("__span . into_context () . apply (async move")); + } + + #[test] + fn expand_fn_probe_user() { + let args = parse_attr! { + #[span_fn("user_span", end_probe = true, user = true)] + }; + + let item_fn = parse_quote! { + fn do_sync() { + do_something_else(); + } + }; + + let actual = expand_from_parsed(args, item_fn).to_string(); + + assert!(actual.contains( + "let mut __span = :: foundations :: telemetry :: tracing :: dual_span (\"user_span\") ;" + )); + assert!(actual.contains("__span . __arm_probe (enabled , span_end_probe) ;")); + } + + #[test] + fn expand_fn_probe_with_crate_path() { + let args = parse_attr! { + #[span_fn("sync_span", end_probe = true, crate_path = "::foo::bar")] + }; + + let item_fn = parse_quote! { + fn do_sync() { + do_something_else(); + } + }; + + let actual = expand_from_parsed(args, item_fn).to_string(); + + assert!(actual.contains( + "let mut __span = :: foo :: bar :: telemetry :: tracing :: span (\"sync_span\") ;" + )); + assert!(actual.contains("__span . __arm_probe (enabled , span_end_probe) ;")); + } + + #[test] + fn expand_fn_probe_with_usdt_provider() { + let args = parse_attr! { + #[span_fn("some::span", end_probe = true, usdt_provider = "myapp")] + }; + + let item_fn = parse_quote! { + fn do_sync() { + do_something_else(); + } + }; + + let actual = expand_from_parsed(args, item_fn).to_string(); + + // The asm template is a nested string literal, so its quotes are + // escaped in the token stream's string representation. + assert!(actual.contains(".asciz \\\"myapp\\\"")); + assert!(actual.contains(".asciz \\\"span_end__some__span\\\"")); + } + + #[test] + fn expand_fn_probe_with_env_usdt_provider() { + unsafe { std::env::set_var(span_with_probe::USDT_PROVIDER_ENV_VAR, "envapp") }; + + let args = parse_attr! { + #[span_fn("some::span", end_probe = true)] + }; + + let item_fn = parse_quote! { + fn do_sync() { + do_something_else(); + } + }; + + let actual = expand_from_parsed(args, item_fn).to_string(); + + assert!(actual.contains(".asciz \\\"envapp\\\"")); + + unsafe { std::env::remove_var(span_with_probe::USDT_PROVIDER_ENV_VAR) }; + } + + #[test] + fn rejects_usdt_provider_without_probe() { + let tokens = quote! { "some::span", usdt_provider = "myapp" }; + let err = match syn::parse2::(tokens) { + Ok(_) => panic!("usdt_provider without probe unexpectedly accepted"), + Err(err) => err, + }; + + assert!( + err.to_string().contains("requires end_probe = true"), + "{err}" + ); + } + + #[test] + fn rejects_invalid_usdt_provider() { + for provider in ["foo:bar", ""] { + let tokens = quote! { "some::span", end_probe = true, usdt_provider = #provider }; + let err = match syn::parse2::(tokens) { + Ok(_) => panic!("provider {provider:?} unexpectedly accepted"), + Err(err) => err, + }; + + assert!( + err.to_string().contains("usdt_provider"), + "provider {provider:?}: {err}" + ); + } + } + + #[test] + fn rejects_probe_with_const_span_name() { + let tokens = quote! { some::module::SPAN, end_probe = true }; + let err = match syn::parse2::(tokens) { + Ok(_) => panic!("const span name with probe unexpectedly accepted"), + Err(err) => err, + }; + + assert!( + err.to_string().contains("string literal span name"), + "{err}" + ); + } } diff --git a/foundations-macros/src/span_with_probe.rs b/foundations-macros/src/span_with_probe.rs new file mode 100644 index 00000000..9df66938 --- /dev/null +++ b/foundations-macros/src/span_with_probe.rs @@ -0,0 +1,317 @@ +use crate::common::parse_optional_trailing_meta_list; +use darling::FromMeta; +use proc_macro::TokenStream; +use proc_macro2::TokenStream as TokenStream2; +use quote::quote; +use syn::parse::{Parse, ParseStream}; +use syn::{LitStr, Path, parse_quote}; + +struct Args { + span_name: LitStr, + options: Options, +} + +#[derive(FromMeta)] +pub(crate) struct Options { + #[darling(default = "Options::default_crate_path")] + crate_path: Path, + + #[darling(default = "Options::default_usdt_provider")] + usdt_provider: LitStr, +} + +impl Options { + fn default_crate_path() -> Path { + parse_quote!(::foundations) + } + + pub(crate) fn default_usdt_provider() -> LitStr { + let provider = + std::env::var(USDT_PROVIDER_ENV_VAR).unwrap_or_else(|_| "foundations".into()); + + LitStr::new(&provider, proc_macro2::Span::call_site()) + } +} + +/// Environment variable that overrides the default USDT provider at compile +/// time; settable per project via `[env]` in `.cargo/config.toml`. +pub(crate) const USDT_PROVIDER_ENV_VAR: &str = "FOUNDATIONS_USDT_PROVIDER"; + +/// Stable-Rust substitute for the unstable `proc_macro::tracked_env`: the +/// `option_env!` read lands in the consumer crate's dep-info, so Cargo +/// rebuilds the call site (rerunning the macro) when the override changes. +pub(crate) fn track_provider_env() -> TokenStream2 { + let env_var = USDT_PROVIDER_ENV_VAR; + + quote!( + const _: Option<&'static str> = option_env!(#env_var); + ) +} + +impl Parse for Args { + fn parse(input: ParseStream) -> syn::Result { + let span_name = input.parse::()?; + let meta_list = parse_optional_trailing_meta_list(&input)?; + let options = Options::from_list(&meta_list)?; + + let provider = &options.usdt_provider; + if !is_valid_usdt_provider(&provider.value()) { + return Err(syn::Error::new( + provider.span(), + "usdt_provider must be non-empty and must not contain `:`", + )); + } + + Ok(Self { span_name, options }) + } +} + +/// libbpf's SEC("usdt/::") auto-attach syntax is +/// colon-delimited with no quoting mechanism, so `:` is rejected outright +/// (bpftrace would accept it in a quoted field, but libbpf would not). +/// Everything else is sanitized away when the provider is embedded into the +/// GAS `.asciz` directive (see [`sanitize`]). +pub(crate) fn is_valid_usdt_provider(provider: &str) -> bool { + !provider.is_empty() && !provider.contains(':') +} + +pub(crate) fn expand(input: TokenStream) -> TokenStream { + let args = syn::parse_macro_input!(input as Args); + + expand_from_parsed(args).into() +} + +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 track_env = track_provider_env(); + + // The USDT machinery is linux/x86_64-only; elsewhere the macro degrades + // to a plain `tracing::span` (no semaphore, no ELF note). The `cfg` must + // be emitted into the expansion: the macro runs on the build host, so the + // target platform is only known when the call site is compiled. + quote!({ + #track_env + + #[allow(unused_mut)] + let mut __span = #crate_path::telemetry::tracing::span(#span_name); + + #[cfg(all(target_os = "linux", target_arch = "x86_64"))] + { + #probe_setup + } + + __span + }) +} + +/// The probe scaffolding shared by `span_with_probe!` and `span_fn`. +pub(crate) fn probe_setup(span_name: &LitStr, usdt_provider: &LitStr) -> TokenStream2 { + let probe_name = probe_name(&span_name.value()); + let template = asm_template(&usdt_provider.value(), &probe_name); + + quote!( + #[unsafe(link_section = ".probes")] + static mut SEMAPHORE: u16 = 0; + + // `#[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) { + unsafe { + ::core::arch::asm!(#template, + sym SEMAPHORE, + in(reg) duration_ns as isize, + options(readonly, nostack, preserves_flags, att_syntax), + ) + } + } + + let enabled = unsafe { ::core::ptr::read_volatile(&raw const SEMAPHORE) } != 0; + + __span.__arm_probe(enabled, span_end_probe); + ) +} + +/// `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 { + let usdt_provider = sanitize(usdt_provider); + + format!( + r#" +990: nop + .pushsection .note.stapsdt,"?","note" + .balign 4 + .4byte 992f-991f, 994f-993f, 3 +991: .asciz "stapsdt" +992: .balign 4 +993: .8byte 990b + .8byte _.stapsdt.base + .8byte {{}} + .asciz "{usdt_provider}" + .asciz "{probe_name}" + .asciz "-8@{{}}" +994: .balign 4 + .popsection +.ifndef _.stapsdt.base + .pushsection .stapsdt.base,"aGR","progbits",.stapsdt.base,comdat + .weak _.stapsdt.base + .hidden _.stapsdt.base +_.stapsdt.base: .space 1 + .size _.stapsdt.base, 1 + .popsection +.endif"# + ) +} + +/// The USDT probe name for a span: `span_end__` +/// (see [`sanitize`]; each `:` of a `::` path separator becomes a `_`). +fn probe_name(span_name: &str) -> String { + format!("span_end__{}", sanitize(span_name)) +} + +/// Sanitizes a string embedded in the `stapsdt` note's GAS `.asciz` +/// directives (probe provider and name): any character that is not an ASCII +/// alphanumeric becomes `_`. The result needs no escaping for GAS and +/// is always addressable by libbpf's colon-delimited +/// `SEC("usdt/::")` auto-attach syntax. +fn sanitize(s: &str) -> String { + s.chars() + .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::test_utils::parse_attr; + + #[test] + fn expand_span_with_probe() { + let args = parse_attr! { + #[span_with_probe("http::client::send_request")] + }; + + let actual = expand_from_parsed(args).to_string(); + + assert!(actual.contains( + "let mut __span = :: foundations :: telemetry :: tracing :: span (\"http::client::send_request\") ;" + )); + assert!(actual.contains("__span . __arm_probe (enabled , span_end_probe) ;")); + assert!(actual.contains(".asciz \\\"span_end__http__client__send_request\\\"")); + 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_crate_path() { + let args = parse_attr! { + #[span_with_probe("sync_span", crate_path = "::foo::bar")] + }; + + let actual = expand_from_parsed(args).to_string(); + + assert!(actual.contains( + "let mut __span = :: foo :: bar :: telemetry :: tracing :: span (\"sync_span\") ;" + )); + } + + #[test] + fn expand_span_with_probe_with_usdt_provider() { + let args = parse_attr! { + #[span_with_probe("some::span", usdt_provider = "myapp")] + }; + + let actual = expand_from_parsed(args).to_string(); + + // The asm template is a nested string literal, so its quotes are + // escaped in the token stream's string representation. + assert!(actual.contains(".asciz \\\"myapp\\\"")); + assert!(actual.contains(".asciz \\\"span_end__some__span\\\"")); + } + + #[test] + fn rejects_invalid_usdt_provider() { + for provider in ["foo:bar", ""] { + let tokens = quote! { "some::span", usdt_provider = #provider }; + let err = match syn::parse2::(tokens) { + Ok(_) => panic!("provider {provider:?} unexpectedly accepted"), + Err(err) => err, + }; + + assert!( + err.to_string().contains("usdt_provider"), + "provider {provider:?}: {err}" + ); + } + } + + #[test] + fn accepts_valid_usdt_provider() { + for provider in ["myapp", "python3.12", "my-app_v2", "foo bar", "foo\"bar"] { + let tokens = quote! { "some::span", usdt_provider = #provider }; + + syn::parse2::(tokens).unwrap(); + } + } + + #[test] + fn usdt_provider_from_env() { + // One test function for all env-var cases: the environment is + // process-global, so spreading these across tests could race under + // a shared-process test runner. + unsafe { std::env::set_var(USDT_PROVIDER_ENV_VAR, "envapp") }; + + let args = parse_attr! { + #[span_with_probe("some::span")] + }; + let actual = expand_from_parsed(args).to_string(); + assert!(actual.contains(".asciz \\\"envapp\\\"")); + + // An invalid override is rejected like an invalid option value. + unsafe { std::env::set_var(USDT_PROVIDER_ENV_VAR, "foo:bar") }; + assert!(syn::parse2::(quote! { "some::span" }).is_err()); + + unsafe { std::env::remove_var(USDT_PROVIDER_ENV_VAR) }; + } + + #[test] + fn sanitizes_probe_name() { + assert_eq!( + probe_name("http::client::send_request"), + "span_end__http__client__send_request" + ); + // A single `:` and any other non-alphanumeric become `_`. + assert_eq!(probe_name("foo:bar"), "span_end__foo_bar"); + assert_eq!(probe_name("foo bar.baz"), "span_end__foo_bar_baz"); + } + + #[test] + fn sanitizes_note_strings() { + // Anything that is not an ASCII alphanumeric becomes `_`, so + // the result is GAS-safe and attachable without any escaping. + assert_eq!(sanitize("foo\nbar\tbaz€"), "foo_bar_baz_"); + assert_eq!(sanitize("foo bar.baz"), "foo_bar_baz"); + assert_eq!(sanitize("a\"b\\c"), "a_b_c"); + } + + #[test] + fn sanitizes_special_chars_in_asm_template() { + let args = parse_attr! { + #[span_with_probe("foo\"bar", usdt_provider = "my\\app")] + }; + + let actual = expand_from_parsed(args).to_string(); + + assert!(actual.contains(r#".asciz \"my_app\""#), "{actual}"); + assert!( + actual.contains(r#".asciz \"span_end__foo_bar\""#), + "{actual}" + ); + } +} diff --git a/foundations/src/telemetry/tracing/internal.rs b/foundations/src/telemetry/tracing/internal.rs index fe7ef969..02ffce97 100644 --- a/foundations/src/telemetry/tracing/internal.rs +++ b/foundations/src/telemetry/tracing/internal.rs @@ -14,6 +14,7 @@ use rand::RngExt as _; use std::borrow::Cow; use std::error::Error; use std::sync::Arc; +use std::time::Instant; pub(crate) type Tracer = cf_rustracing::Tracer, SpanContextState>; @@ -74,6 +75,34 @@ pub(crate) struct SharedSpan { // NOTE: store sampling flag separately, so we don't need to acquire lock // every time we need to check the flag. pub(crate) is_sampled: bool, + /// USDT span probe state, recorded when the span's probe semaphore is + /// non-zero (a tracer is attached), regardless of sampling. Shared by all + /// clones, so the `span_end__*` probe fires exactly once when the last + /// clone of the span drops. + pub(crate) probe: Option>, +} + +/// Probe state for a single span invocation. Dropping it fires the span's +/// `span_end__*` USDT probe with the span duration in nanoseconds. +#[derive(Debug)] +pub(crate) struct SpanProbe { + start: Instant, + end_probe: fn(u64), +} + +impl SpanProbe { + pub(crate) fn new(end_probe: fn(u64)) -> Self { + Self { + start: Instant::now(), + end_probe, + } + } +} + +impl Drop for SpanProbe { + fn drop(&mut self) { + (self.end_probe)(self.start.elapsed().as_nanos() as u64); + } } /// Wraps a span and registers it with the internal harness's `active_roots` for live tracking. @@ -83,6 +112,7 @@ pub(crate) fn shared_span(span: Span) -> SharedSpan { SharedSpan { inner: SharedSpanHandle::new(span), is_sampled, + probe: None, } } @@ -98,7 +128,11 @@ pub(crate) fn user_shared_span(span: Span) -> SharedSpan { SharedSpanHandle::Inactive }; - SharedSpan { inner, is_sampled } + SharedSpan { + inner, + is_sampled, + probe: None, + } } pub fn write_current_span(write_fn: impl FnOnce(&mut Span)) { diff --git a/foundations/src/telemetry/tracing/mod.rs b/foundations/src/telemetry/tracing/mod.rs index 0d06fb02..5b27971f 100644 --- a/foundations/src/telemetry/tracing/mod.rs +++ b/foundations/src/telemetry/tracing/mod.rs @@ -25,7 +25,9 @@ mod output_otlp_uds; mod traceparent; use self::init::TracingHarness; -use self::internal::{SharedSpan, create_span, current_span, shared_span, span_trace_id}; +use self::internal::{ + SharedSpan, SpanProbe, create_span, current_span, shared_span, span_trace_id, +}; #[cfg(feature = "user-tracing")] use self::internal::{ SharedSpanHandle, child_user_span, current_user_span, start_user_trace, user_shared_span, @@ -131,6 +133,24 @@ pub fn get_active_traces() -> String { /// } /// ``` /// +/// # Emitting a USDT probe at span end +/// +/// ``` +/// use foundations::telemetry::tracing; +/// +/// #[tracing::span_fn("foo", end_probe = true)] +/// fn foo() { +/// // Does something... +/// } +/// ``` +/// +/// With `end_probe = true`, the span additionally sets up a per-span USDT probe fired with the span +/// duration when the span ends, with the same semantics as the `span_with_probe!` macro. The span +/// name must be a string literal in this case. The provider defaults to the +/// `FOUNDATIONS_USDT_PROVIDER` environment variable at compile time (settable per project via +/// `[env]` in `.cargo/config.toml`), or `"foundations"` when unset, and can be customized per span +/// with `usdt_provider = "..."`. +/// /// # Renamed or reexported crate /// /// The macro will fail to compile if `foundations` crate is reexported. However, the crate path @@ -168,6 +188,14 @@ pub fn get_active_traces() -> String { /// [async_trait]: https://crates.io/crates/async-trait pub use foundations_macros::span_fn; +/// [`span()`] variant that additionally arms a per-span USDT probe +/// (`span_end__`) fired when the span ends and a tracer has +/// attached to it. +/// +/// See the macro documentation in `foundations-macros` for details. +#[doc(inline)] +pub use foundations_macros::span_with_probe; + /// A handle for the scope in which tracing span is active. /// /// Scope ends when the handle is dropped. @@ -235,6 +263,20 @@ impl SpanScope { ctx } + + /// Arms the span's USDT probe when `record_probe_start` is true (the + /// span's probe semaphore, a `static` in the `.probes` section bumped by + /// the tracer on attach, is non-zero). `end_probe` is the address of a + /// per-span function containing the probe's NOP placeholder; it is called + /// with the span duration in nanoseconds when the last clone of the span + /// drops. Arming records the start timestamp unconditionally with respect + /// to sampling, so probes work even when span tracing is disabled. + #[doc(hidden)] + pub fn __arm_probe(&mut self, record_probe_start: bool, end_probe: fn(u64)) { + if record_probe_start { + self.span.probe = Some(Arc::new(SpanProbe::new(end_probe))); + } + } } /// A handle for the scope in which a user-tracing span is active. @@ -306,6 +348,7 @@ impl UserSpan { span: SharedSpan { inner: SharedSpanHandle::Inactive, is_sampled: false, + probe: None, }, } } @@ -418,6 +461,13 @@ impl DualSpanScope { ctx } + + /// Arms the internal span's USDT probe when `record_probe_start` is true. + /// The user span is never probed. + #[doc(hidden)] + pub fn __arm_probe(&mut self, record_probe_start: bool, end_probe: fn(u64)) { + self.inner.__arm_probe(record_probe_start, end_probe); + } } /// Options for a new trace. @@ -1984,3 +2034,56 @@ mod user_tracing_tests { ); } } + +#[cfg(test)] +mod probe_tests { + use super::span; + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn noop_probe(_duration_ns: u64) {} + + #[test] + fn arm_probe_only_when_enabled() { + // No telemetry context is installed: spans are unsampled. Probe + // arming must not depend on sampling. + let mut scope = span("test::probe"); + assert!(scope.span.probe.is_none()); + + scope.__arm_probe(false, noop_probe); + assert!(scope.span.probe.is_none()); + + scope.__arm_probe(true, noop_probe); + assert!(scope.span.probe.is_some()); + } + + #[test] + fn end_probe_fires_once_when_last_clone_drops() { + static FIRES: AtomicUsize = AtomicUsize::new(0); + + fn counting_probe(_duration_ns: u64) { + FIRES.fetch_add(1, Ordering::Relaxed); + } + + let mut scope = span("test::probe"); + scope.__arm_probe(true, counting_probe); + let clone = scope.span.clone(); + + drop(scope); + assert_eq!(FIRES.load(Ordering::Relaxed), 0); + + drop(clone); + assert_eq!(FIRES.load(Ordering::Relaxed), 1); + } + + #[cfg(feature = "user-tracing")] + #[test] + fn dual_scope_probe_arms_internal_span_only() { + let mut scope = super::dual_span("test::probe"); + assert!(scope.inner.span.probe.is_none()); + // No user trace is active, so no user span is created. + assert!(scope.user.is_none()); + + scope.__arm_probe(true, noop_probe); + assert!(scope.inner.span.probe.is_some()); + } +} From 7ec8942c29b11deb6c5ded9de32d8f6a4a27c0e6 Mon Sep 17 00:00:00 2001 From: Jakub Sitnicki Date: Fri, 28 Aug 2026 10:23:09 +0200 Subject: [PATCH 2/2] feat(examples): add span_with_probe demo workload Runs spans instrumented with span_with_probe! and span_fn's end_probe option in a loop, with distinct delay ranges so their durations can be told apart. Comes with a bpftrace script that attaches to the span end probes and prints duration histograms. --- examples/Cargo.toml | 4 ++ examples/README.md | 30 ++++++++++ examples/span_with_probe/main.rs | 69 ++++++++++++++++++++++ examples/span_with_probe/span_durations.bt | 29 +++++++++ 4 files changed, 132 insertions(+) create mode 100644 examples/span_with_probe/main.rs create mode 100644 examples/span_with_probe/span_durations.bt diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 2cd2f6d8..9c1968b9 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -22,3 +22,7 @@ tokio = { workspace = true, features = ["full"]} [[example]] name = "http_server" path = "http_server/main.rs" + +[[example]] +name = "span_with_probe" +path = "span_with_probe/main.rs" diff --git a/examples/README.md b/examples/README.md index 2612bafa..c19b04f6 100644 --- a/examples/README.md +++ b/examples/README.md @@ -21,5 +21,35 @@ cargo run --example http_server -- --config http_server/example_conf.yaml ``` +## `span_with_probe` + + +Demo workload for per-span USDT probes. Runs spans instrumented with the +`span_with_probe!` macro and `span_fn`'s `end_probe = true` option in a loop; +attach with bpftrace to get duration histograms: + +``` +cargo run --example span_with_probe +sudo bpftrace examples/span_with_probe/span_durations.bt -p +``` + +Sample output: + +``` +Attaching to span end probes, durations in milliseconds... +Hit Ctrl-C to end and print histograms. +^C + +@long_task_ms: +[32, 64) 15 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@| +[64, 128) 12 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ | + +@short_task_ms: +[4, 8) 2 |@@@@@ | +[8, 16) 8 |@@@@@@@@@@@@@@@@@@@@@@@ | +[16, 32) 18 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@| +``` + + diff --git a/examples/span_with_probe/main.rs b/examples/span_with_probe/main.rs new file mode 100644 index 00000000..a8d9fed3 --- /dev/null +++ b/examples/span_with_probe/main.rs @@ -0,0 +1,69 @@ +//! Demo workload for the `span_with_probe!` macro and `span_fn`'s +//! `end_probe = true` option. +//! +//! Instrumented spans with distinct delay ranges run in a loop. Attach +//! with bpftrace to get duration histograms for all spans (from the repo +//! root, so the probe path resolves): +//! +//! ```text +//! bpftrace examples/span_with_probe/span_durations.bt -p +//! ``` +//! +//! No telemetry context is installed: every span is unsampled and tracing is +//! effectively disabled. This is intentional — the USDT probe must fire +//! regardless of span sampling. + +use foundations::telemetry::tracing::{span_fn, span_with_probe}; +use std::io::Write as _; +use std::time::Duration; + +/// Span with a 5-25ms delay range. +async fn short_task(iter: u64) { + let work_fut = async move { + // Simulate variable work so durations show up as a histogram. + tokio::time::sleep(Duration::from_millis(5 + iter % 20)).await; + }; + + span_with_probe!("example::short_task") + .into_context() + .apply(work_fut) + .await; +} + +/// Span with a distinct delay range (50-70ms) so the two can be told apart +/// in a duration histogram. +async fn long_task(iter: u64) { + let work_fut = async move { + tokio::time::sleep(Duration::from_millis(50 + iter % 20)).await; + }; + + span_with_probe!("example::long_task") + .into_context() + .apply(work_fut) + .await; +} + +/// Same probing via `span_fn`'s `end_probe = true` option, with its own +/// delay range (100-120ms). +#[span_fn("example::attr_task", end_probe = true)] +async fn attr_task(iter: u64) { + tokio::time::sleep(Duration::from_millis(100 + iter % 20)).await; +} + +#[tokio::main(flavor = "current_thread")] +async fn main() { + println!("pid {}", std::process::id()); + + let mut iter = 0; + loop { + short_task(iter).await; + long_task(iter).await; + attr_task(iter).await; + + print!("\riteration {iter}"); + std::io::stdout().flush().unwrap(); + + iter += 1; + tokio::time::sleep(Duration::from_millis(100)).await; + } +} diff --git a/examples/span_with_probe/span_durations.bt b/examples/span_with_probe/span_durations.bt new file mode 100644 index 00000000..6adb482c --- /dev/null +++ b/examples/span_with_probe/span_durations.bt @@ -0,0 +1,29 @@ +/* + * 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. + * + * 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 + * + * Probes receive the span duration in nanoseconds as arg0; recorded here in + * milliseconds. + */ +BEGIN { + printf("Attaching to span end probes, durations in milliseconds...\n"); + printf("Hit Ctrl-C to end and print histograms.\n"); +} + +usdt:target/debug/examples/span_with_probe:foundations:span_end__example__short_task { + @short_task_ms = hist((uint64)arg0 / 1000000); +} + +usdt:target/debug/examples/span_with_probe:foundations:span_end__example__long_task { + @long_task_ms = hist((uint64)arg0 / 1000000); +} + +usdt:target/debug/examples/span_with_probe:foundations:span_end__example__attr_task { + @attr_task_ms = hist((uint64)arg0 / 1000000); +}