perf(context): split context into flags and inner context

Author: bantonsson

opentelemetry/src/context.rs

@@ -93,11 +93,18 @@ thread_local! {
 /// assert_eq!(current.get::<ValueB>(), None);
 /// ```
 #[derive(Clone, Default)]
+#[cfg_attr(target_pointer_width = "64", repr(align(16)))]
+#[cfg_attr(target_pointer_width = "32", repr(align(8)))]
 pub struct Context {
+    pub(crate) inner: Option<Arc<InnerContext>>,
+    flags: ContextFlags,
+}
+
+#[derive(Default)]
+pub(crate) struct InnerContext {
     #[cfg(feature = "trace")]
     pub(crate) span: Option<Arc<SynchronizedSpan>>,
     entries: Option<Arc<EntryMap>>,
-    suppress_telemetry: bool,
 }
 
 type EntryMap = HashMap<TypeId, Arc<dyn Any + Sync + Send>, BuildHasherDefault<IdHasher>>;
@@ -198,7 +205,9 @@ impl Context {
     /// assert_eq!(cx.get::<MyUser>(), None);
     /// ```
     pub fn get<T: 'static>(&self) -> Option<&T> {
-        self.entries
+        self.inner
+            .as_ref()?
+            .entries
             .as_ref()?
             .get(&TypeId::of::<T>())?
             .downcast_ref()
@@ -232,20 +241,29 @@ impl Context {
     /// assert_eq!(cx_with_a_and_b.get::<ValueB>(), Some(&ValueB(42)));
     /// ```
     pub fn with_value<T: 'static + Send + Sync>(&self, value: T) -> Self {
-        let entries = if let Some(current_entries) = &self.entries {
-            let mut inner_entries = (**current_entries).clone();
-            inner_entries.insert(TypeId::of::<T>(), Arc::new(value));
-            Some(Arc::new(inner_entries))
-        } else {
+        fn new_entries<T: 'static + Send + Sync>(value: T) -> Option<Arc<EntryMap>> {
             let mut entries = EntryMap::default();
             entries.insert(TypeId::of::<T>(), Arc::new(value));
             Some(Arc::new(entries))
+        }
+        let (entries, span) = if let Some(inner) = &self.inner {
+            if let Some(current_entries) = &inner.entries {
+                let mut inner_entries = (**current_entries).clone();
+                inner_entries.insert(TypeId::of::<T>(), Arc::new(value));
+                (Some(Arc::new(inner_entries)), &inner.span)
+            } else {
+                (new_entries(value), &inner.span)
+            }
+        } else {
+            (new_entries(value), &None)
         };
         Context {
-            entries,
-            #[cfg(feature = "trace")]
-            span: self.span.clone(),
-            suppress_telemetry: self.suppress_telemetry,
+            inner: Some(Arc::new(InnerContext {
+                entries,
+                #[cfg(feature = "trace")]
+                span: span.clone(),
+            })),
+            flags: self.flags,
         }
     }
 
@@ -335,16 +353,14 @@ impl Context {
     /// Returns whether telemetry is suppressed in this context.
     #[inline]
     pub fn is_telemetry_suppressed(&self) -> bool {
-        self.suppress_telemetry
+        self.flags.is_telemetry_suppressed()
     }
 
     /// Returns a new context with telemetry suppression enabled.
     pub fn with_telemetry_suppressed(&self) -> Self {
         Context {
-            entries: self.entries.clone(),
-            #[cfg(feature = "trace")]
-            span: self.span.clone(),
-            suppress_telemetry: true,
+            inner: self.inner.clone(),
+            flags: self.flags.with_telemetry_suppressed(),
         }
     }
 
@@ -410,19 +426,45 @@ impl Context {
 
     #[cfg(feature = "trace")]
     pub(crate) fn current_with_synchronized_span(value: SynchronizedSpan) -> Self {
-        Self::map_current(|cx| Context {
-            span: Some(Arc::new(value)),
-            entries: cx.entries.clone(),
-            suppress_telemetry: cx.suppress_telemetry,
+        Self::map_current(|cx| {
+            if let Some(inner) = &cx.inner {
+                Context {
+                    inner: Some(Arc::new(InnerContext {
+                        span: Some(Arc::new(value)),
+                        entries: inner.entries.clone(),
+                    })),
+                    flags: cx.flags,
+                }
+            } else {
+                Context {
+                    inner: Some(Arc::new(InnerContext {
+                        span: Some(Arc::new(value)),
+                        entries: None,
+                    })),
+                    flags: ContextFlags::new(),
+                }
+            }
         })
     }
 
     #[cfg(feature = "trace")]
     pub(crate) fn with_synchronized_span(&self, value: SynchronizedSpan) -> Self {
-        Context {
-            span: Some(Arc::new(value)),
-            entries: self.entries.clone(),
-            suppress_telemetry: self.suppress_telemetry,
+        if let Some(inner) = &self.inner {
+            Context {
+                inner: Some(Arc::new(InnerContext {
+                    span: Some(Arc::new(value)),
+                    entries: inner.entries.clone(),
+                })),
+                flags: self.flags,
+            }
+        } else {
+            Context {
+                inner: Some(Arc::new(InnerContext {
+                    span: Some(Arc::new(value)),
+                    entries: None,
+                })),
+                flags: ContextFlags::new(),
+            }
         }
     }
 }
@@ -432,25 +474,67 @@ impl fmt::Debug for Context {
         let mut dbg = f.debug_struct("Context");
 
         #[cfg(feature = "trace")]
-        let mut entries = self.entries.as_ref().map_or(0, |e| e.len());
+        let mut entries = self
+            .inner
+            .as_ref()
+            .map_or(0, |i| i.entries.as_ref().map_or(0, |e| e.len()));
         #[cfg(feature = "trace")]
         {
-            if let Some(span) = &self.span {
+            if let Some(Some(span)) = self.inner.as_ref().map(|i| i.span.as_ref()) {
                 dbg.field("span", &span.span_context());
                 entries += 1;
             } else {
                 dbg.field("span", &"None");
             }
         }
         #[cfg(not(feature = "trace"))]
-        let entries = self.entries.as_ref().map_or(0, |e| e.len());
+        let entries = self
+            .inner
+            .as_ref()
+            .map_or(0, |i| i.entries.as_ref().map_or(0, |e| e.len()));
 
         dbg.field("entries count", &entries)
-            .field("suppress_telemetry", &self.suppress_telemetry)
+            .field("flags", &self.flags)
             .finish()
     }
 }
 
+/// Bit flags for context state.
+#[derive(Clone, Copy, Default)]
+struct ContextFlags(u16);
+
+impl ContextFlags {
+    const SUPPRESS_TELEMETRY: u16 = 1 << 0;
+
+    /// Creates a new ContextFlags with all flags cleared.
+    #[inline(always)]
+    const fn new() -> Self {
+        ContextFlags(0)
+    }
+
+    /// Returns true if telemetry suppression is enabled.
+    #[inline(always)]
+    const fn is_telemetry_suppressed(self) -> bool {
+        (self.0 & Self::SUPPRESS_TELEMETRY) != 0
+    }
+
+    /// Returns a new ContextFlags with telemetry suppression enabled.
+    #[inline(always)]
+    const fn with_telemetry_suppressed(self) -> Self {
+        ContextFlags(self.0 | Self::SUPPRESS_TELEMETRY)
+    }
+}
+
+impl fmt::Debug for ContextFlags {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("ContextFlags(")?;
+        if self.is_telemetry_suppressed() {
+            f.write_str("TELEMETRY_SUPPRESSED")?;
+        }
+        f.write_str(")")
+    }
+}
+
 /// A guard that resets the current context to the prior context when dropped.
 #[derive(Debug)]
 pub struct ContextGuard {

opentelemetry/src/trace/context.rs

@@ -308,15 +308,15 @@ impl TraceContextExt for Context {
     }
 
     fn span(&self) -> SpanRef<'_> {
-        if let Some(span) = self.span.as_ref() {
+        if let Some(Some(span)) = self.inner.as_ref().map(|i| i.span.as_ref()) {
             SpanRef(span)
         } else {
             SpanRef(&NOOP_SPAN)
         }
     }
 
     fn has_active_span(&self) -> bool {
-        self.span.is_some()
+        self.inner.as_ref().map_or(false, |i| i.span.is_some())
     }
 
     fn with_remote_span_context(&self, span_context: crate::trace::SpanContext) -> Self {

opentelemetry/src/trace/span_context.rs

@@ -444,9 +444,9 @@ mod tests {
         let cx = Context::current();
         assert_eq!(
             format!("{cx:?}"),
-            "Context { span: \"None\", entries count: 0, suppress_telemetry: false }"
+            "Context { span: \"None\", entries count: 0, flags: ContextFlags() }"
         );
-        let cx = Context::current().with_remote_span_context(SpanContext::NONE);
+        let cx = Context::current().with_remote_span_context(SpanContext::NONE).with_telemetry_suppressed();
         assert_eq!(
             format!("{cx:?}"),
             "Context { \
@@ -457,7 +457,7 @@ mod tests {
                        is_remote: false, \
                        trace_state: TraceState(None) \
                      }, \
-               entries count: 1, suppress_telemetry: false \
+               entries count: 1, flags: ContextFlags(TELEMETRY_SUPPRESSED) \
              }"
         );
     }