Add capped sampling rate increases

Author: raphaelgavache

dd-trace-core/src/main/java/datadog/trace/common/sampling/RateByServiceTraceSampler.java

@@ -10,6 +10,7 @@
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.function.Function;
+import java.util.function.LongSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -24,8 +25,11 @@ public class RateByServiceTraceSampler implements Sampler, PrioritySampler, Remo
   public static final String SAMPLING_AGENT_RATE = "_dd.agent_psr";
 
   private static final double DEFAULT_RATE = 1.0;
+  static final long RAMP_UP_INTERVAL_NANOS = 1_000_000_000L;
 
   private volatile RateSamplersByEnvAndService serviceRates = new RateSamplersByEnvAndService();
+  private long lastCappedNanos;
+  LongSupplier nanoTimeSupplier = System::nanoTime;
 
   @Override
   public <T extends CoreSpan<T>> boolean sample(final T span) {
@@ -62,6 +66,16 @@ private <T extends CoreSpan<T>> String getSpanEnv(final T span) {
     return span.getTag("env", "");
   }
 
+  static double cappedRate(double oldRate, double newRate, boolean canIncrease) {
+    if (newRate <= oldRate || oldRate == 0) {
+      return newRate;
+    }
+    if (!canIncrease) {
+      return oldRate;
+    }
+    return Math.min(oldRate * 2, newRate);
+  }
+
   @Override
   public void onResponse(
       final String endpoint, final Map<String, Map<String, Number>> responseJson) {
@@ -72,6 +86,13 @@ public void onResponse(
     }
 
     log.debug("Update service sampler rates: {} -> {}", endpoint, responseJson);
+
+    final RateSamplersByEnvAndService currentSnapshot = serviceRates;
+    final long now = nanoTimeSupplier.getAsLong();
+    final boolean canIncrease =
+        lastCappedNanos == 0 || (now - lastCappedNanos) >= RAMP_UP_INTERVAL_NANOS;
+    boolean anyCapped = false;
+
     final TreeMap<String, TreeMap<String, RateSampler>> updatedEnvServiceRates =
         new TreeMap<>(String::compareToIgnoreCase);
 
@@ -84,17 +105,33 @@ public void onResponse(
 
       EnvAndService envAndService = EnvAndService.fromString(entry.getKey());
       if (envAndService.isFallback()) {
-        fallbackSampler = RateByServiceTraceSampler.createRateSampler(rate);
+        double oldRate = currentSnapshot.getFallbackSampler().getSampleRate();
+        double effective = cappedRate(oldRate, rate, canIncrease);
+        if (effective != rate) {
+          anyCapped = true;
+        }
+        fallbackSampler = RateByServiceTraceSampler.createRateSampler(effective);
       } else {
+        double oldRate =
+            currentSnapshot
+                .getSampler(envAndService.lowerEnv, envAndService.lowerService)
+                .getSampleRate();
+        double effective = cappedRate(oldRate, rate, canIncrease);
+        if (effective != rate) {
+          anyCapped = true;
+        }
         Map<String, RateSampler> serviceRates =
             updatedEnvServiceRates.computeIfAbsent(
                 envAndService.lowerEnv, env -> new TreeMap<>(String::compareToIgnoreCase));
 
         serviceRates.computeIfAbsent(
             envAndService.lowerService,
-            service -> RateByServiceTraceSampler.createRateSampler(rate));
+            service -> RateByServiceTraceSampler.createRateSampler(effective));
       }
     }
+    if (anyCapped) {
+      lastCappedNanos = now;
+    }
     serviceRates = new RateSamplersByEnvAndService(updatedEnvServiceRates, fallbackSampler);
   }
 
@@ -128,6 +165,10 @@ private static final class RateSamplersByEnvAndService {
       this.fallbackSampler = fallbackSampler;
     }
 
