fix

Author: raphaelgavache

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

@@ -25,6 +25,7 @@ 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;
+  private static final double MAX_RATE_INCREASE_FACTOR = 2.0;
   static final long RAMP_UP_INTERVAL_NANOS = 1_000_000_000L;
 
   private volatile RateSamplersByEnvAndService serviceRates = new RateSamplersByEnvAndService();
@@ -66,14 +67,12 @@ 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);
+  static boolean shouldCap(double oldRate, double newRate) {
+    return oldRate != 0 && newRate > oldRate * MAX_RATE_INCREASE_FACTOR;
+  }
+
+  static double cappedRate(double oldRate) {
+    return oldRate * MAX_RATE_INCREASE_FACTOR;
   }
 
   @Override
@@ -106,30 +105,35 @@ public void onResponse(
       EnvAndService envAndService = EnvAndService.fromString(entry.getKey());
       if (envAndService.isFallback()) {
         double oldRate = currentSnapshot.getFallbackSampler().getSampleRate();
-        double effective = cappedRate(oldRate, rate, canIncrease);
-        if (effective != rate) {
+        if (canIncrease && shouldCap(oldRate, rate)) {
+          rate = cappedRate(oldRate);
           anyCapped = true;
+        } else if (!canIncrease && shouldCap(oldRate, rate)) {
+          rate = oldRate;
         }
-        fallbackSampler = RateByServiceTraceSampler.createRateSampler(effective);
+        fallbackSampler = RateByServiceTraceSampler.createRateSampler(rate);
       } else {
         double oldRate =
             currentSnapshot
                 .getSampler(envAndService.lowerEnv, envAndService.lowerService)
                 .getSampleRate();
-        double effective = cappedRate(oldRate, rate, canIncrease);
-        if (effective != rate) {
+        if (canIncrease && shouldCap(oldRate, rate)) {
+          rate = cappedRate(oldRate);
           anyCapped = true;
+        } else if (!canIncrease && shouldCap(oldRate, rate)) {
+          rate = oldRate;
         }
+        final double effectiveRate = rate;
         Map<String, RateSampler> serviceRates =
             updatedEnvServiceRates.computeIfAbsent(
                 envAndService.lowerEnv, env -> new TreeMap<>(String::compareToIgnoreCase));
 
         serviceRates.computeIfAbsent(
             envAndService.lowerService,
-            service -> RateByServiceTraceSampler.createRateSampler(effective));
+            service -> RateByServiceTraceSampler.createRateSampler(effectiveRate));
       }
     }
-    if (anyCapped) {
+    if (canIncrease && anyCapped) {
       lastCappedNanos = now;
     }
     serviceRates = new RateSamplersByEnvAndService(updatedEnvServiceRates, fallbackSampler);

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

@@ -222,32 +222,31 @@ class RateByServiceTraceSamplerTest extends DDCoreSpecification {
     'manual.keep' | true     | PrioritySampling.USER_KEEP
   }
 
-  def "cappedRate returns new rate when decreasing"() {
+  def "shouldCap returns false when rate decreases or stays same"() {
     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
+    !RateByServiceTraceSampler.shouldCap(0.8, 0.4)
+    !RateByServiceTraceSampler.shouldCap(0.5, 0.5)
+    !RateByServiceTraceSampler.shouldCap(0.5, 1.0) // 1.0 <= 0.5 * 2, no cap needed
   }
 
-  def "cappedRate returns new rate when old rate is zero"() {
+  def "shouldCap returns false when old rate is zero"() {
     expect:
-    RateByServiceTraceSampler.cappedRate(0.0, 0.5, true) == 0.5
-    RateByServiceTraceSampler.cappedRate(0.0, 0.5, false) == 0.5
+    !RateByServiceTraceSampler.shouldCap(0.0, 0.5)
+    !RateByServiceTraceSampler.shouldCap(0.0, 1.0)
   }
 
-  def "cappedRate caps increase to 2x when canIncrease"() {
+  def "shouldCap returns true when new rate exceeds 2x old rate"() {
     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
+    RateByServiceTraceSampler.shouldCap(0.1, 1.0)
+    RateByServiceTraceSampler.shouldCap(0.2, 0.8)
+    RateByServiceTraceSampler.shouldCap(0.1, 0.3)
   }
 
-  def "cappedRate holds old rate when canIncrease is false"() {
+  def "cappedRate returns 2x old rate"() {
     expect:
-    RateByServiceTraceSampler.cappedRate(0.1, 1.0, false) == 0.1
-    RateByServiceTraceSampler.cappedRate(0.2, 0.8, false) == 0.2
+    RateByServiceTraceSampler.cappedRate(0.1) == 0.2
+    RateByServiceTraceSampler.cappedRate(0.2) == 0.4
+    RateByServiceTraceSampler.cappedRate(0.4) == 0.8
   }
 
   def "ramp-up caps rate increases at 2x per interval"() {
@@ -351,6 +350,43 @@ class RateByServiceTraceSamplerTest extends DDCoreSpecification {
     Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.4) < tolerance
   }
 
+  def "cooldown not reset by blocked increase"() {
+    setup:
+    RateByServiceTraceSampler serviceSampler = new RateByServiceTraceSampler()
+    long currentTime = 1_000_000_000L
+    serviceSampler.nanoTimeSupplier = { -> currentTime }
+    def tolerance = 0.01
+
+    // Set initial low rate
+    String response = '{"rate_by_service": {"service:foo,env:bar":0.01}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    expect:
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.01) < tolerance
+
+    when: "wait for cooldown, apply increase: 0.01 -> 0.02"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS
+    response = '{"rate_by_service": {"service:foo,env:bar":1.0}}'
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate is capped at 2x = 0.02"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.02) < tolerance
+
+    when: "before cooldown elapses, send another increase - rate should be held and lastCapped NOT reset"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS / 2
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate stays at 0.02 (cooldown)"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.02) < tolerance
+
+    when: "wait remaining half of cooldown from the original cap - should allow next ramp-up"
+    currentTime += RateByServiceTraceSampler.RAMP_UP_INTERVAL_NANOS / 2
+    serviceSampler.onResponse("traces", serializer.fromJson(response))
+
+    then: "rate doubles to 0.04 because lastCapped was NOT reset by the blocked increase"
+    Math.abs(serviceSampler.serviceRates.getSampler("bar", "foo").sampleRate - 0.04) < tolerance
+  }
+
   def "not setting forced tracing via tag or setting it wrong value not causing exception"() {
     setup:
     def sampler = new RateByServiceTraceSampler()