Cache SharedSpan::is_sampled for deferred user tracing spans

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<BoxSampler<SpanContextState>, SpanContextState>;
@@ -75,22 +76,83 @@ impl From<SharedSpanHandle> for Arc<RwLock<Span>> {
     }
 }
 
-#[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
     /// clone of the span drops.
     pub(crate) probe: Option<Arc<SpanProbe>>,
 }
 
+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<Cow<'static, str>>) -> SharedSpan {
     match current_span() {
-        Some(span) if span.is_sampled => span,
+        Some(span) if span.is_sampled() => span,
         _ => return shared_span(Span::inactive()),
     };
 

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.