Add option to export unsampled spans from span processors (#6057)

Co-authored-by: jack-berg <[email protected]>
open-telemetry · Jan 4, 2024 · 07351a2 · 07351a2
1 parent f4b5bbe
commit 07351a2
Show file tree

Hide file tree

Showing 8 changed files with 185 additions and 59 deletions.
diff --git a/docs/apidiffs/current_vs_latest/opentelemetry-sdk-trace.txt b/docs/apidiffs/current_vs_latest/opentelemetry-sdk-trace.txt
@@ -1,2 +1,12 @@
 Comparing source compatibility of  against 
-No changes.
+***  MODIFIED CLASS: PUBLIC FINAL io.opentelemetry.sdk.trace.export.BatchSpanProcessorBuilder  (not serializable)
+	===  CLASS FILE FORMAT VERSION: 52.0 <- 52.0
+	+++  NEW METHOD: PUBLIC(+) io.opentelemetry.sdk.trace.export.BatchSpanProcessorBuilder setExportUnsampledSpans(boolean)
+***  MODIFIED CLASS: PUBLIC FINAL io.opentelemetry.sdk.trace.export.SimpleSpanProcessor  (not serializable)
+	===  CLASS FILE FORMAT VERSION: 52.0 <- 52.0
+	+++  NEW METHOD: PUBLIC(+) STATIC(+) io.opentelemetry.sdk.trace.export.SimpleSpanProcessorBuilder builder(io.opentelemetry.sdk.trace.export.SpanExporter)
++++  NEW CLASS: PUBLIC(+) FINAL(+) io.opentelemetry.sdk.trace.export.SimpleSpanProcessorBuilder  (not serializable)
+	+++  CLASS FILE FORMAT VERSION: 52.0 <- n.a.
+	+++  NEW SUPERCLASS: java.lang.Object
+	+++  NEW METHOD: PUBLIC(+) io.opentelemetry.sdk.trace.export.SimpleSpanProcessor build()
+	+++  NEW METHOD: PUBLIC(+) io.opentelemetry.sdk.trace.export.SimpleSpanProcessorBuilder setExportUnsampledSpans(boolean)
diff --git a/sdk/all/src/test/java/io/opentelemetry/sdk/OpenTelemetrySdkTest.java b/sdk/all/src/test/java/io/opentelemetry/sdk/OpenTelemetrySdkTest.java
@@ -415,7 +415,7 @@ void stringRepresentation() {
                 + "resource=Resource{schemaUrl=null, attributes={service.name=\"otel-test\"}}, "
                 + "spanLimitsSupplier=SpanLimitsValue{maxNumberOfAttributes=128, maxNumberOfEvents=128, maxNumberOfLinks=128, maxNumberOfAttributesPerEvent=128, maxNumberOfAttributesPerLink=128, maxAttributeValueLength=2147483647}, "
                 + "sampler=ParentBased{root:AlwaysOnSampler,remoteParentSampled:AlwaysOnSampler,remoteParentNotSampled:AlwaysOffSampler,localParentSampled:AlwaysOnSampler,localParentNotSampled:AlwaysOffSampler}, "
-                + "spanProcessor=SimpleSpanProcessor{spanExporter=MultiSpanExporter{spanExporters=[MockSpanExporter{}, MockSpanExporter{}]}}"
+                + "spanProcessor=SimpleSpanProcessor{spanExporter=MultiSpanExporter{spanExporters=[MockSpanExporter{}, MockSpanExporter{}]}, exportUnsampledSpans=false}"
                 + "}, "
                 + "meterProvider=SdkMeterProvider{"
                 + "clock=SystemClock{}, "

diff --git a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessor.java b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessor.java
@@ -53,6 +53,7 @@ public final class BatchSpanProcessor implements SpanProcessor {
       AttributeKey.booleanKey("dropped");
   private static final String SPAN_PROCESSOR_TYPE_VALUE = BatchSpanProcessor.class.getSimpleName();
 
+  private final boolean exportUnsampledSpans;
   private final Worker worker;
   private final AtomicBoolean isShutdown = new AtomicBoolean(false);
 
@@ -69,11 +70,13 @@ public static BatchSpanProcessorBuilder builder(SpanExporter spanExporter) {
 
   BatchSpanProcessor(
       SpanExporter spanExporter,
+      boolean exportUnsampledSpans,
       MeterProvider meterProvider,
       long scheduleDelayNanos,
       int maxQueueSize,
       int maxExportBatchSize,
       long exporterTimeoutNanos) {
+    this.exportUnsampledSpans = exportUnsampledSpans;
     this.worker =
         new Worker(
             spanExporter,
@@ -96,10 +99,9 @@ public boolean isStartRequired() {
 
   @Override
   public void onEnd(ReadableSpan span) {
-    if (span == null || !span.getSpanContext().isSampled()) {
-      return;
+    if (span != null && (exportUnsampledSpans || span.getSpanContext().isSampled())) {
+      worker.addSpan(span);
     }
-    worker.addSpan(span);
   }
 
   @Override
@@ -135,6 +137,8 @@ public String toString() {
     return "BatchSpanProcessor{"
         + "spanExporter="
         + worker.spanExporter
+        + ", exportUnsampledSpans="
+        + exportUnsampledSpans
         + ", scheduleDelayNanos="
         + worker.scheduleDelayNanos
         + ", maxExportBatchSize="

diff --git a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorBuilder.java b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorBuilder.java
@@ -25,6 +25,7 @@ public final class BatchSpanProcessorBuilder {
   static final int DEFAULT_EXPORT_TIMEOUT_MILLIS = 30_000;
 
   private final SpanExporter spanExporter;
+  private boolean exportUnsampledSpans;
   private long scheduleDelayNanos = TimeUnit.MILLISECONDS.toNanos(DEFAULT_SCHEDULE_DELAY_MILLIS);
   private int maxQueueSize = DEFAULT_MAX_QUEUE_SIZE;
   private int maxExportBatchSize = DEFAULT_MAX_EXPORT_BATCH_SIZE;
@@ -35,6 +36,15 @@ public final class BatchSpanProcessorBuilder {
     this.spanExporter = requireNonNull(spanExporter, "spanExporter");
   }
 
+  /**
+   * Sets whether unsampled spans should be exported. If unset, defaults to exporting only sampled
+   * spans.
+   */
+  public BatchSpanProcessorBuilder setExportUnsampledSpans(boolean exportUnsampledSpans) {
+    this.exportUnsampledSpans = exportUnsampledSpans;
+    return this;
+  }
+
   /**
    * Sets the delay interval between two consecutive exports. If unset, defaults to {@value
    * DEFAULT_SCHEDULE_DELAY_MILLIS}ms.
@@ -146,6 +156,7 @@ int getMaxExportBatchSize() {
   public BatchSpanProcessor build() {
     return new BatchSpanProcessor(
         spanExporter,
+        exportUnsampledSpans,
         meterProvider,
         scheduleDelayNanos,
         maxQueueSize,

diff --git a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/SimpleSpanProcessor.java b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/SimpleSpanProcessor.java
@@ -36,7 +36,7 @@ public final class SimpleSpanProcessor implements SpanProcessor {
   private static final Logger logger = Logger.getLogger(SimpleSpanProcessor.class.getName());
 
   private final SpanExporter spanExporter;
-  private final boolean sampled;
+  private final boolean exportUnsampledSpans;
   private final Set<CompletableResultCode> pendingExports =
       Collections.newSetFromMap(new ConcurrentHashMap<>());
   private final AtomicBoolean isShutdown = new AtomicBoolean(false);
@@ -53,12 +53,17 @@ public final class SimpleSpanProcessor implements SpanProcessor {
    */
   public static SpanProcessor create(SpanExporter exporter) {
     requireNonNull(exporter, "exporter");
-    return new SimpleSpanProcessor(exporter, /* sampled= */ true);
+    return builder(exporter).build();
   }
 
-  SimpleSpanProcessor(SpanExporter spanExporter, boolean sampled) {
+  public static SimpleSpanProcessorBuilder builder(SpanExporter exporter) {
+    requireNonNull(exporter, "exporter");
+    return new SimpleSpanProcessorBuilder(exporter);
+  }
+
+  SimpleSpanProcessor(SpanExporter spanExporter, boolean exportUnsampledSpans) {
     this.spanExporter = requireNonNull(spanExporter, "spanExporter");
-    this.sampled = sampled;
+    this.exportUnsampledSpans = exportUnsampledSpans;
   }
 
   @Override
@@ -73,22 +78,21 @@ public boolean isStartRequired() {
 
   @Override
   public void onEnd(ReadableSpan span) {
-    if (sampled && !span.getSpanContext().isSampled()) {
-      return;
-    }
-    try {
-      List<SpanData> spans = Collections.singletonList(span.toSpanData());
-      CompletableResultCode result = spanExporter.export(spans);
-      pendingExports.add(result);
-      result.whenComplete(
-          () -> {
-            pendingExports.remove(result);
-            if (!result.isSuccess()) {
-              logger.log(Level.FINE, "Exporter failed");
-            }
-          });
-    } catch (RuntimeException e) {
-      logger.log(Level.WARNING, "Exporter threw an Exception", e);
+    if (span != null && (exportUnsampledSpans || span.getSpanContext().isSampled())) {
+      try {
+        List<SpanData> spans = Collections.singletonList(span.toSpanData());
+        CompletableResultCode result = spanExporter.export(spans);
+        pendingExports.add(result);
+        result.whenComplete(
+            () -> {
+              pendingExports.remove(result);
+              if (!result.isSuccess()) {
+                logger.log(Level.FINE, "Exporter failed");
+              }
+            });
+      } catch (RuntimeException e) {
+        logger.log(Level.WARNING, "Exporter threw an Exception", e);
+      }
     }
   }
 
@@ -128,6 +132,11 @@ public CompletableResultCode forceFlush() {
 
   @Override
   public String toString() {
-    return "SimpleSpanProcessor{" + "spanExporter=" + spanExporter + '}';
+    return "SimpleSpanProcessor{"
+        + "spanExporter="
+        + spanExporter
+        + ", exportUnsampledSpans="
+        + exportUnsampledSpans
+        + '}';
   }
 }
diff --git a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/SimpleSpanProcessorBuilder.java b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/export/SimpleSpanProcessorBuilder.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.trace.export;
+
+import static java.util.Objects.requireNonNull;
+
+/** Builder class for {@link SimpleSpanProcessor}. */
+public final class SimpleSpanProcessorBuilder {
+  private final SpanExporter spanExporter;
+  private boolean exportUnsampledSpans;
+
+  SimpleSpanProcessorBuilder(SpanExporter spanExporter) {
+    this.spanExporter = requireNonNull(spanExporter, "spanExporter");
+  }
+
+  /**
+   * Sets whether unsampled spans should be exported. If unset, defaults to exporting only sampled
+   * spans.
+   */
+  public SimpleSpanProcessorBuilder setExportUnsampledSpans(boolean exportUnsampledSpans) {
+    this.exportUnsampledSpans = exportUnsampledSpans;
+    return this;
+  }
+
+  /**
+   * Returns a new {@link SimpleSpanProcessor} with the configuration of this builder.
+   *
+   * @return a new {@link SimpleSpanProcessor}.
+   */
+  public SimpleSpanProcessor build() {
+    return new SimpleSpanProcessor(spanExporter, exportUnsampledSpans);
+  }
+}
diff --git a/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorTest.java b/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorTest.java
@@ -517,6 +517,38 @@ void exportNotSampledSpans_recordOnly() {
     assertThat(exported).containsExactly(span.toSpanData());
   }
 
+  @Test
+  void exportUnsampledSpans_recordOnly() {
+    WaitingSpanExporter waitingSpanExporter =
+        new WaitingSpanExporter(1, CompletableResultCode.ofSuccess());
+
+    when(mockSampler.shouldSample(any(), any(), any(), any(), any(), anyList()))
+        .thenReturn(SamplingResult.recordOnly());
+    sdkTracerProvider =
+        SdkTracerProvider.builder()
+            .addSpanProcessor(
+                BatchSpanProcessor.builder(waitingSpanExporter)
+                    .setExportUnsampledSpans(true)
+                    .setScheduleDelay(MAX_SCHEDULE_DELAY_MILLIS, TimeUnit.MILLISECONDS)
+                    .build())
+            .setSampler(mockSampler)
+            .build();
+
+    ReadableSpan span1 = createEndedSpan(SPAN_NAME_1);
+    when(mockSampler.shouldSample(any(), any(), any(), any(), any(), anyList()))
+        .thenReturn(SamplingResult.recordAndSample());
+    ReadableSpan span2 = createEndedSpan(SPAN_NAME_2);
+
+    // Spans are recorded and exported in the same order as they are ended, we test that a non
+    // exported span is not exported by creating and ending a sampled span after a non sampled span
+    // and checking that the first exported span is the sampled span (the non sampled did not get
+    // exported).
+    List<SpanData> exported = waitingSpanExporter.waitForExport();
+    // Need to check this because otherwise the variable span1 is unused, other option is to not
+    // have a span1 variable.
+    assertThat(exported).containsExactly(span1.toSpanData(), span2.toSpanData());
+  }
+
   @Test
   @Timeout(10)
   @SuppressLogger(SdkTracerProvider.class)
@@ -569,6 +601,7 @@ void stringRepresentation() {
         .hasToString(
             "BatchSpanProcessor{"
                 + "spanExporter=mockSpanExporter, "
+                + "exportUnsampledSpans=false, "
                 + "scheduleDelayNanos=5000000000, "
                 + "maxExportBatchSize=512, "
                 + "exporterTimeoutNanos=30000000000}");

diff --git a/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/SimpleSpanProcessorTest.java b/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/SimpleSpanProcessorTest.java
@@ -59,7 +59,12 @@ class SimpleSpanProcessorTest {
           SpanId.getInvalid(),
           TraceFlags.getSampled(),
           TraceState.getDefault());
-  private static final SpanContext NOT_SAMPLED_SPAN_CONTEXT = SpanContext.getInvalid();
+  private static final SpanContext NOT_SAMPLED_SPAN_CONTEXT =
+      SpanContext.create(
+          TraceId.getInvalid(),
+          SpanId.getInvalid(),
+          TraceFlags.getDefault(),
+          TraceState.getDefault());
 
   private SpanProcessor simpleSampledSpansProcessor;
 
@@ -100,29 +105,29 @@ void onEndSync_NotSampledSpan() {
   }
 
   @Test
-  void onEndSync_OnlySampled_NotSampledSpan() {
+  void onEndSync_ExportUnsampledSpans_NotSampledSpan() {
+    SpanData spanData = TestUtils.makeBasicSpan();
     when(readableSpan.getSpanContext()).thenReturn(NOT_SAMPLED_SPAN_CONTEXT);
-    when(readableSpan.toSpanData())
-        .thenReturn(TestUtils.makeBasicSpan())
-        .thenThrow(new RuntimeException());
-    SpanProcessor simpleSpanProcessor = SimpleSpanProcessor.create(spanExporter);
+    when(readableSpan.toSpanData()).thenReturn(spanData);
+    SpanProcessor simpleSpanProcessor =
+        SimpleSpanProcessor.builder(spanExporter).setExportUnsampledSpans(true).build();
     simpleSpanProcessor.onEnd(readableSpan);
-    verifyNoInteractions(spanExporter);
+    verify(spanExporter).export(Collections.singletonList(spanData));
   }
 
   @Test
-  void onEndSync_OnlySampled_SampledSpan() {
+  void onEndSync_ExportUnsampledSpans_SampledSpan() {
+    SpanData spanData = TestUtils.makeBasicSpan();
     when(readableSpan.getSpanContext()).thenReturn(SAMPLED_SPAN_CONTEXT);
-    when(readableSpan.toSpanData())
-        .thenReturn(TestUtils.makeBasicSpan())
-        .thenThrow(new RuntimeException());
-    SpanProcessor simpleSpanProcessor = SimpleSpanProcessor.create(spanExporter);
+    when(readableSpan.toSpanData()).thenReturn(spanData);
+    SpanProcessor simpleSpanProcessor =
+        SimpleSpanProcessor.builder(spanExporter).setExportUnsampledSpans(true).build();
     simpleSpanProcessor.onEnd(readableSpan);
-    verify(spanExporter).export(Collections.singletonList(TestUtils.makeBasicSpan()));
+    verify(spanExporter).export(Collections.singletonList(spanData));
   }
 
   @Test
-  void tracerSdk_NotSampled_Span() {
+  void tracerSdk_SampledSpan() {
     WaitingSpanExporter waitingSpanExporter =
         new WaitingSpanExporter(1, CompletableResultCode.ofSuccess());
 
@@ -159,25 +164,43 @@ void tracerSdk_NotSampled_Span() {
   }
 
   @Test
-  void tracerSdk_NotSampled_RecordingEventsSpan() {
-    // TODO(bdrutu): Fix this when Sampler return RECORD_ONLY option.
-    /*
-    tracer.addSpanProcessor(
-        BatchSpanProcessor.builder(waitingSpanExporter)
-            .setScheduleDelayMillis(MAX_SCHEDULE_DELAY_MILLIS)
-            .reportOnlySampled(false)
-            .build());
-
-    io.opentelemetry.trace.Span span =
-        tracer
-            .spanBuilder("FOO")
-            .setSampler(Samplers.neverSample())
-            .startSpanWithSampler();
-    span.end();
-
-    List<SpanData> exported = waitingSpanExporter.waitForExport(1);
-    assertThat(exported).containsExactly(((ReadableSpan) span).toSpanData());
-    */
+  void tracerSdk_ExportUnsampledSpans_NotSampledSpan() {
+    WaitingSpanExporter waitingSpanExporter =
+        new WaitingSpanExporter(1, CompletableResultCode.ofSuccess());
+
+    SdkTracerProvider sdkTracerProvider =
+        SdkTracerProvider.builder()
+            .addSpanProcessor(
+                SimpleSpanProcessor.builder(waitingSpanExporter)
+                    .setExportUnsampledSpans(true)
+                    .build())
+            .setSampler(mockSampler)
+            .build();
+
+    when(mockSampler.shouldSample(any(), any(), any(), any(), any(), anyList()))
+        .thenReturn(SamplingResult.drop());
+
+    try {
+      Tracer tracer = sdkTracerProvider.get(getClass().getName());
+      tracer.spanBuilder(SPAN_NAME).startSpan();
+      tracer.spanBuilder(SPAN_NAME).startSpan();
+
+      when(mockSampler.shouldSample(any(), any(), any(), any(), any(), anyList()))
+          .thenReturn(SamplingResult.recordOnly());
+      Span span = tracer.spanBuilder(SPAN_NAME).startSpan();
+      span.end();
+
+      // Spans are recorded and exported in the same order as they are ended, we test that a non
+      // sampled span is not exported by creating and ending a sampled span after a non sampled span
+      // and checking that the first exported span is the sampled span (the non sampled did not get
+      // exported).
+      List<SpanData> exported = waitingSpanExporter.waitForExport();
+      // Need to check this because otherwise the variable span1 is unused, other option is to not
+      // have a span1 variable.
+      assertThat(exported).containsExactly(((ReadableSpan) span).toSpanData());
+    } finally {
+      sdkTracerProvider.shutdown();
+    }
   }
 
   @Test