opentelemetry: Add PreSampledTracer interface (tokio-rs#962)

## Motivation

Currently the `OpenTelemetryLayer` is coupled to the opentelemetry SDK
in a way that makes using alternate API implementations difficult (they
would have to use the SDK's sampler and id generator) as well as causing
duplication when specifying sampler configuration (it must currently be
set on both the tracer provider configuration as well as the layer
configuration, which is difficult to remember to keep in sync).

## Solution

This change creates an interface for working with pre-sampled tracers,
and implements it for the current opentelemetry SDK. This simplifies the
layer as it no longer needs to track an id generator or sampler, and
allows other otel API implementations to implement the trait rather than
rely on the current SDK.

Co-authored-by: Eliza Weisman <eliza@buoyant.io>

Author: jtescher

tracing-opentelemetry/src/layer.rs

@@ -1,5 +1,5 @@
-use opentelemetry::api::{Context as OtelContext, IdGenerator, TraceContextExt};
-use opentelemetry::{api, sdk};
+use crate::PreSampledTracer;
+use opentelemetry::api::{self, Context as OtelContext, TraceContextExt};
 use std::any::TypeId;
 use std::fmt;
 use std::marker;
@@ -20,10 +20,8 @@ static SPAN_KIND_FIELD: &str = "otel.kind";
 ///
 /// [OpenTelemetry]: https://opentelemetry.io
 /// [tracing]: https://github.com/tokio-rs/tracing
-pub struct OpenTelemetryLayer<S, T: api::Tracer> {
+pub struct OpenTelemetryLayer<S, T> {
     tracer: T,
-    sampler: Box<dyn sdk::ShouldSample>,
-    id_generator: sdk::IdGenerator,
 
     get_context: WithContext,
     _registry: marker::PhantomData<S>,
@@ -34,7 +32,7 @@ where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
     fn default() -> Self {
-        OpenTelemetryLayer::new(api::NoopTracer {}, sdk::Sampler::AlwaysOn)
+        OpenTelemetryLayer::new(api::NoopTracer {})
     }
 }
 
@@ -50,8 +48,7 @@ where
 ///
 /// // Use the tracing subscriber `Registry`, or any other subscriber
 /// // that impls `LookupSpan`
-/// let subscriber = Registry::default()
-///     .with(tracing_opentelemetry::layer());
+/// let subscriber = Registry::default().with(tracing_opentelemetry::layer());
 /// # drop(subscriber);
 /// ```
 pub fn layer<S>() -> OpenTelemetryLayer<S, api::NoopTracer>
@@ -70,7 +67,7 @@ pub(crate) struct WithContext(
     fn(
         &tracing::Dispatch,
         &span::Id,
-        f: &mut dyn FnMut(&mut api::SpanBuilder, &dyn sdk::ShouldSample),
+        f: &mut dyn FnMut(&mut api::SpanBuilder, &dyn PreSampledTracer),
     ),
 );
 
@@ -81,63 +78,12 @@ impl WithContext {
         &self,
         dispatch: &'a tracing::Dispatch,
         id: &span::Id,
-        mut f: impl FnMut(&mut api::SpanBuilder, &dyn sdk::ShouldSample),
+        mut f: impl FnMut(&mut api::SpanBuilder, &dyn PreSampledTracer),
     ) {
         (self.0)(dispatch, id, &mut f)
     }
 }
 
-pub(crate) fn build_span_context(
-    builder: &mut api::SpanBuilder,
-    sampler: &dyn sdk::ShouldSample,
-) -> api::SpanContext {
-    let span_id = builder.span_id.expect("Builders must have id");
-    let (trace_id, trace_flags) = builder
-        .parent_context
-        .as_ref()
-        .filter(|parent_context| parent_context.is_valid())
-        .map(|parent_context| (parent_context.trace_id(), parent_context.trace_flags()))
-        .unwrap_or_else(|| {
-            let trace_id = builder.trace_id.expect("trace_id should exist");
-
-            // ensure sampling decision is recorded so all span contexts have consistent flags
-            let sampling_decision = if let Some(result) = builder.sampling_result.as_ref() {
-                result.decision.clone()
-            } else {
-                let mut result = sampler.should_sample(
-                    builder.parent_context.as_ref(),
-                    trace_id,
-                    &builder.name,
-                    builder
-                        .span_kind
-                        .as_ref()
-                        .unwrap_or(&api::SpanKind::Internal),
-                    builder.attributes.as_ref().unwrap_or(&Vec::new()),
-                    builder.links.as_ref().unwrap_or(&Vec::new()),
-                );
-
-                // Record additional attributes resulting from sampling
-                if let Some(attributes) = &mut builder.attributes {
-                    attributes.append(&mut result.attributes)
-                } else {
-                    builder.attributes = Some(result.attributes);
-                }
-
-                result.decision
-            };
-
-            let trace_flags = if sampling_decision == sdk::SamplingDecision::RecordAndSampled {
-                api::TRACE_FLAG_SAMPLED
-            } else {
-                0
-            };
-
-            (trace_id, trace_flags)
-        });
-
-    api::SpanContext::new(trace_id, span_id, trace_flags, false)
-}
-
 fn str_to_span_kind(s: &str) -> Option<api::SpanKind> {
     if s.eq_ignore_ascii_case("SERVER") {
         Some(api::SpanKind::Server)
@@ -235,13 +181,12 @@ impl<'a> field::Visit for SpanAttributeVisitor<'a> {
 impl<S, T> OpenTelemetryLayer<S, T>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
-    T: api::Tracer + 'static,
+    T: api::Tracer + PreSampledTracer + 'static,
 {
-    /// Set the [`Tracer`] and [`Sampler`] that this layer will use to produce and
-    /// track OpenTelemetry [`Span`]s.
+    /// Set the [`Tracer`] that this layer will use to produce and track
+    /// OpenTelemetry [`Span`]s.
     ///
     /// [`Tracer`]: https://docs.rs/opentelemetry/latest/opentelemetry/api/trace/tracer/trait.Tracer.html
-    /// [`Sampler`]: https://docs.rs/opentelemetry/latest/opentelemetry/api/trace/sampler/trait.Sampler.html
     /// [`Span`]: https://docs.rs/opentelemetry/latest/opentelemetry/api/trace/span/trait.Span.html
     ///
     /// # Examples
@@ -273,26 +218,17 @@ where
     /// // Get a tracer from the provider for a component
     /// let tracer = provider.get_tracer("component-name");
     ///
-    /// // The probability sampler can be used to export a percentage of spans
-    /// let sampler = sdk::Sampler::Probability(0.33);
-    ///
     /// // Create a layer with the configured tracer
-    /// let otel_layer = OpenTelemetryLayer::new(tracer, sampler);
+    /// let otel_layer = OpenTelemetryLayer::new(tracer);
     ///
     /// // Use the tracing subscriber `Registry`, or any other subscriber
     /// // that impls `LookupSpan`
-    /// let subscriber = Registry::default()
-    ///     .with(otel_layer);
+    /// let subscriber = Registry::default().with(otel_layer);
     /// # drop(subscriber);
     /// ```
-    pub fn new<Sampler>(tracer: T, sampler: Sampler) -> Self
-    where
-        Sampler: sdk::ShouldSample + 'static,
-    {
+    pub fn new(tracer: T) -> Self {
         OpenTelemetryLayer {
             tracer,
-            sampler: Box::new(sampler),
-            id_generator: sdk::IdGenerator::default(),
             get_context: WithContext(Self::get_context),
             _registry: marker::PhantomData,
         }
@@ -343,52 +279,15 @@ where
     /// ```
     pub fn with_tracer<Tracer>(self, tracer: Tracer) -> OpenTelemetryLayer<S, Tracer>
     where
-        Tracer: api::Tracer + 'static,
+        Tracer: api::Tracer + PreSampledTracer + 'static,
     {
         OpenTelemetryLayer {
             tracer,
-            sampler: self.sampler,
-            id_generator: self.id_generator,
             get_context: WithContext(OpenTelemetryLayer::<S, Tracer>::get_context),
             _registry: self._registry,
         }
     }
 
-    /// Set the [`Sampler`] to configure the logic around which [`Span`]s are
-    /// exported.
-    ///
-    /// [`Sampler`]: https://docs.rs/opentelemetry/latest/opentelemetry/api/trace/sampler/trait.Sampler.html
-    /// [`Span`]: https://docs.rs/opentelemetry/latest/opentelemetry/api/trace/span/trait.Span.html
-    ///
-    /// # Examples
-    ///
-    /// ```rust,no_run
-    /// use opentelemetry::sdk;
-    /// use tracing_subscriber::layer::SubscriberExt;
-    /// use tracing_subscriber::Registry;
-    ///
-    /// // The probability sampler can be used to export a percentage of spans
-    /// let sampler = sdk::Sampler::Probability(0.33);
-    ///
-    /// // Create a layer with the configured sampler
-    /// let otel_layer = tracing_opentelemetry::layer().with_sampler(sampler);
-    ///
-    /// // Use the tracing subscriber `Registry`, or any other subscriber
-    /// // that impls `LookupSpan`
-    /// let subscriber = Registry::default()
-    ///     .with(otel_layer);
-    /// # drop(subscriber);
-    /// ```
-    pub fn with_sampler<Sampler>(self, sampler: Sampler) -> Self
-    where
-        Sampler: sdk::ShouldSample + 'static,
-    {
-        OpenTelemetryLayer {
-            sampler: Box::new(sampler),
-            ..self
-        }
-    }
-
     /// Retrieve the parent OpenTelemetry [`SpanContext`] from the current
     /// tracing [`span`] through the [`Registry`]. This [`SpanContext`]
     /// links spans to their parent for proper hierarchical visualization.
@@ -407,14 +306,14 @@ where
             let mut extensions = span.extensions_mut();
             extensions
                 .get_mut::<api::SpanBuilder>()
-                .map(|builder| build_span_context(builder, self.sampler.as_ref()))
+                .map(|builder| self.tracer.sampled_span_context(builder))
         // Else if the span is inferred from context, look up any available current span.
         } else if attrs.is_contextual() {
             ctx.lookup_current().and_then(|span| {
                 let mut extensions = span.extensions_mut();
                 extensions
                     .get_mut::<api::SpanBuilder>()
-                    .map(|builder| build_span_context(builder, self.sampler.as_ref()))
+                    .map(|builder| self.tracer.sampled_span_context(builder))
             })
         // Explicit root spans should have no parent context.
         } else {
@@ -425,7 +324,7 @@ where
     fn get_context(
         dispatch: &tracing::Dispatch,
         id: &span::Id,
-        f: &mut dyn FnMut(&mut api::SpanBuilder, &dyn sdk::ShouldSample),
+        f: &mut dyn FnMut(&mut api::SpanBuilder, &dyn PreSampledTracer),
     ) {
         let subscriber = dispatch
             .downcast_ref::<S>()
@@ -439,15 +338,15 @@ where
 
         let mut extensions = span.extensions_mut();
         if let Some(builder) = extensions.get_mut::<api::SpanBuilder>() {
-            f(builder, layer.sampler.as_ref());
+            f(builder, &layer.tracer);
         }
     }
 }
 
 impl<S, T> Layer<S> for OpenTelemetryLayer<S, T>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
-    T: api::Tracer + 'static,
+    T: api::Tracer + PreSampledTracer + 'static,
 {
     /// Creates an [OpenTelemetry `Span`] for the corresponding [tracing `Span`].
     ///
@@ -462,7 +361,7 @@ where
             .span_builder(attrs.metadata().name())
             .with_start_time(SystemTime::now())
             // Eagerly assign span id so children have stable parent id
-            .with_span_id(self.id_generator.new_span_id());
+            .with_span_id(self.tracer.new_span_id());
 
         // Set optional parent span context from attrs
         builder.parent_context = self.parent_span_context(attrs, &ctx);
@@ -473,7 +372,7 @@ where
             if existing_otel_span_context.is_valid() {
                 builder.trace_id = Some(existing_otel_span_context.trace_id());
             } else {
-                builder.trace_id = Some(self.id_generator.new_trace_id());
+                builder.trace_id = Some(self.tracer.new_trace_id());
             }
         }
 
@@ -507,7 +406,7 @@ where
             .get_mut::<api::SpanBuilder>()
             .expect("Missing SpanBuilder span extensions");
 
-        let follows_context = build_span_context(follows_builder, self.sampler.as_ref());
+        let follows_context = self.tracer.sampled_span_context(follows_builder);
         let follows_link = api::Link::new(follows_context, Vec::new());
         if let Some(ref mut links) = builder.links {
             links.push(follows_link);
@@ -613,6 +512,18 @@ mod tests {
         }
     }
 
+    impl PreSampledTracer for TestTracer {
+        fn sampled_span_context(&self, _builder: &mut api::SpanBuilder) -> api::SpanContext {
+            api::SpanContext::empty_context()
+        }
+        fn new_trace_id(&self) -> api::TraceId {
+            api::TraceId::invalid()
+        }
+        fn new_span_id(&self) -> api::SpanId {
+            api::SpanId::invalid()
+        }
+    }
+
     #[derive(Debug, Clone)]
     struct TestSpan(api::SpanContext);
     impl api::Span for TestSpan {

tracing-opentelemetry/src/lib.rs

@@ -46,7 +46,7 @@
 //! may still occur on the path to 1.0. You can follow the changes via the
 //! [spec repository] to track progress toward stabilization.
 //!
-//! [spec repository]: (https://github.com/open-telemetry/opentelemetry-specification)
+//! [spec repository]: https://github.com/open-telemetry/opentelemetry-specification
 //!
 //! ## Examples
 //!
@@ -100,6 +100,9 @@
 mod layer;
 /// Span extension which enables OpenTelemetry context management.
 mod span_ext;
+/// Protocols for OpenTelemetry Tracers that are compatible with Tracing
+mod tracer;
 
 pub use layer::{layer, OpenTelemetryLayer};
 pub use span_ext::OpenTelemetrySpanExt;
+pub use tracer::PreSampledTracer;

tracing-opentelemetry/src/span_ext.rs

@@ -1,4 +1,4 @@
-use crate::layer::{build_span_context, WithContext};
+use crate::layer::WithContext;
 use opentelemetry::api;
 use opentelemetry::api::TraceContextExt;
 
@@ -76,7 +76,7 @@ impl OpenTelemetrySpanExt for tracing::Span {
     fn set_parent(&self, parent_context: &api::Context) {
         self.with_subscriber(move |(id, subscriber)| {
             if let Some(get_context) = subscriber.downcast_ref::<WithContext>() {
-                get_context.with_context(subscriber, id, move |builder, _sampler| {
+                get_context.with_context(subscriber, id, move |builder, _tracer| {
                     builder.parent_context = parent_context.remote_span_context().cloned()
                 });
             }
@@ -87,8 +87,8 @@ impl OpenTelemetrySpanExt for tracing::Span {
         let mut span_context = None;
         self.with_subscriber(|(id, subscriber)| {
             if let Some(get_context) = subscriber.downcast_ref::<WithContext>() {
-                get_context.with_context(subscriber, id, |builder, sampler| {
-                    span_context = Some(build_span_context(builder, sampler));
+                get_context.with_context(subscriber, id, |builder, tracer| {
+                    span_context = Some(tracer.sampled_span_context(builder));
                 })
             }
         });