test: configure dynamic sampler

custom/src/main/java/com/grafana/extensions/sampler/DynamicSampler.java

@@ -20,14 +20,14 @@ public class DynamicSampler {
   private final Set<String> sampledTraces = new ConcurrentSkipListSet<>();
   public static final Logger logger = Logger.getLogger(DynamicSampler.class.getName());
   private final int windowSize;
-  private final double thresholdVal;
+  private final double threshold;
   private final Map<String, MovingAverage> movingAvgs = new ConcurrentHashMap<>();
   private static DynamicSampler INSTANCE;
 
   private DynamicSampler(ConfigProperties properties) {
     // read properties and configure dynamic sampling
-    this.thresholdVal = properties.getDouble("threshold", 1.5);
-    this.windowSize = properties.getInt("window", 5);
+    this.threshold = properties.getDouble("threshold", 1.3); // 30%
+    this.windowSize = properties.getInt("window", 3);
   }
 
   public static void configure(ConfigProperties properties) {
@@ -42,6 +42,11 @@ public void setSampled(String traceId) {
     sampledTraces.add(traceId);
   }
 
+  // for testing
+  public void setMovingAvg(String spanName, MovingAverage ma) {
+    this.movingAvgs.put(spanName, ma);
+  }
+
   boolean isSampled(String traceId) {
     return sampledTraces.contains(traceId);
   }
@@ -74,7 +79,7 @@ boolean shouldSample(ReadableSpan span) {
         Level.INFO,
         "spanName {0} - windowSize {1}: {2}",
         new Object[] {span.getName(), windowSize, span.getAttributes()});
-    long duration = (span.getLatencyNanos()) / 1_000_000;
+    long duration = span.getLatencyNanos();
     MovingAverage currMovingAvg =
         movingAvgs.computeIfAbsent(spanName, ma -> new MovingAverage(windowSize));
     currMovingAvg.add(duration);
@@ -83,12 +88,16 @@ boolean shouldSample(ReadableSpan span) {
       logger.log(
           Level.INFO,
           "avg {0} * threshold {1} = {2}, duration {3}",
-          new Object[] {avg, thresholdVal, avg * thresholdVal, duration});
+          new Object[] {avg, threshold, avg * threshold, duration});
       // discard
-      if (duration < avg * thresholdVal) {
+      if (duration < avg * threshold) {
         return false;
       }
     }
+    logger.log(
+        Level.INFO,
+        "sending span part of Trace: {0} - {1}",
+        new Object[] {span.toSpanData().getTraceId(), duration});
     return true;
   }
 }

custom/src/main/java/com/grafana/extensions/util/MovingAverage.java

@@ -7,6 +7,7 @@
 
 import java.util.LinkedList;
 import java.util.Queue;
+import java.util.concurrent.ThreadLocalRandom;
 
 public class MovingAverage {
   private final Queue<Long> window = new LinkedList<>();
@@ -17,6 +18,14 @@ public MovingAverage(int size) {
     this.size = size;
   }
 
+  public static MovingAverage getPrepopulatedMovingAvgForTest(int size, int lowerBound) {
+    MovingAverage ma = new MovingAverage(size);
+    for(int i=0;i<size;i++) {
+      ma.add(ThreadLocalRandom.current().nextLong(lowerBound, 30_000_000));
+    }
+    return ma;
+  }
+
   public int getCount() {
     return this.window.size();
   }

custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderReaderTest.java

@@ -8,10 +8,14 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.grafana.extensions.sampler.DynamicSampler;
+import com.grafana.extensions.util.MovingAverage;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
+import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
+import java.util.Collections;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -22,21 +26,30 @@ class ServerTimingHeaderReaderTest {
 
   @RegisterExtension InstrumentationExtension testing = LibraryInstrumentationExtension.create();
 
+  @BeforeAll
+  static void initialize() {
+    DynamicSampler.configure(DefaultConfigProperties.createFromMap(Collections.emptyMap()));
+  }
+
   @BeforeEach
   void setUp() {
     DynamicSampler.getInstance().clear();
   }
 
   @Test
   void notSampled() {
+    String spanName = "server";
+    MovingAverage testMovingAvg = MovingAverage.getPrepopulatedMovingAvgForTest(3, 11_900_000);
+    DynamicSampler.getInstance().setMovingAvg(spanName, testMovingAvg);
     testing.runWithSpan(
-        "server",
+        spanName,
         () -> {
+          String traceId = Span.current().getSpanContext().getTraceId();
           String serverTiming = ServerTimingHeaderCustomizer.toHeaderValue(Context.current());
-
           serverTimingHeaderReader.consume(
               new StringHttpCommonAttributesGetter(serverTiming), "request", "response");
-          assertThat(DynamicSampler.getInstance().getSampledTraces()).isEmpty();
+          System.out.println(serverTiming);
+          assertThat(DynamicSampler.getInstance().getSampledTraces()).doesNotContain(traceId);
         });
   }
 

custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderTest.java

@@ -10,15 +10,19 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.grafana.extensions.sampler.DynamicSampler;
+import com.grafana.extensions.util.MovingAverage;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
 import io.opentelemetry.sdk.resources.Resource;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Consumer;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -31,6 +35,11 @@ class ServerTimingHeaderTest {
 
   record TestCase(Resource want, ConfigProperties config) {}
 
+  @BeforeAll
+  static void initialize() {
+    DynamicSampler.configure(DefaultConfigProperties.createFromMap(Collections.emptyMap()));
+  }
+
   @BeforeEach
   void setUp() {
     ServerTimingHeaderCustomizer.enabled = true;
@@ -48,6 +57,8 @@ void shouldNotSetAnyHeadersWithoutValidCurrentSpan() {
 
   @Test
   void shouldSetHeaders() {
+    MovingAverage testMovingAvg = MovingAverage.getPrepopulatedMovingAvgForTest(3, 11_900_000);
+    DynamicSampler.getInstance().setMovingAvg("server", testMovingAvg);
     assertSetHeader("00", span -> {});
     assertSetHeader(
         "01", span -> DynamicSampler.getInstance().setSampled(span.getSpanContext().getTraceId()));
@@ -75,8 +86,9 @@ private void assertSetHeader(String traceFlags, Consumer<Span> spanConsumer) {
             + "-"
             + traceFlags
             + "\"";
+
     assertThat(headers)
-        .containsEntry(SERVER_TIMING, serverTimingHeaderValue)
+                .containsEntry(SERVER_TIMING, serverTimingHeaderValue)
         .containsEntry(EXPOSE_HEADERS, SERVER_TIMING);
   }
 }