Context holds SynchronizedSpan directly, not via HashMap (#1268)

- lookup on TypeId and downcasting isn't necesary
- Context::with_value and current_with_value are more efficient
  as they no longer clone and overwrite the entry in the map
  which represens the current span

Author: shaun-cox

opentelemetry/src/context.rs

@@ -1,3 +1,4 @@
+use crate::trace::context::SynchronizedSpan;
 use std::any::{Any, TypeId};
 use std::cell::RefCell;
 use std::collections::HashMap;
@@ -74,6 +75,7 @@ thread_local! {
 /// ```
 #[derive(Clone, Default)]
 pub struct Context {
+    pub(super) span: Option<Arc<SynchronizedSpan>>,
     entries: HashMap<TypeId, Arc<dyn Any + Sync + Send>, BuildHasherDefault<IdHasher>>,
 }
 
@@ -303,6 +305,20 @@ impl Context {
             _marker: PhantomData,
         }
     }
+
+    pub(super) fn current_with_synchronized_span(value: SynchronizedSpan) -> Self {
+        Context {
+            span: Some(Arc::new(value)),
+            entries: Context::map_current(|cx| cx.entries.clone()),
+        }
+    }
+
+    pub(super) fn with_synchronized_span(&self, value: SynchronizedSpan) -> Self {
+        Context {
+            span: Some(Arc::new(value)),
+            entries: self.entries.clone(),
+        }
+    }
 }
 
 impl fmt::Debug for Context {

opentelemetry/src/trace/context.rs

@@ -26,13 +26,31 @@ static NOOP_SPAN: Lazy<SynchronizedSpan> = Lazy::new(|| SynchronizedSpan {
 pub struct SpanRef<'a>(&'a SynchronizedSpan);
 
 #[derive(Debug)]
-struct SynchronizedSpan {
+pub(crate) struct SynchronizedSpan {
     /// Immutable span context
     span_context: SpanContext,
     /// Mutable span inner that requires synchronization
     inner: Option<Mutex<global::BoxedSpan>>,
 }
 
+impl From<SpanContext> for SynchronizedSpan {
+    fn from(value: SpanContext) -> Self {
+        Self {
+            span_context: value,
+            inner: None,
+        }
+    }
+}
+
+impl<T: Span + Send + Sync + 'static> From<T> for SynchronizedSpan {
+    fn from(value: T) -> Self {
+        Self {
+            span_context: value.span_context().clone(),
+            inner: Some(Mutex::new(global::BoxedSpan::new(value))),
+        }
+    }
+}
+
 impl SpanRef<'_> {
     fn with_inner_mut<F: FnOnce(&mut global::BoxedSpan)>(&self, f: F) {
         if let Some(ref inner) = self.0.inner {
@@ -253,36 +271,27 @@ pub trait TraceContextExt {
 
 impl TraceContextExt for Context {
     fn current_with_span<T: crate::trace::Span + Send + Sync + 'static>(span: T) -> Self {
-        Context::current_with_value(SynchronizedSpan {
-            span_context: span.span_context().clone(),
-            inner: Some(Mutex::new(global::BoxedSpan::new(span))),
-        })
+        Context::current_with_synchronized_span(span.into())
     }
 
     fn with_span<T: crate::trace::Span + Send + Sync + 'static>(&self, span: T) -> Self {
-        self.with_value(SynchronizedSpan {
-            span_context: span.span_context().clone(),
-            inner: Some(Mutex::new(global::BoxedSpan::new(span))),
-        })
+        self.with_synchronized_span(span.into())
     }
 
     fn span(&self) -> SpanRef<'_> {
-        if let Some(span) = self.get::<SynchronizedSpan>() {
+        if let Some(span) = self.span.as_ref() {
             SpanRef(span)
         } else {
             SpanRef(&NOOP_SPAN)
         }
     }
 
     fn has_active_span(&self) -> bool {
-        self.get::<SynchronizedSpan>().is_some()
+        self.span.is_some()
     }
 
     fn with_remote_span_context(&self, span_context: crate::trace::SpanContext) -> Self {
-        self.with_value(SynchronizedSpan {
-            span_context,
-            inner: None,
-        })
+        self.with_synchronized_span(span_context.into())
     }
 }
 

opentelemetry/src/trace/mod.rs

@@ -166,7 +166,7 @@ use std::borrow::Cow;
 use std::time;
 use thiserror::Error;
 
-mod context;
+pub(crate) mod context;
 pub mod noop;
 mod span;
 mod span_context;