perf: Optimize cloning of Context since it is immutable

Author: bantonsson

opentelemetry/src/context.rs

@@ -92,10 +92,15 @@ thread_local! {
 /// ```
 #[derive(Clone, Default)]
 pub struct Context {
+    pub(super) inner: Option<Arc<InnerContext>>,
+}
+
+#[derive(Clone, Default)]
+pub(super) struct InnerContext {
     #[cfg(feature = "trace")]
     pub(crate) span: Option<Arc<SynchronizedSpan>>,
     entries: Option<Arc<EntryMap>>,
-    suppress_telemetry: bool,
+    context_flags: ContextFlags,
 }
 
 type EntryMap = HashMap<TypeId, Arc<dyn Any + Sync + Send>, BuildHasherDefault<IdHasher>>;
@@ -196,7 +201,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()
@@ -230,20 +237,36 @@ 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))
+        let (span, mut entries, context_flags) = if let Some(inner) = &self.inner {
+            #[cfg(feature = "trace")]
+            let span = inner.span.clone();
+            #[cfg(not(feature = "trace"))]
+            let span:Option<()> = None;
+
+            if let Some(entries) = &inner.entries {
+                (
+                    span,
+                    (**entries).clone(),
+                    inner.context_flags.clone(),
+                )
+            } else {
+                (
+                    span,
+                    EntryMap::default(),
+                    inner.context_flags.clone(),
+                )
+            }
         } else {
-            let mut entries = EntryMap::default();
-            entries.insert(TypeId::of::<T>(), Arc::new(value));
-            Some(Arc::new(entries))
+            (None, EntryMap::default(), ContextFlags::default())
         };
+        entries.insert(TypeId::of::<T>(), Arc::new(value));
         Context {
-            entries,
-            #[cfg(feature = "trace")]
-            span: self.span.clone(),
-            suppress_telemetry: self.suppress_telemetry,
+            inner: Some(Arc::new(InnerContext {
+                #[cfg(feature = "trace")]
+                span,
+                entries: Some(Arc::new(entries)),
+                context_flags,
+            })),
         }
     }
 
@@ -333,16 +356,49 @@ impl Context {
     /// Returns whether telemetry is suppressed in this context.
     #[inline]
     pub fn is_telemetry_suppressed(&self) -> bool {
-        self.suppress_telemetry
+        if let Some(inner) = &self.inner {
+            inner.context_flags.suppress_telemetry()
+        } else {
+            false
+        }
     }
 
     /// Returns a new context with telemetry suppression enabled.
     pub fn with_telemetry_suppressed(&self) -> Self {
-        Context {
-            entries: self.entries.clone(),
+        if self.is_telemetry_suppressed() {
+            return self.clone();
+        }
+        let (span, entries, context_flags) = if let Some(inner) = &self.inner {
             #[cfg(feature = "trace")]
-            span: self.span.clone(),
-            suppress_telemetry: true,
+            {
+                (
+                    inner.span.clone(),
+                    inner.entries.clone(),
+                    inner.context_flags.with_telemetry_suppressed(),
+                )
+            }
+            #[cfg(not(feature = "trace"))]
+            {
+                (
+                    Option::<()>::None,
+                    inner.entries.clone(),
+                    inner.context_flags.with_telemetry_suppressed(),
+                )
+            }
+        } else {
+            (
+                None,
+                None,
+                ContextFlags::default().with_telemetry_suppressed(),
+            )
+        };
+        Context {
+            inner: Some(Arc::new(InnerContext {
+                #[cfg(feature = "trace")]
+                span,
+                entries,
+                context_flags,
+            })),
         }
     }
 
@@ -408,30 +464,42 @@ 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| cx.with_synchronized_span(value))
     }
 
     #[cfg(feature = "trace")]
     pub(crate) fn with_synchronized_span(&self, value: SynchronizedSpan) -> Self {
+        let (entries, context_flags) = if let Some(inner) = &self.inner {
+            (inner.entries.clone(), inner.context_flags.clone())
+        } else {
+            (None, ContextFlags::default())
+        };
         Context {
-            span: Some(Arc::new(value)),
-            entries: self.entries.clone(),
-            suppress_telemetry: self.suppress_telemetry,
+            inner: Some(Arc::new(InnerContext {
+                span: Some(Arc::new(value)),
+                entries,
+                context_flags,
+            })),
         }
     }
 }
 
 impl fmt::Debug for Context {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut dbg = f.debug_struct("Context");
-        let mut entries = self.entries.as_ref().map_or(0, |e| e.len());
+        let (mut entries, suppress_telemetry) = self.inner.as_ref().map_or((0, false), |inner| {
+            (
+                inner.entries.as_ref().map_or(0, |e| e.len()),
+                inner.context_flags.suppress_telemetry(),
+            )
+        });
         #[cfg(feature = "trace")]
         {
-            if let Some(span) = &self.span {
+            if let Some(span) = self
+                .inner
+                .as_ref()
+                .map_or(None, |inner| inner.span.as_ref())
+            {
                 dbg.field("span", &span.span_context());
                 entries += 1;
             } else {
@@ -440,7 +508,7 @@ impl fmt::Debug for Context {
         }
 
         dbg.field("entries count", &entries)
-            .field("suppress_telemetry", &self.suppress_telemetry)
+            .field("suppress_telemetry", &suppress_telemetry)
             .finish()
     }
 }
@@ -599,6 +667,23 @@ impl Default for ContextStack {
     }
 }
 
+#[derive(Clone, Default)]
+struct ContextFlags(u32);
+
+impl ContextFlags {
+    const SUPPRESS_TELEMETRY: u32 = 0b0000_0001;
+
+    #[inline(always)]
+    fn suppress_telemetry(&self) -> bool {
+        self.0 & ContextFlags::SUPPRESS_TELEMETRY != 0
+    }
+
+    #[inline(always)]
+    fn with_telemetry_suppressed(&self) -> Self {
+        ContextFlags(self.0 | ContextFlags::SUPPRESS_TELEMETRY)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

opentelemetry/src/trace/context.rs

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