diff --git a/opentelemetry/src/context.rs b/opentelemetry/src/context.rs index 79c7d0c0c4..65c2f17364 100644 --- a/opentelemetry/src/context.rs +++ b/opentelemetry/src/context.rs @@ -79,9 +79,11 @@ thread_local! { pub struct Context { #[cfg(feature = "trace")] pub(super) span: Option>, - entries: HashMap, BuildHasherDefault>, + entries: Option>, } +type EntryMap = HashMap, BuildHasherDefault>; + impl Context { /// Creates an empty `Context`. /// @@ -111,7 +113,7 @@ impl Context { /// do_work() /// ``` pub fn current() -> Self { - Context::map_current(|cx| cx.clone()) + Self::map_current(|cx| cx.clone()) } /// Applies a function to the current context returning its value. @@ -153,12 +155,7 @@ impl Context { /// assert_eq!(all_current_and_b.get::(), Some(&ValueB(42))); /// ``` pub fn current_with_value(value: T) -> Self { - let mut new_context = Context::current(); - new_context - .entries - .insert(TypeId::of::(), Arc::new(value)); - - new_context + Self::map_current(|cx| cx.with_value(value)) } /// Returns a reference to the entry for the corresponding value type. @@ -184,8 +181,9 @@ impl Context { /// ``` pub fn get(&self) -> Option<&T> { self.entries - .get(&TypeId::of::()) - .and_then(|rc| rc.downcast_ref()) + .as_ref()? + .get(&TypeId::of::())? + .downcast_ref() } /// Returns a copy of the context with the new value included. @@ -216,12 +214,20 @@ impl Context { /// assert_eq!(cx_with_a_and_b.get::(), Some(&ValueB(42))); /// ``` pub fn with_value(&self, value: T) -> Self { - let mut new_context = self.clone(); - new_context - .entries - .insert(TypeId::of::(), Arc::new(value)); - - new_context + let entries = if let Some(current_entries) = &self.entries { + let mut inner_entries = (**current_entries).clone(); + inner_entries.insert(TypeId::of::(), Arc::new(value)); + Some(Arc::new(inner_entries)) + } else { + let mut entries = EntryMap::default(); + entries.insert(TypeId::of::(), Arc::new(value)); + Some(Arc::new(entries)) + }; + Context { + entries, + #[cfg(feature = "trace")] + span: self.span.clone(), + } } /// Replaces the current context on this thread with this context. @@ -327,7 +333,7 @@ impl Context { 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.len(); + let mut entries = self.entries.as_ref().map_or(0, |e| e.len()); #[cfg(feature = "trace")] { if let Some(span) = &self.span {