From 3fd770e7e749ab5b32f656db642e3b8d17917cc3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Leo=20Bl=C3=B6cher?= Date: Thu, 3 Sep 2026 15:20:09 +0100 Subject: [PATCH] Cache `SharedSpan::is_sampled` for deferred user tracing spans Relaxed atomic load and stores are just regular load and stores on x86 and arm64, so this adds zero cost for regular spans and allows us to avoid some more lock accesses for deferred user tracing spans. --- foundations/src/telemetry/tracing/internal.rs | 102 ++++++++++++------ foundations/src/telemetry/tracing/mod.rs | 21 ++-- 2 files changed, 80 insertions(+), 43 deletions(-) diff --git a/foundations/src/telemetry/tracing/internal.rs b/foundations/src/telemetry/tracing/internal.rs index bc2cbb66..e8b0f6d6 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::sync::atomic::{AtomicU8, Ordering}; use std::time::Instant; pub(crate) type Tracer = cf_rustracing::Tracer, SpanContextState>; @@ -75,15 +76,16 @@ impl From for Arc> { } } -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct SharedSpan { // NOTE: we intentionally use a lock without poisoning here to not // panic the threads if they just share telemetry with failed thread. pub(crate) inner: SharedSpanHandle, // NOTE: store sampling flag separately, so we don't need to acquire lock // every time we need to check the flag. Deferred user roots are the exception: - // their initially inactive span can be replaced in place. - pub(crate) is_sampled: bool, + // their initially inactive span can be replaced in place, which is represented + // by the `0xFF` sentinel (`DEFERRED_SAMPLING`). + is_sampled: AtomicU8, /// 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 @@ -91,6 +93,66 @@ pub(crate) struct SharedSpan { pub(crate) probe: Option>, } +impl SharedSpan { + /// Sentinel value to indicate a deferred span in `is_sampled`. + const DEFERRED_SAMPLING: u8 = 0xFF; + + /// Creates a [`SharedSpan`] equivalent to [`Span::inactive()`]. + #[cfg(feature = "user-tracing")] + pub(crate) const fn inactive() -> Self { + Self { + inner: SharedSpanHandle::Inactive, + is_sampled: AtomicU8::new(false as u8), + probe: None, + } + } + + /// Creates a [`SharedSpan`] whose root can be activated in place + /// after contexts have captured it. + #[cfg(feature = "user-tracing")] + pub(crate) fn deferred() -> Self { + Self { + inner: SharedSpanHandle::Deferred(Arc::new(RwLock::new(Span::inactive()))), + is_sampled: AtomicU8::new(Self::DEFERRED_SAMPLING), + probe: None, + } + } + + #[inline] + pub(crate) fn is_sampled(&self) -> bool { + match self.is_sampled.load(Ordering::Relaxed) { + 0 => return false, + 1 => return true, + // Deferred SharedSpan, look inside the lock + _v => { + debug_assert_eq!( + _v, + Self::DEFERRED_SAMPLING, + "unexpected value in SharedSpan::is_sampled", + ); + } + } + + let is_sampled = self.inner.with_read(|span| span.is_sampled()); + if is_sampled { + // We never de-initialize a span inside SharedSpan, so we can + // save the result once its true. + self.is_sampled.store(true as u8, Ordering::Relaxed); + } + is_sampled + } +} + +impl Clone for SharedSpan { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + is_sampled: AtomicU8::new(self.is_sampled.load(Ordering::Relaxed)), + probe: self.probe.clone(), + } + } +} + /// 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)] @@ -120,7 +182,7 @@ pub(crate) fn shared_span(span: Span) -> SharedSpan { SharedSpan { inner: SharedSpanHandle::new(span), - is_sampled, + is_sampled: AtomicU8::new(is_sampled as u8), probe: None, } } @@ -139,36 +201,16 @@ pub(crate) fn user_shared_span(span: Span) -> SharedSpan { SharedSpan { inner, - is_sampled, - probe: None, - } -} - -/// Creates a user span whose root can be activated in place after contexts have captured it. -#[cfg(feature = "user-tracing")] -pub(crate) fn deferred_user_span() -> SharedSpan { - SharedSpan { - inner: SharedSpanHandle::Deferred(Arc::new(RwLock::new(Span::inactive()))), - is_sampled: false, + is_sampled: AtomicU8::new(is_sampled as u8), probe: None, } } -#[cfg(feature = "user-tracing")] -impl SharedSpan { - pub(crate) fn user_is_sampled(&self) -> bool { - match &self.inner { - SharedSpanHandle::Deferred(span) => span.read().is_sampled(), - _ => self.is_sampled, - } - } -} - pub fn write_current_span(write_fn: impl FnOnce(&mut Span)) { // Check the cached flag before touching the lock. Writing to an unsampled span is a no-op // anyway, so this only avoids taking a write guard for nothing. let span = match current_span() { - Some(span) if span.is_sampled => span, + Some(span) if span.is_sampled() => span, _ => return, }; @@ -212,7 +254,7 @@ pub(crate) fn start_trace( } let mut current_span = match current_span() { - Some(current_span) if current_span.is_sampled => current_span, + Some(current_span) if current_span.is_sampled() => current_span, _ => return span_builder.start(), }; @@ -252,7 +294,7 @@ pub(crate) fn child_user_span( #[cfg(feature = "user-tracing")] pub fn write_current_user_span(write_fn: impl FnOnce(&mut Span)) { let span = match current_user_span() { - Some(span) if span.user_is_sampled() => span, + Some(span) if span.is_sampled() => span, _ => return, }; @@ -286,7 +328,7 @@ pub(crate) fn activate_deferred_user_trace( return; } - *span = start_user_trace(name, Arc::clone(&routing), inbound); + *span = start_user_trace(name, routing, inbound); } /// Starts a root user span on the user harness, optionally continuing the inbound W3C trace. @@ -361,7 +403,7 @@ fn link_new_trace_with_current( pub(crate) fn fork_trace(fork_name: impl Into>) -> SharedSpan { match current_span() { - Some(span) if span.is_sampled => span, + Some(span) if span.is_sampled() => span, _ => return shared_span(Span::inactive()), }; diff --git a/foundations/src/telemetry/tracing/mod.rs b/foundations/src/telemetry/tracing/mod.rs index 3456d5b6..a0023de7 100644 --- a/foundations/src/telemetry/tracing/mod.rs +++ b/foundations/src/telemetry/tracing/mod.rs @@ -30,8 +30,8 @@ use self::internal::{ }; #[cfg(feature = "user-tracing")] use self::internal::{ - SharedSpanHandle, activate_deferred_user_trace, child_user_span, current_user_span, - deferred_user_span, start_user_trace, user_shared_span, + activate_deferred_user_trace, child_user_span, current_user_span, start_user_trace, + user_shared_span, }; use super::TelemetryContext; use super::scope::Scope; @@ -335,7 +335,7 @@ pub struct UserSpan { #[cfg(feature = "user-tracing")] impl UserSpan { #[inline] - pub(crate) fn from_shared(span: SharedSpan) -> Self { + pub(crate) const fn from_shared(span: SharedSpan) -> Self { Self { span } } @@ -344,13 +344,7 @@ impl UserSpan { /// Useful as a placeholder before a trace has been started, and as the resting state for a /// span slot that has been finished. pub const fn inactive() -> Self { - Self { - span: SharedSpan { - inner: SharedSpanHandle::Inactive, - is_sampled: false, - probe: None, - }, - } + Self::from_shared(SharedSpan::inactive()) } /// Creates a root user span that can be activated after telemetry contexts have captured it. @@ -359,7 +353,7 @@ impl UserSpan { /// reported. Activation updates shared state in place, so scopes and [`TelemetryContext`] /// values created from this handle beforehand observe the active span without being replaced. pub fn deferred() -> Self { - Self::from_shared(deferred_user_span()) + Self::from_shared(SharedSpan::deferred()) } /// Activates this deferred root, continuing the inbound W3C trace from `inbound` when given. @@ -422,8 +416,9 @@ impl UserSpan { /// /// Eager spans read cached state. Deferred roots inspect their shared span slot so contexts /// captured before activation observe the updated state. + #[inline] pub fn is_sampled(&self) -> bool { - self.span.user_is_sampled() + self.span.is_sampled() } /// W3C `traceparent` for this span, for outbound propagation to the next hop. @@ -520,7 +515,7 @@ pub struct StartTraceOptions { /// This is useful to do a cheap check before commencing more expensive work, /// for example to set up span tags. pub fn span_is_sampled() -> bool { - matches!(current_span(), Some(span) if span.is_sampled) + matches!(current_span(), Some(span) if span.is_sampled()) } /// Returns a trace ID of the current span.