Fix API Security sampling in standalone mode (#10165)

What Does This Do
Fix the API security sampling in standalone mode (APM tracing disabled).
Since we decided that, at the end of the request, if the pre-sample chooses to retain the trace we will add the asm.keep and propagation tags for ASM in standalone mode, we also need to keep that same decision in the sample method. This avoids inconsistencies in concurrent threads under high load scenarios where the post-processor might be delayed.

Added a bit of extra complexity to get more detailed logs through logSamplingDecision. This should make it easier to troubleshoot future sampling issues

Motivation
API Security standalone system tests were failing intermittently in CI with _sampling_priority_v1 set to 2, in traces that should be not retained due to API Security sampling

Related with APPSEC-57815

Additional Notes
tested in https://github.com/DataDog/system-tests/actions/runs/20127087811/job/57783265656

Author: jandro996

dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySamplerImpl.java

@@ -4,6 +4,8 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.time.SystemTimeSource;
 import datadog.trace.api.time.TimeSource;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.util.Deque;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
@@ -73,8 +75,15 @@ public boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx) {
       return false;
     }
     if (counter.tryAcquire()) {
-      log.debug("API security sampling is required for this request (presampled)");
       ctx.setKeepOpenForApiSecurityPostProcessing(true);
+      if (!Config.get().isApmTracingEnabled()) {
+        boolean sampled = updateApiAccessIfExpired(hash);
+        if (sampled) {
+          logSamplingDecision("preSampleRequest", hash);
+        }
+        return sampled;
+      }
+      logSamplingDecision("preSampleRequest", hash);
       return true;
     }
     return false;
@@ -91,7 +100,11 @@ public boolean sampleRequest(AppSecRequestContext ctx) {
       // This should never happen, it should have been short-circuited before.
       return false;
     }
-    return updateApiAccessIfExpired(hash);
+    boolean sampled = Config.get().isApmTracingEnabled() ? updateApiAccessIfExpired(hash) : true;
+    if (sampled) {
+      logSamplingDecision("sampleRequest", hash);
+    }
+    return sampled;
   }
 
   @Override
@@ -168,4 +181,23 @@ private long computeApiHash(final String route, final String method, final int s
     result = 31 * result + statusCode;
     return result;
   }
+
+  private void logSamplingDecision(final String method, final long hash) {
+    if (!log.isDebugEnabled()) {
+      return;
+    }
+    AgentSpan activeSpan = AgentTracer.get().activeSpan();
+
+    if (activeSpan != null) {
+      log.debug(
+          "API security sampling decision in {}: hash={}, traceId={}, spanId={}",
+          method,
+          hash,
+          activeSpan.getTraceId(),
+          activeSpan.getSpanId());
+    } else {
+      log.debug(
+          "API security sampling decision in {}: hash={}, traceId=null, spanId=null", method, hash);
+    }
+  }
 }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecuritySamplerTest.groovy

@@ -209,6 +209,161 @@ class ApiSecuritySamplerTest extends DDSpecification {
     }
   }
 
+  void 'preSampleRequest with tracing disabled updates access map immediately'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'first request is presampled with tracing disabled'
+    def preSampled = sampler.preSampleRequest(ctx)
+
+    then: 'request is sampled and access map is updated immediately'
+    preSampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.containsKey(ctx.getApiSecurityEndpointHash())
+
+    when: 'second request for same endpoint is attempted'
+    def ctx2 = createContext('route1', 'GET', 200)
+    sampler.releaseOne()
+    def preSampled2 = sampler.preSampleRequest(ctx2)
+
+    then: 'second request is not sampled because endpoint was already updated in first preSampleRequest'
+    !preSampled2
+  }
+
+  void 'sampleRequest with tracing disabled returns true without updating access map'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    ctx.setApiSecurityEndpointHash(42L)
+    ctx.setKeepOpenForApiSecurityPostProcessing(true)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'sampleRequest is called with tracing disabled'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'request is sampled without updating access map'
+    sampled
+    sampler.accessMap.size() == 0
+  }
+
+  void 'preSampleRequest with tracing enabled does not update access map immediately'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'true')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'request is presampled with tracing enabled'
+    def preSampled = sampler.preSampleRequest(ctx)
+
+    then: 'request is sampled but access map is NOT updated yet'
+    preSampled
+    sampler.accessMap.size() == 0
+
+    when: 'sampleRequest is called to finalize sampling'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'access map is updated in sampleRequest'
+    sampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.containsKey(ctx.getApiSecurityEndpointHash())
+  }
+
+  void 'sampleRequest with tracing enabled updates access map'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'true')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    ctx.setApiSecurityEndpointHash(42L)
+    ctx.setKeepOpenForApiSecurityPostProcessing(true)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'sampleRequest is called with tracing enabled'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'request is sampled and access map is updated'
+    sampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.containsKey(42L)
+
+    when: 'second request for same endpoint is made'
+    def sampled2 = sampler.sampleRequest(ctx)
+
+    then: 'second request is not sampled'
+    !sampled2
+  }
+
+  void 'concurrent requests with tracing disabled do not see expired state'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx1 = createContext('route1', 'GET', 200)
+    def ctx2 = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'first request is presampled'
+    def preSampled1 = sampler.preSampleRequest(ctx1)
+
+    then: 'first request is sampled and access map is updated immediately'
+    preSampled1
+    ctx1.getApiSecurityEndpointHash() != null
+    sampler.accessMap.size() == 1
+
+    when: 'concurrent second request tries to presample same endpoint'
+    sampler.releaseOne()
+    def preSampled2 = sampler.preSampleRequest(ctx2)
+
+    then: 'second request is not sampled because endpoint is already in access map'
+    !preSampled2
+  }
+
+  void 'full flow with tracing disabled updates map only in preSampleRequest'() {
+    given:
+    injectSysConfig('dd.apm.tracing.enabled', 'false')
+    rebuildConfig()
+    def ctx = createContext('route1', 'GET', 200)
+    final timeSource = new ControllableTimeSource()
+    timeSource.set(0)
+    final long expirationTimeInMs = 10_000
+    def sampler = new ApiSecuritySamplerImpl(10, expirationTimeInMs, timeSource)
+
+    when: 'request goes through full sampling flow with tracing disabled'
+    def preSampled = sampler.preSampleRequest(ctx)
+
+    then: 'preSampleRequest returns true and updates access map'
+    preSampled
+    sampler.accessMap.size() == 1
+    def hash = ctx.getApiSecurityEndpointHash()
+    sampler.accessMap.containsKey(hash)
+
+    when: 'sampleRequest is called'
+    def sampled = sampler.sampleRequest(ctx)
+
+    then: 'sampleRequest returns true without modifying access map'
+    sampled
+    sampler.accessMap.size() == 1
+    sampler.accessMap.get(hash) == 0L // Still has the value from preSampleRequest
+  }
+
   private static AppSecRequestContext createContext(final String route, final String method, int statusCode) {
     final AppSecRequestContext ctx = new AppSecRequestContext()
     ctx.setRoute(route)