feat: configure custom sampler

grafana · Dec 2, 2024 · 0f895d1 · 0f895d1
1 parent 0eb9fa8
commit 0f895d1
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 16 deletions.
diff --git a/custom/src/main/java/com/grafana/extensions/GrafanaAutoConfigCustomizerProvider.java b/custom/src/main/java/com/grafana/extensions/GrafanaAutoConfigCustomizerProvider.java
@@ -8,8 +8,10 @@
 import com.grafana.extensions.filter.MetricsCustomizer;
 import com.grafana.extensions.instrumentations.TestedInstrumentationsCustomizer;
 import com.grafana.extensions.resources.ResourceCustomizer;
+import com.grafana.extensions.sampling.MovingAverageThresholdSampler;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
+
 import java.util.HashMap;
 import java.util.Map;
 
@@ -18,10 +20,11 @@ public class GrafanaAutoConfigCustomizerProvider implements AutoConfigurationCus
   @Override
   public void customize(AutoConfigurationCustomizer autoConfiguration) {
     autoConfiguration
-        .addPropertiesSupplier(GrafanaAutoConfigCustomizerProvider::getDefaultProperties)
-        .addPropertiesCustomizer(TestedInstrumentationsCustomizer::customizeProperties)
-        .addMeterProviderCustomizer(MetricsCustomizer::configure)
-        .addResourceCustomizer(ResourceCustomizer::truncate);
+      .addPropertiesSupplier(GrafanaAutoConfigCustomizerProvider::getDefaultProperties)
+      .addPropertiesCustomizer(TestedInstrumentationsCustomizer::customizeProperties)
+      .addMeterProviderCustomizer(MetricsCustomizer::configure)
+      .addSamplerCustomizer(MovingAverageThresholdSampler::configure)
+      .addResourceCustomizer(ResourceCustomizer::truncate);
   }
 
   private static Map<String, String> getDefaultProperties() {

diff --git a/custom/src/main/java/com/grafana/extensions/sampling/MovingAverageThresholdSampler.java b/custom/src/main/java/com/grafana/extensions/sampling/MovingAverageThresholdSampler.java
@@ -9,46 +9,54 @@
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class MovingAverageThresholdSampler implements Sampler {
   private final Map<String, MovingAverage> movingAvgs = new ConcurrentHashMap<>();
-  private final double threshHoldVal;
+  private final double thresholdVal;
   private final int windowSize;
 
-  public MovingAverageThresholdSampler(double threshHoldVal, int windowSize) {
-    this.threshHoldVal = threshHoldVal;
+  protected MovingAverageThresholdSampler(double thresholdVal, int windowSize) {
+    this.thresholdVal = thresholdVal;
     this.windowSize = windowSize;
   }
 
+  public static Sampler configure(Sampler sampler, ConfigProperties properties) {
+    double threshold = properties.getDouble("threshold", 1.5);
+    int windowSize = properties.getInt("window", 20);
+    return new MovingAverageThresholdSampler(threshold, windowSize);
+  }
+
   @Override
   public SamplingResult shouldSample(
-      Context context,
-      String traceId,
-      String spanName,
-      SpanKind spanKind,
-      Attributes attributes,
-      List<LinkData> list) {
+    Context context,
+    String traceId,
+    String spanName,
+    SpanKind spanKind,
+    Attributes attributes,
+    List<LinkData> list) {
     Long duration = attributes.get(AttributeKey.longKey("duration"));
     if (duration == null) {
       return SamplingResult.recordAndSample();
     }
 
     MovingAverage currMovingAvg =
-        movingAvgs.computeIfAbsent(spanName, ma -> new MovingAverage(windowSize));
+      movingAvgs.computeIfAbsent(spanName, ma -> new MovingAverage(windowSize));
     // record until window is full
     if (currMovingAvg.getCount() < windowSize) {
       currMovingAvg.addAndCalcAverage(duration);
       return SamplingResult.recordAndSample();
     }
 
     double avg = currMovingAvg.addAndCalcAverage(duration);
-    if (duration < avg * threshHoldVal) {
+    if (duration < avg * thresholdVal) {
       return SamplingResult.drop();
     }
 
@@ -57,6 +65,7 @@ public SamplingResult shouldSample(
 
   @Override
   public String getDescription() {
-    return "";
+    return "MovingAverageThresholdSampler";
   }
+
 }