Sample traces based on the parent (#4514)

Author: sandhose

crates/cli/src/telemetry.rs

@@ -119,10 +119,14 @@ fn otlp_tracer_provider(
     let batch_processor =
         BatchSpanProcessor::builder(exporter, opentelemetry_sdk::runtime::Tokio).build();
 
+    // We sample traces based on the parent if we have one, and if not, we
+    // sample a ratio based on the configured sample rate
+    let sampler = Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(sample_rate)));
+
     let tracer_provider = SdkTracerProvider::builder()
         .with_span_processor(batch_processor)
         .with_resource(resource())
-        .with_sampler(Sampler::TraceIdRatioBased(sample_rate))
+        .with_sampler(sampler)
         .build();
 
     Ok(tracer_provider)

crates/context/src/fmt.rs

@@ -4,7 +4,10 @@
 // Please see LICENSE in the repository root for full details.
 
 use console::{Color, Style};
-use opentelemetry::{TraceId, trace::TraceContextExt};
+use opentelemetry::{
+    TraceId,
+    trace::{SamplingDecision, TraceContextExt},
+};
 use tracing::{Level, Subscriber};
 use tracing_opentelemetry::OtelData;
 use tracing_subscriber::{
@@ -128,18 +131,29 @@ where
         // If we have a OTEL span, we can add the trace ID to the end of the log line
         if let Some(span) = ctx.lookup_current() {
             if let Some(otel) = span.extensions().get::<OtelData>() {
-                // If it is the root span, the trace ID will be in the span builder. Else, it
-                // will be in the parent OTEL context
-                let trace_id = otel
+                let parent_cx_span = otel.parent_cx.span();
+                let sc = parent_cx_span.span_context();
+
+                // Check if the span is sampled, first from the span builder,
+                // then from the parent context if nothing is set there
+                if otel
                     .builder
-                    .trace_id
-                    .unwrap_or_else(|| otel.parent_cx.span().span_context().trace_id());
-                if trace_id != TraceId::INVALID {
-                    let label = Style::new()
-                        .italic()
-                        .force_styling(ansi)
-                        .apply_to("trace.id");
-                    write!(&mut writer, " {label}={trace_id}")?;
+                    .sampling_result
+                    .as_ref()
+                    .map_or(sc.is_sampled(), |r| {
+                        r.decision == SamplingDecision::RecordAndSample
+                    })
+                {
+                    // If it is the root span, the trace ID will be in the span builder. Else, it
+                    // will be in the parent OTEL context
+                    let trace_id = otel.builder.trace_id.unwrap_or(sc.trace_id());
+                    if trace_id != TraceId::INVALID {
+                        let label = Style::new()
+                            .italic()
+                            .force_styling(ansi)
+                            .apply_to("trace.id");
+                        write!(&mut writer, " {label}={trace_id}")?;
+                    }
                 }
             }
         }