+    RateSampler getFallbackSampler() {
+      return fallbackSampler;
+    }
+
     // used in tests only
     RateSampler getSampler(EnvAndService envAndService) {
       return getSampler(envAndService.lowerEnv, envAndService.lowerService);

dd-trace-core/src/test/groovy/datadog/trace/common/sampling/RateByServiceTraceSamplerTest.groovy

@@ -222,6 +222,135 @@ class RateByServiceTraceSamplerTest extends DDCoreSpecification {
     'manual.keep' | true     | PrioritySampling.USER_KEEP
   }
 
+  def "cappedRate returns new rate when decreasing"() {
+    expect:
+    RateByServiceTraceSampler.cappedRate(0.8, 0.4, true) == 0.4
+    RateByServiceTraceSampler.cappedRate(0.8, 0.4, false) == 0.4
+    RateByServiceTraceSampler.cappedRate(0.5, 0.5, true) == 0.5
+    RateByServiceTraceSampler.cappedRate(0.5, 0.5, false) == 0.5
+  }
+
+  def "cappedRate returns new rate when old rate is zero"() {
+    expect:
+    RateByServiceTraceSampler.cappedRate(0.0, 0.5, true) == 0.5
+    RateByServiceTraceSampler.cappedRate(0.0, 0.5, false) == 0.5
+  }
+
+  def "cappedRate caps increase to 2x when canIncrease"() {
+    expect:
+    RateByServiceTraceSampler.cappedRate(0.1, 1.0, true) == 0.2
+    RateByServiceTraceSampler.cappedRate(0.2, 1.0, true) == 0.4
+    RateByServiceTraceSampler.cappedRate(0.4, 1.0, true) == 0.8
+    RateByServiceTraceSampler.cappedRate(0.4, 0.5, true) == 0.5
+  }
+
+  def "cappedRate holds old rate when canIncrease is false"() {
+    expect:
+    RateByServiceTraceSampler.cappedRate(0.1, 1.0, false) == 0.1
+    RateByServiceTraceSampler.cappedRate(0.2, 0.8, false) == 0.2
+  }
+
+  def "ramp-up caps rate increases at 2x per interval"() {
+    setup:
+    RateByServiceTraceSampler serviceSampler = new RateByServiceTraceSampler()
+    long currentTime = 1_000_000_000L
+    serviceSampler.nanoTimeSupplier = { -> currentTime }
+    def tolerance = 0.01
+
+    // Set initial rate to 0.1
+    String response = '{"rate_by_service": {"service:foo,env:bar":0.1, "service:,env:":0.1}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    expect:
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.1) < tolerance
+
+    when: "agent restart sends rate 1.0, first interval"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    response = '{"rate_by_service": {"service:foo,env:bar":1.0, "service:,env:":1.0}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate is capped at 2x = 0.2"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.2) < tolerance
+    Math.abs(serviceSampler.serviceRates.getFallbackSampler().sampleRate - 0.2) < tolerance
+
+    when: "second interval"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate doubles to 0.4"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.4) < tolerance
+    Math.abs(serviceSampler.serviceRates.getFallbackSampler().sampleRate - 0.4) < tolerance
+
+    when: "third interval"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate doubles to 0.8"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.8) < tolerance
+    Math.abs(serviceSampler.serviceRates.getFallbackSampler().sampleRate - 0.8) < tolerance
+
+    when: "fourth interval"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate reaches target 1.0 (2x=1.6 > 1.0)"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 1.0) < tolerance
+    Math.abs(serviceSampler.serviceRates.getFallbackSampler().sampleRate - 1.0) < tolerance
+  }
+
+  def "ramp-down applies immediately"() {
+    setup:
+    RateByServiceTraceSampler serviceSampler = new RateByServiceTraceSampler()
+    long currentTime = 1_000_000_000L
+    serviceSampler.nanoTimeSupplier = { -> currentTime }
+    def tolerance = 0.01
+
+    // Set initial rate to 0.8
+    String response = '{"rate_by_service": {"service:foo,env:bar":0.8, "service:,env:":0.8}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    when: "rate decreases to 0.2"
+    response = '{"rate_by_service": {"service:foo,env:bar":0.2, "service:,env:":0.2}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "decrease is applied immediately"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.2) < tolerance
+    Math.abs(serviceSampler.serviceRates.getFallbackSampler().sampleRate - 0.2) < tolerance
+  }
+
+  def "rate increase blocked during cooldown"() {
+    setup:
+    RateByServiceTraceSampler serviceSampler = new RateByServiceTraceSampler()
+    long currentTime = 1_000_000_000L
+    serviceSampler.nanoTimeSupplier = { -> currentTime }
+    def tolerance = 0.01
+
+    // Set initial rate to 0.1
+    String response = '{"rate_by_service": {"service:foo,env:bar":0.1}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    when: "rate jumps, first capped increase"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    response = '{"rate_by_service": {"service:foo,env:bar":1.0}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "capped to 0.2"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.2) < tolerance
+
+    when: "try again immediately (within cooldown)"
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate stays at 0.2 because cooldown hasn't elapsed"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.2) < tolerance
+
+    when: "after cooldown elapsed"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate doubles to 0.4"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.4) < tolerance
+  }
+
   def "not setting forced tracing via tag or setting it wrong value not causing exception"() {
     setup:
     def sampler = new RateByServiceTraceSampler()