From cf2b471e25e7892ad642dd597e7bab8194bebfc0 Mon Sep 17 00:00:00 2001
From: Julian Tescher <jatescher@gmail.com>
Date: Mon, 28 Dec 2020 15:25:42 -0800
Subject: [PATCH] Use default context if unset for builders (#401)

---
 opentelemetry/src/sdk/trace/tracer.rs | 79 +++++++++++++++++----------
 1 file changed, 49 insertions(+), 30 deletions(-)

diff --git a/opentelemetry/src/sdk/trace/tracer.rs b/opentelemetry/src/sdk/trace/tracer.rs
index 29d54cab20..5c229f1a84 100644
--- a/opentelemetry/src/sdk/trace/tracer.rs
+++ b/opentelemetry/src/sdk/trace/tracer.rs
@@ -67,7 +67,7 @@ impl Tracer {
     #[allow(clippy::too_many_arguments)]
     fn make_sampling_decision(
         &self,
-        parent_cx: Option<&Context>,
+        parent_cx: &Context,
         trace_id: TraceId,
         name: &str,
         span_kind: &SpanKind,
@@ -77,8 +77,14 @@ impl Tracer {
         let provider = self.provider()?;
         let sampler = &provider.config().default_sampler;
 
-        let sampling_result =
-            sampler.should_sample(parent_cx, trace_id, name, span_kind, attributes, links);
+        let sampling_result = sampler.should_sample(
+            Some(parent_cx),
+            trace_id,
+            name,
+            span_kind,
+            attributes,
+            links,
+        );
 
         self.process_sampling_result(sampling_result, parent_cx)
     }
@@ -86,7 +92,7 @@ impl Tracer {
     fn process_sampling_result(
         &self,
         sampling_result: SamplingResult,
-        parent_cx: Option<&Context>,
+        parent_cx: &Context,
     ) -> Option<(u8, Vec<KeyValue>, TraceState)> {
         match sampling_result {
             SamplingResult {
@@ -98,9 +104,7 @@ impl Tracer {
                 attributes,
                 trace_state,
             } => {
-                let trace_flags = parent_cx
-                    .map(|ctx| ctx.span().span_context().trace_flags())
-                    .unwrap_or(0);
+                let trace_flags = parent_cx.span().span_context().trace_flags();
                 Some((trace_flags & !TRACE_FLAG_SAMPLED, attributes, trace_state))
             }
             SamplingResult {
@@ -108,9 +112,7 @@ impl Tracer {
                 attributes,
                 trace_state,
             } => {
-                let trace_flags = parent_cx
-                    .map(|ctx| ctx.span().span_context().trace_flags())
-                    .unwrap_or(0);
+                let trace_flags = parent_cx.span().span_context().trace_flags();
                 Some((trace_flags | TRACE_FLAG_SAMPLED, attributes, trace_state))
             }
         }
@@ -174,24 +176,26 @@ impl crate::trace::Tracer for Tracer {
         let mut flags = 0;
         let mut span_trace_state = Default::default();
 
-        let parent_cx = builder.parent_context.take().map(|cx| {
-            // Sampling expects to be able to access the parent span via `span` so wrap remote span
-            // context in a wrapper span if necessary. Remote span contexts will be passed to
-            // subsequent context's, so wrapping is only necessary if there is no active span.
-            match cx.remote_span_context() {
-                Some(remote_sc) if !cx.has_active_span() => {
-                    cx.with_span(Span::new(remote_sc.clone(), None, self.clone()))
+        let parent_cx = builder
+            .parent_context
+            .take()
+            .map(|cx| {
+                // Sampling expects to be able to access the parent span via `span` so wrap remote span
+                // context in a wrapper span if necessary. Remote span contexts will be passed to
+                // subsequent context's, so wrapping is only necessary if there is no active span.
+                match cx.remote_span_context() {
+                    Some(remote_sc) if !cx.has_active_span() => {
+                        cx.with_span(Span::new(remote_sc.clone(), None, self.clone()))
+                    }
+                    _ => cx,
                 }
-                _ => cx,
-            }
-        });
-        let parent_span_context = parent_cx.as_ref().and_then(|parent_cx| {
-            if parent_cx.has_active_span() {
-                Some(parent_cx.span().span_context())
-            } else {
-                None
-            }
-        });
+            })
+            .unwrap_or_else(Context::current);
+        let parent_span_context = if parent_cx.has_active_span() {
+            Some(parent_cx.span().span_context())
+        } else {
+            None
+        };
         // Build context for sampling decision
         let (no_parent, trace_id, parent_span_id, remote_parent, parent_trace_flags) =
             parent_span_context
@@ -221,10 +225,10 @@ impl crate::trace::Tracer for Tracer {
         // * There is no parent or a remote parent, in which case make decision now
         // * There is a local parent, in which case defer to the parent's decision
         let sampling_decision = if let Some(sampling_result) = builder.sampling_result.take() {
-            self.process_sampling_result(sampling_result, parent_cx.as_ref())
+            self.process_sampling_result(sampling_result, &parent_cx)
         } else if no_parent || remote_parent {
             self.make_sampling_decision(
-                parent_cx.as_ref(),
+                &parent_cx,
                 trace_id,
                 &builder.name,
                 &span_kind,
@@ -288,7 +292,7 @@ impl crate::trace::Tracer for Tracer {
 
         // Call `on_start` for all processors
         for processor in provider.span_processors() {
-            processor.on_start(&span, parent_cx.as_ref().unwrap_or(&Context::new()))
+            processor.on_start(&span, &parent_cx)
         }
 
         span
@@ -377,4 +381,19 @@ mod tests {
 
         assert!(!span.span_context().is_sampled());
     }
+
+    #[test]
+    fn uses_current_context_for_builders_if_unset() {
+        let sampler = Sampler::ParentBased(Box::new(Sampler::AlwaysOn));
+        let config = Config::default().with_default_sampler(sampler);
+        let tracer_provider = sdk::trace::TracerProvider::builder()
+            .with_config(config)
+            .build();
+
+        let _attached = Context::current_with_span(TestSpan(SpanContext::empty_context())).attach();
+        let tracer = tracer_provider.get_tracer("test", None);
+        let span = tracer.span_builder("must_not_be_sampled").start(&tracer);
+
+        assert!(!span.span_context().is_sampled());
+    }
 }