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); +} 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()); + } +}