fix upward propagation

grafana · Dec 5, 2024 · e131a92 · e131a92
1 parent 4d1fa19
commit e131a92
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 8 deletions.
diff --git a/custom/src/main/java/com/grafana/extensions/sampler/DynamicSampler.java b/custom/src/main/java/com/grafana/extensions/sampler/DynamicSampler.java
@@ -116,9 +116,11 @@ private String evaluateReason(ReadableSpan span, String traceId) {
     return null;
   }
 
-  // visible for testing
-  public void clear() {
+  public void resetForTest() {
     sampledTraces.clear();
+    spansByTrace.clear();
+    movingAvgs.clear();
+    firstSampledCallback.clear();
   }
 
   void removeTrace(String traceId) {

diff --git a/custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderReaderTest.java b/custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderReaderTest.java
@@ -14,6 +14,7 @@
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
 import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
+import io.opentelemetry.sdk.trace.ReadableSpan;
 import java.util.Collections;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
@@ -33,7 +34,7 @@ static void initialize() {
 
   @BeforeEach
   void setUp() {
-    DynamicSampler.getInstance().clear();
+    DynamicSampler.getInstance().resetForTest();
   }
 
   @Test
@@ -44,6 +45,7 @@ void notSampled() {
     testing.runWithSpan(
         spanName,
         () -> {
+          DynamicSampler.getInstance().registerNewSpan((ReadableSpan) Span.current());
           String traceId = Span.current().getSpanContext().getTraceId();
           String serverTiming = ServerTimingHeaderCustomizer.toHeaderValue(Context.current());
           serverTimingHeaderReader.consume(
@@ -58,12 +60,13 @@ void sampled() {
     testing.runWithSpan(
         "server",
         () -> {
+          DynamicSampler.getInstance().registerNewSpan((ReadableSpan) Span.current());
           String traceId = Span.current().getSpanContext().getTraceId();
           DynamicSampler.getInstance().setSampled(traceId, "test");
           String serverTiming = ServerTimingHeaderCustomizer.toHeaderValue(Context.current());
 
           // remove the traceId to see that it is added back by the reader
-          DynamicSampler.getInstance().clear();
+          DynamicSampler.getInstance().resetForTest();
           serverTimingHeaderReader.consume(
               new StringHttpCommonAttributesGetter(serverTiming), "request", "response");
           assertThat(DynamicSampler.getInstance().getSampledTraces()).contains(traceId);

diff --git a/custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderTest.java b/custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderTest.java
@@ -18,6 +18,8 @@
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
 import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -43,7 +45,7 @@ static void initialize() {
   @BeforeEach
   void setUp() {
     ServerTimingHeaderCustomizer.enabled = true;
-    DynamicSampler.getInstance().clear();
+    DynamicSampler.getInstance().resetForTest();
   }
 
   @Test
@@ -62,8 +64,11 @@ void shouldSetHeaders() {
     assertSetHeader("00", span -> {});
     assertSetHeader(
         "01",
-        span ->
-            DynamicSampler.getInstance().setSampled(span.getSpanContext().getTraceId(), "test"));
+        span -> {
+          DynamicSampler.getInstance().registerNewSpan((ReadableSpan) Span.current());
+          DynamicSampler.getInstance().setSampled(span.getSpanContext().getTraceId(), "test");
+          DynamicSampler.getInstance().evaluateSampled((ReadWriteSpan) span);
+        });
   }
 
   private void assertSetHeader(String traceFlags, Consumer<Span> spanConsumer) {
@@ -73,11 +78,15 @@ private void assertSetHeader(String traceFlags, Consumer<Span> spanConsumer) {
         testing.runWithSpan(
             "server",
             () -> {
-              spanConsumer.accept(Span.current());
               serverTiming.customize(Context.current(), headers, Map::put);
+              spanConsumer.accept(Span.current());
               return Span.current().getSpanContext();
             });
 
+    if (traceFlags.equals("00")) {
+      assertThat(headers).isEmpty();
+      return;
+    }
     assertThat(headers).hasSize(2);
 
     var serverTimingHeaderValue =