Use Clock in ASM Standalone Sampler (#7270)

What Does This Do
Use clock.millis() instead os System.currentTimeMillis() in the ASM Standalone Sampler
fix the rate limiter to 1 minute as is defined in the RFC
Motivation
Mitigate flakiness and facilitate test implementation

APPSEC-10459

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

@@ -2,25 +2,29 @@
 
 import datadog.trace.api.sampling.SamplingMechanism;
 import datadog.trace.core.CoreSpan;
+import java.time.Clock;
 import java.util.concurrent.atomic.AtomicLong;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /**
- * This class is designed to sample traces based on a fixed time rate in milliseconds. However, it's
- * important to note that the precision of this class is not in the millisecond range due to the use
- * of System.currentTimeMillis().
+ * This class is designed to only allow 1 APM trace per minute as standalone ASM is only interested
+ * in the traces containing ASM events. But the service catalog and the billing need a continuous
+ * ingestion of at least at 1 trace per minute to consider a service as being live and billable. In
+ * the absence of ASM events, no APM traces must be sent, so we need to let some regular APM traces
+ * go through, even in the absence of ASM events.
  */
 public class AsmStandaloneSampler implements Sampler, PrioritySampler {
 
   private static final Logger log = LoggerFactory.getLogger(AsmStandaloneSampler.class);
+  private static final int RATE_IN_MILLISECONDS = 60000; // 1 minute
 
   private final AtomicLong lastSampleTime;
-  private final int rateInMilliseconds;
+  private final Clock clock;
 
-  public AsmStandaloneSampler(int rateInMilliseconds) {
-    this.rateInMilliseconds = rateInMilliseconds;
-    this.lastSampleTime = new AtomicLong(System.currentTimeMillis() - rateInMilliseconds);
+  public AsmStandaloneSampler(final Clock clock) {
+    this.clock = clock;
+    this.lastSampleTime = new AtomicLong(clock.millis() - RATE_IN_MILLISECONDS);
   }
 
   @Override
@@ -43,9 +47,9 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
   }
 
   private boolean shouldSample() {
-    long now = System.currentTimeMillis();
+    long now = clock.millis();
     return lastSampleTime.updateAndGet(
-            lastTime -> now - lastTime >= rateInMilliseconds ? now : lastTime)
+            lastTime -> now - lastTime >= RATE_IN_MILLISECONDS ? now : lastTime)
         == now;
   }
 }

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

@@ -10,6 +10,7 @@
 import datadog.trace.api.sampling.SamplingMechanism;
 import datadog.trace.api.sampling.SamplingRule;
 import datadog.trace.core.CoreSpan;
+import java.time.Clock;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -20,15 +21,6 @@
 /** Main interface to sample a collection of traces. */
 public interface Sampler {
 
-  /**
-   * To only allow 1 APM trace per minute as standalone ASM is only interested in the traces
-   * containing ASM events. But the service catalog and the billing need a continuous ingestion of
-   * at least at 1 trace per minute to consider a service as being live and billable. In the absence
-   * of ASM events, no APM traces must be sent, so we need to let some regular APM traces go
-   * through, even in the absence of ASM events.
-   */
-  int ASM_STANDALONE_BILLING_SAMPLER_RATE = 60000; // 1 minute
-
   /**
    * Sample a collection of traces based on the parent span
    *
@@ -45,7 +37,7 @@ public static Sampler forConfig(final Config config, final TraceConfig traceConf
       if (config != null) {
         if (config.isAppSecStandaloneEnabled()) {
           log.debug("APM is disabled. Only 1 trace per minute will be sent.");
-          return new AsmStandaloneSampler(ASM_STANDALONE_BILLING_SAMPLER_RATE);
+          return new AsmStandaloneSampler(Clock.systemUTC());
         }
         final Map<String, String> serviceRules = config.getTraceSamplingServiceRules();
         final Map<String, String> operationRules = config.getTraceSamplingOperationRules();

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

@@ -4,41 +4,52 @@ import datadog.trace.common.writer.ListWriter
 import datadog.trace.core.test.DDCoreSpecification
 import datadog.trace.api.sampling.PrioritySampling
 
+import java.time.Clock
+import java.util.concurrent.atomic.AtomicLong
+
 class AsmStandaloneSamplerTest extends DDCoreSpecification{
 
   def writer = new ListWriter()
 
   void "test setSamplingPriority"(){
     setup:
-    final rate = 2000 // 1 trace each 2 seconds
-    def sampler = new AsmStandaloneSampler(rate)
+    def current = new AtomicLong(System.currentTimeMillis())
+    final Clock clock = Mock(Clock) {
+      millis() >> {
+        current.get()
+      }
+    }
+    def sampler = new AsmStandaloneSampler(clock)
     def tracer = tracerBuilder().writer(writer).sampler(sampler).build()
 
     when:
     def span1 = tracer.buildSpan("test").start()
     sampler.setSamplingPriority(span1)
 
     then:
+    1 * clock.millis() >> {
+      current.updateAndGet(value -> value + 1000)
+    } // increment in one second
     span1.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
 
     when:
     def span2 = tracer.buildSpan("test2").start()
     sampler.setSamplingPriority(span2)
 
     then:
+    1 * clock.millis() >> { current.updateAndGet(value -> value + 1000) } // increment in one second
     span2.getSamplingPriority() == PrioritySampling.SAMPLER_DROP
 
     when:
     def span3 = tracer.buildSpan("test3").start()
-
-    then: "we wait for one second"
-    Thread.sleep(rate)
     sampler.setSamplingPriority(span3)
 
-    and:
+    then: "Mock one minute later"
+    clock.millis() >> { current.updateAndGet(value -> value + 60000) } // increment in one minute
     span3.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
 
     cleanup:
     tracer.close()
   }
+
 }

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

@@ -15,6 +15,5 @@ class SamplerTest extends DDSpecification{
 
     then:
     sampler instanceof AsmStandaloneSampler
-    ((AsmStandaloneSampler)sampler).rateInMilliseconds == Sampler.ASM_STANDALONE_BILLING_SAMPLER_RATE
   }
 }