perf: Optimize cloning of Context since it is immutable

Author: bantonsson

opentelemetry/src/context.rs

@@ -94,6 +94,11 @@ 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>>,
@@ -198,7 +203,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 +239,30 @@ 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))
+        if let Some(inner) = &self.inner {
+            let (mut entries, suppress_telemetry) = if let Some(entries) = &inner.entries {
+                ((**entries).clone(), inner.suppress_telemetry)
+            } else {
+                (EntryMap::default(), inner.suppress_telemetry)
+            };
+            entries.insert(TypeId::of::<T>(), Arc::new(value));
+            Context {
+                inner: Some(Arc::new(InnerContext {
+                    #[cfg(feature = "trace")]
+                    span: inner.span.clone(),
+                    entries: Some(Arc::new(entries)),
+                    suppress_telemetry,
+                })),
+            }
         } else {
             let mut entries = EntryMap::default();
             entries.insert(TypeId::of::<T>(), Arc::new(value));
-            Some(Arc::new(entries))
-        };
-        Context {
-            entries,
-            #[cfg(feature = "trace")]
-            span: self.span.clone(),
-            suppress_telemetry: self.suppress_telemetry,
+            Context {
+                inner: Some(Arc::new(InnerContext {
+                    entries: Some(Arc::new(entries)),
+                    ..Default::default()
+                })),
+            }
         }
     }
 
@@ -335,16 +352,35 @@ 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.suppress_telemetry
+        } else {
+            false
+        }
     }
 
     /// 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,
+        if self.is_telemetry_suppressed() {
+            return self.clone();
+        }
+
+        if let Some(inner) = &self.inner {
+            Context {
+                inner: Some(Arc::new(InnerContext {
+                    #[cfg(feature = "trace")]
+                    span: inner.span.clone(),
+                    entries: inner.entries.clone(),
+                    suppress_telemetry: true,
+                })),
+            }
+        } else {
+            Context {
+                inner: Some(Arc::new(InnerContext {
+                    suppress_telemetry: true,
+                    ..Default::default()
+                })),
+            }
         }
     }
 
@@ -410,43 +446,53 @@ 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, suppress_telemetry) = if let Some(inner) = &self.inner {
+            (inner.entries.clone(), inner.suppress_telemetry)
+        } else {
+            (None, false)
+        };
         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,
+                suppress_telemetry,
+            })),
         }
     }
 }
 
 impl fmt::Debug for Context {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut dbg = f.debug_struct("Context");
+        let mut entries = 0;
+        let mut suppress_telemetry = false;
 
-        #[cfg(feature = "trace")]
-        let mut entries = self.entries.as_ref().map_or(0, |e| e.len());
-        #[cfg(feature = "trace")]
-        {
-            if let Some(span) = &self.span {
-                dbg.field("span", &span.span_context());
-                entries += 1;
-            } else {
+        if let Some(inner) = &self.inner {
+            entries = inner.entries.as_ref().map_or(0, |e| e.len());
+            #[cfg(feature = "trace")]
+            {
+                if let Some(span) = inner.span.as_ref() {
+                    dbg.field("span", &span.span_context());
+                    entries += 1;
+                } else {
+                    dbg.field("span", &"None");
+                }
+            }
+            suppress_telemetry = inner.suppress_telemetry;
+        } else {
+            #[cfg(feature = "trace")]
+            {
                 dbg.field("span", &"None");
             }
         }
-        #[cfg(not(feature = "trace"))]
-        let entries = self.entries.as_ref().map_or(0, |e| e.len());
 
         dbg.field("entries count", &entries)
-            .field("suppress_telemetry", &self.suppress_telemetry)
+            .field("suppress_telemetry", &suppress_telemetry)
             .finish()
     }
 }

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() {
+            inner.span.is_some()
+        } else {
+            false
+        }
     }
 
     fn with_remote_span_context(&self, span_context: crate::trace::SpanContext) -> Self {