Store parent context in span builder (#399)

jtescher · web-flow · commit 601d2973e660 · 2020-12-27T16:37:08.000-08:00
The span builder currently holds a parent `SpanContext`, and `Context`
information is passed separately when starting a span directly or via a
builder. This poses a few problems, first you can assign a context which
contains an active span _as well as_ a builder with a parent span
context, which can lead to confusion, and secondly we have to construct
a wrapper context when sampling even in cases when the passed in context
already contains an active span.

This patch resolves this by changing the span builder's parent context
to instead store a `Context` directly. This allows the builder methods
to have a sigle context to look at when searchig for parent span
contexts, and allows the wrapper context to onlly be created for
sampling when it is facilitating a remote parent.
diff --git a/examples/aws-xray/src/server.rs b/examples/aws-xray/src/server.rs
@@ -20,7 +20,7 @@ async fn handle(req: Request<Body>) -> Result<Response<Body>, Infallible> {
         .to_str()
         .unwrap();
 
-    let span = global::tracer("example/server").start_from_context("hello", &parent_context);
+    let span = global::tracer("example/server").start_with_context("hello", parent_context);
     span.add_event(format!("Handling - {}", x_amzn_trace_id), Vec::new());
 
     Ok(Response::new(
diff --git a/examples/grpc/src/server.rs b/examples/grpc/src/server.rs
@@ -25,7 +25,7 @@ impl Greeter for MyGreeter {
         request: Request<HelloRequest>, // Accept request of type HelloRequest
     ) -> Result<Response<HelloReply>, Status> {
         let parent_cx = global::get_text_map_propagator(|prop| prop.extract(request.metadata()));
-        let span = global::tracer("greeter").start_from_context("Processing reply", &parent_cx);
+        let span = global::tracer("greeter").start_with_context("Processing reply", parent_cx);
         span.set_attribute(KeyValue::new("request", format!("{:?}", request)));
 
         // Return an instance of type HelloReply
diff --git a/examples/http/src/server.rs b/examples/http/src/server.rs
@@ -13,7 +13,7 @@ use std::{convert::Infallible, net::SocketAddr};
 
 async fn handle(req: Request<Body>) -> Result<Response<Body>, Infallible> {
     let parent_cx = global::get_text_map_propagator(|propagator| propagator.extract(req.headers()));
-    let span = global::tracer("example/server").start_from_context("hello", &parent_cx);
+    let span = global::tracer("example/server").start_with_context("hello", parent_cx);
     span.add_event("handling this...".to_string(), Vec::new());
 
     Ok(Response::new("Hello, World!".into()))
diff --git a/opentelemetry/src/api/trace/noop.rs b/opentelemetry/src/api/trace/noop.rs
@@ -129,12 +129,13 @@ impl trace::Tracer for NoopTracer {
         trace::NoopSpan::new()
     }
 
-    /// Starts a new `NoopSpan` in a given context.
+    /// Starts a new `NoopSpan` with a given context.
     ///
     /// If the context contains a valid span context, it is propagated.
-    fn start_from_context(&self, name: &str, cx: &Context) -> Self::Span {
-        let builder = self.span_builder(name);
-        self.build_with_context(builder, cx)
+    fn start_with_context(&self, name: &str, cx: Context) -> Self::Span {
+        let mut builder = self.span_builder(name);
+        builder.parent_context = Some(cx);
+        self.build(builder)
     }
 
     /// Starts a `SpanBuilder`.
@@ -145,18 +146,14 @@ impl trace::Tracer for NoopTracer {
     /// Builds a `NoopSpan` from a `SpanBuilder`.
     ///
     /// If the span builder or context contains a valid span context, it is propagated.
-    fn build_with_context(&self, mut builder: trace::SpanBuilder, cx: &Context) -> Self::Span {
-        let parent_span_context = builder
-            .parent_context
-            .take()
-            .or_else(|| {
-                if cx.has_active_span() {
-                    Some(cx.span().span_context().clone())
-                } else {
-                    None
-                }
-            })
-            .or_else(|| cx.remote_span_context().cloned());
+    fn build(&self, mut builder: trace::SpanBuilder) -> Self::Span {
+        let parent_span_context = builder.parent_context.take().and_then(|parent_cx| {
+            if parent_cx.has_active_span() {
+                Some(parent_cx.span().span_context().clone())
+            } else {
+                parent_cx.remote_span_context().cloned()
+            }
+        });
         if let Some(span_context) = parent_span_context {
             trace::NoopSpan { span_context }
         } else {
@@ -190,6 +187,7 @@ impl SpanExporter for NoopSpanExporter {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::testing::trace::TestSpan;
     use crate::trace::{self, Span, Tracer};
 
     fn valid_span_context() -> trace::SpanContext {
@@ -205,15 +203,17 @@ mod tests {
     #[test]
     fn noop_tracer_defaults_to_invalid_span() {
         let tracer = NoopTracer::new();
-        let span = tracer.start_from_context("foo", &Context::new());
+        let span = tracer.start_with_context("foo", Context::new());
         assert!(!span.span_context().is_valid());
     }
 
     #[test]
     fn noop_tracer_propagates_valid_span_context_from_builder() {
         let tracer = NoopTracer::new();
-        let builder = tracer.span_builder("foo").with_parent(valid_span_context());
-        let span = tracer.build_with_context(builder, &Context::new());
+        let builder = tracer
+            .span_builder("foo")
+            .with_parent_context(Context::current_with_span(TestSpan(valid_span_context())));
+        let span = tracer.build(builder);
         assert!(span.span_context().is_valid());
     }
 
@@ -223,15 +223,15 @@ mod tests {
         let cx = Context::new().with_span(NoopSpan {
             span_context: valid_span_context(),
         });
-        let span = tracer.start_from_context("foo", &cx);
+        let span = tracer.start_with_context("foo", cx);
         assert!(span.span_context().is_valid());
     }
 
     #[test]
     fn noop_tracer_propagates_valid_span_context_from_remote_span_context() {
         let tracer = NoopTracer::new();
         let cx = Context::new().with_remote_span_context(valid_span_context());
-        let span = tracer.start_from_context("foo", &cx);
+        let span = tracer.start_with_context("foo", cx);
         assert!(span.span_context().is_valid());
     }
 }
diff --git a/opentelemetry/src/api/trace/tracer.rs b/opentelemetry/src/api/trace/tracer.rs
@@ -1,8 +1,6 @@
 use crate::sdk;
 use crate::{
-    trace::{
-        Event, Link, Span, SpanContext, SpanId, SpanKind, StatusCode, TraceContextExt, TraceId,
-    },
+    trace::{Event, Link, Span, SpanId, SpanKind, StatusCode, TraceContextExt, TraceId},
     Context, KeyValue,
 };
 use std::fmt;
@@ -35,19 +33,20 @@ use std::time::SystemTime;
 /// Spans can be created and nested manually:
 ///
 /// ```
-/// use opentelemetry::{global, trace::{Span, Tracer}};
+/// use opentelemetry::{global, trace::{Span, Tracer, TraceContextExt}, Context};
 ///
 /// let tracer = global::tracer("my-component");
 ///
 /// let parent = tracer.start("foo");
+/// let parent_cx = Context::current_with_span(parent);
 /// let child = tracer.span_builder("bar")
-///     .with_parent(parent.span_context().clone())
+///     .with_parent_context(parent_cx.clone())
 ///     .start(&tracer);
 ///
 /// // ...
 ///
 /// child.end();
-/// parent.end();
+/// drop(parent_cx) // end parent
 /// ```
 ///
 /// Spans can also use the current thread's [`Context`] to track which span is active:
@@ -189,10 +188,10 @@ pub trait Tracer: fmt::Debug + 'static {
     /// `is_remote` to true on a parent `SpanContext` so `Span` creation knows if the
     /// parent is remote.
     fn start(&self, name: &str) -> Self::Span {
-        self.start_from_context(name, &Context::current())
+        self.start_with_context(name, Context::current())
     }
 
-    /// Starts a new `Span` in a given context
+    /// Starts a new `Span` with a given context
     ///
     /// By default the currently active `Span` is set as the new `Span`'s
     /// parent. The `Tracer` MAY provide other default options for newly
@@ -215,20 +214,15 @@ pub trait Tracer: fmt::Debug + 'static {
     /// created in another process. Each propagators' deserialization must set
     /// `is_remote` to true on a parent `SpanContext` so `Span` creation knows if the
     /// parent is remote.
-    fn start_from_context(&self, name: &str, context: &Context) -> Self::Span;
+    fn start_with_context(&self, name: &str, context: Context) -> Self::Span;
 
     /// Creates a span builder
     ///
     /// An ergonomic way for attributes to be configured before the `Span` is started.
     fn span_builder(&self, name: &str) -> SpanBuilder;
 
     /// Create a span from a `SpanBuilder`
-    fn build(&self, builder: SpanBuilder) -> Self::Span {
-        self.build_with_context(builder, &Context::current())
-    }
-
-    /// Create a span from a `SpanBuilder`
-    fn build_with_context(&self, builder: SpanBuilder, cx: &Context) -> Self::Span;
+    fn build(&self, builder: SpanBuilder) -> Self::Span;
 
     /// Start a new span and execute the given closure with reference to the span's
     /// context.
@@ -335,8 +329,8 @@ pub trait Tracer: fmt::Debug + 'static {
 /// ```
 #[derive(Clone, Debug, Default)]
 pub struct SpanBuilder {
-    /// Parent `SpanContext`
-    pub parent_context: Option<SpanContext>,
+    /// Parent `Context`
+    pub parent_context: Option<Context>,
     /// Trace id, useful for integrations with external tracing systems.
     pub trace_id: Option<TraceId>,
     /// Span id, useful for integrations with external tracing systems.
@@ -385,7 +379,7 @@ impl SpanBuilder {
     }
 
     /// Assign parent context
-    pub fn with_parent(self, parent_context: SpanContext) -> Self {
+    pub fn with_parent_context(self, parent_context: Context) -> Self {
         SpanBuilder {
             parent_context: Some(parent_context),
             ..self
diff --git a/opentelemetry/src/global/trace.rs b/opentelemetry/src/global/trace.rs
@@ -92,7 +92,7 @@ impl trace::Tracer for BoxedTracer {
     /// trace. A span is said to be a _root span_ if it does not have a parent. Each
     /// trace includes a single root span, which is the shared ancestor of all other
     /// spans in the trace.
-    fn start_from_context(&self, name: &str, cx: &Context) -> Self::Span {
+    fn start_with_context(&self, name: &str, cx: Context) -> Self::Span {
         BoxedSpan(self.0.start_with_context_boxed(name, cx))
     }
 
@@ -104,8 +104,8 @@ impl trace::Tracer for BoxedTracer {
     }
 
     /// Create a span from a `SpanBuilder`
-    fn build_with_context(&self, builder: trace::SpanBuilder, cx: &Context) -> Self::Span {
-        BoxedSpan(self.0.build_with_context_boxed(builder, cx))
+    fn build(&self, builder: trace::SpanBuilder) -> Self::Span {
+        BoxedSpan(self.0.build_boxed(builder))
     }
 }
 
@@ -120,11 +120,11 @@ pub trait GenericTracer: fmt::Debug + 'static {
 
     /// Returns a trait object so the underlying implementation can be swapped
     /// out at runtime.
-    fn start_with_context_boxed(&self, name: &str, cx: &Context) -> Box<DynSpan>;
+    fn start_with_context_boxed(&self, name: &str, cx: Context) -> Box<DynSpan>;
 
     /// Returns a trait object so the underlying implementation can be swapped
     /// out at runtime.
-    fn build_with_context_boxed(&self, builder: trace::SpanBuilder, cx: &Context) -> Box<DynSpan>;
+    fn build_boxed(&self, builder: trace::SpanBuilder) -> Box<DynSpan>;
 }
 
 impl<S, T> GenericTracer for T
@@ -139,14 +139,14 @@ where
 
     /// Returns a trait object so the underlying implementation can be swapped
     /// out at runtime.
-    fn start_with_context_boxed(&self, name: &str, cx: &Context) -> Box<DynSpan> {
-        Box::new(self.start_from_context(name, cx))
+    fn start_with_context_boxed(&self, name: &str, cx: Context) -> Box<DynSpan> {
+        Box::new(self.start_with_context(name, cx))
     }
 
     /// Returns a trait object so the underlying implementation can be swapped
     /// out at runtime.
-    fn build_with_context_boxed(&self, builder: trace::SpanBuilder, cx: &Context) -> Box<DynSpan> {
-        Box::new(self.build_with_context(builder, cx))
+    fn build_boxed(&self, builder: trace::SpanBuilder) -> Box<DynSpan> {
+        Box::new(self.build(builder))
     }
 }
 
diff --git a/opentelemetry/src/sdk/trace/sampler.rs b/opentelemetry/src/sdk/trace/sampler.rs
@@ -113,19 +113,21 @@ impl ShouldSample for Sampler {
             // Never sample the trace
             Sampler::AlwaysOff => SamplingDecision::Drop,
             // The parent decision if sampled; otherwise the decision of delegate_sampler
-            Sampler::ParentBased(delegate_sampler) => parent_context.map_or(
-                delegate_sampler
-                    .should_sample(parent_context, trace_id, name, span_kind, attributes, links)
-                    .decision,
-                |ctx| {
-                    let parent_span_context = ctx.span().span_context();
-                    if parent_span_context.is_sampled() {
-                        SamplingDecision::RecordAndSample
-                    } else {
-                        SamplingDecision::Drop
-                    }
-                },
-            ),
+            Sampler::ParentBased(delegate_sampler) => {
+                parent_context.filter(|cx| cx.has_active_span()).map_or(
+                    delegate_sampler
+                        .should_sample(parent_context, trace_id, name, span_kind, attributes, links)
+                        .decision,
+                    |ctx| {
+                        let parent_span_context = ctx.span().span_context();
+                        if parent_span_context.is_sampled() {
+                            SamplingDecision::RecordAndSample
+                        } else {
+                            SamplingDecision::Drop
+                        }
+                    },
+                )
+            }
             // Probabilistically sample the trace.
             Sampler::TraceIdRatioBased(prob) => {
                 if *prob >= 1.0 {
@@ -270,4 +272,20 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn filter_parent_sampler_for_active_spans() {
+        let sampler = Sampler::ParentBased(Box::new(Sampler::AlwaysOn));
+        let cx = Context::current_with_value("some_value");
+        let result = sampler.should_sample(
+            Some(&cx),
+            TraceId::from_u128(1),
+            "should sample",
+            &SpanKind::Internal,
+            &[],
+            &[],
+        );
+
+        assert_eq!(result.decision, SamplingDecision::RecordAndSample);
+    }
 }
diff --git a/opentelemetry/src/sdk/trace/span_processor.rs b/opentelemetry/src/sdk/trace/span_processor.rs
@@ -695,10 +695,8 @@ mod tests {
         D: Fn(time::Duration) -> DS + 'static + Send + Sync,
         DS: Future<Output = ()> + Send + Sync + 'static,
     {
-        async fn export(&mut self, batch: Vec<SpanData>) -> ExportResult {
-            println!("Accepting {} spans", batch.len());
+        async fn export(&mut self, _batch: Vec<SpanData>) -> ExportResult {
             (self.delay_fn)(self.delay_for).await;
-            println!("Finish exporting, return result from exporter");
             Ok(())
         }
     }
diff --git a/opentelemetry/src/sdk/trace/tracer.rs b/opentelemetry/src/sdk/trace/tracer.rs

Original file line number	Diff line number	Diff line change
`@@ -92,7 +92,7 @@ impl trace::Tracer for BoxedTracer {`
`92`	`92`	`/// trace. A span is said to be a _root span_ if it does not have a parent. Each`
`93`	`93`	`/// trace includes a single root span, which is the shared ancestor of all other`
`94`	`94`	`/// spans in the trace.`
`95`		`- fn start_from_context(&self, name: &str, cx: &Context) -> Self::Span {`
	`95`	`+ fn start_with_context(&self, name: &str, cx: Context) -> Self::Span {`
`96`	`96`	`BoxedSpan(self.0.start_with_context_boxed(name, cx))`
`97`	`97`	`}`
`98`	`98`
`@@ -104,8 +104,8 @@ impl trace::Tracer for BoxedTracer {`
`104`	`104`	`}`
`105`	`105`
`106`	`106`	/// Create a span from a `SpanBuilder`
`107`		`- fn build_with_context(&self, builder: trace::SpanBuilder, cx: &Context) -> Self::Span {`
`108`		`- BoxedSpan(self.0.build_with_context_boxed(builder, cx))`
	`107`	`+ fn build(&self, builder: trace::SpanBuilder) -> Self::Span {`
	`108`	`+ BoxedSpan(self.0.build_boxed(builder))`
`109`	`109`	`}`
`110`	`110`	`}`
`111`	`111`
`@@ -120,11 +120,11 @@ pub trait GenericTracer: fmt::Debug + 'static {`
`120`	`120`
`121`	`121`	`/// Returns a trait object so the underlying implementation can be swapped`
`122`	`122`	`/// out at runtime.`
`123`		`- fn start_with_context_boxed(&self, name: &str, cx: &Context) -> Box<DynSpan>;`
	`123`	`+ fn start_with_context_boxed(&self, name: &str, cx: Context) -> Box<DynSpan>;`
`124`	`124`
`125`	`125`	`/// Returns a trait object so the underlying implementation can be swapped`
`126`	`126`	`/// out at runtime.`
`127`		`- fn build_with_context_boxed(&self, builder: trace::SpanBuilder, cx: &Context) -> Box<DynSpan>;`
	`127`	`+ fn build_boxed(&self, builder: trace::SpanBuilder) -> Box<DynSpan>;`
`128`	`128`	`}`
`129`	`129`
`130`	`130`	`impl<S, T> GenericTracer for T`
`@@ -139,14 +139,14 @@ where`
`139`	`139`
`140`	`140`	`/// Returns a trait object so the underlying implementation can be swapped`
`141`	`141`	`/// out at runtime.`
`142`		`- fn start_with_context_boxed(&self, name: &str, cx: &Context) -> Box<DynSpan> {`
`143`		`- Box::new(self.start_from_context(name, cx))`
	`142`	`+ fn start_with_context_boxed(&self, name: &str, cx: Context) -> Box<DynSpan> {`
	`143`	`+ Box::new(self.start_with_context(name, cx))`
`144`	`144`	`}`
`145`	`145`
`146`	`146`	`/// Returns a trait object so the underlying implementation can be swapped`
`147`	`147`	`/// out at runtime.`
`148`		`- fn build_with_context_boxed(&self, builder: trace::SpanBuilder, cx: &Context) -> Box<DynSpan> {`
`149`		`- Box::new(self.build_with_context(builder, cx))`
	`148`	`+ fn build_boxed(&self, builder: trace::SpanBuilder) -> Box<DynSpan> {`
	`149`	`+ Box::new(self.build(builder))`
`150`	`150`	`}`
`151`	`151`	`}`
`152`	`152`
Original file line number	Diff line number	Diff line change
`@@ -695,10 +695,8 @@ mod tests {`
`695`	`695`	`D: Fn(time::Duration) -> DS + 'static + Send + Sync,`
`696`	`696`	`DS: Future<Output = ()> + Send + Sync + 'static,`
`697`	`697`	`{`
`698`		`- async fn export(&mut self, batch: Vec<SpanData>) -> ExportResult {`
`699`		`- println!("Accepting {} spans", batch.len());`
	`698`	`+ async fn export(&mut self, _batch: Vec<SpanData>) -> ExportResult {`
`700`	`699`	`(self.delay_fn)(self.delay_for).await;`
`701`		`- println!("Finish exporting, return result from exporter");`
`702`	`700`	`Ok(())`
`703`	`701`	`}`
`704`	`702`	`}`