Fix after merge

Author: manuel-alvarez-alvarez

dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java

@@ -1,7 +1,5 @@
 package com.datadog.appsec;
 
-import com.datadog.appsec.api.security.ApiSecurityDownstreamSampler;
-import com.datadog.appsec.api.security.ApiSecurityDownstreamSamplerImpl;
 import com.datadog.appsec.api.security.ApiSecuritySampler;
 import com.datadog.appsec.api.security.ApiSecuritySamplerImpl;
 import com.datadog.appsec.api.security.AppSecSpanPostProcessor;
@@ -47,8 +45,6 @@ public class AppSecSystem {
   private static Runnable RESET_SUBSCRIPTION_SERVICE;
   private static final AtomicBoolean API_SECURITY_INITIALIZED = new AtomicBoolean(false);
   private static volatile ApiSecuritySampler API_SECURITY_SAMPLER = new ApiSecuritySampler.NoOp();
-  private static volatile ApiSecurityDownstreamSampler API_SECURITY_DOWNSTREAM_SAMPLER =
-      new ApiSecurityDownstreamSampler.NoOp();
 
   public static void start(SubscriptionService gw, SharedCommunicationObjects sco) {
     try {
@@ -90,7 +86,6 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
             gw,
             REPLACEABLE_EVENT_PRODUCER,
             () -> API_SECURITY_SAMPLER,
-            () -> API_SECURITY_DOWNSTREAM_SAMPLER,
             APP_SEC_CONFIG_SERVICE.getTraceSegmentPostProcessors());
 
     loadModules(
@@ -217,9 +212,6 @@ private static void maybeInitializeApiSecurity() {
         SpanPostProcessor.Holder.INSTANCE =
             new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
         API_SECURITY_SAMPLER = requestSampler;
-
-        final double rate = Config.get().getApiSecurityDownstreamRequestAnalysisSampleRate();
-        API_SECURITY_DOWNSTREAM_SAMPLER = new ApiSecurityDownstreamSamplerImpl(rate);
       }
     }
   }

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

@@ -10,6 +10,8 @@ public interface ApiSecurityDownstreamSampler {
 
   class NoOp implements ApiSecurityDownstreamSampler {
 
+    public static final NoOp INSTANCE = new NoOp();
+
     @Override
     public boolean sampleHttpClientRequest(AppSecRequestContext ctx, long requestId) {
       return false;

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

@@ -1,6 +1,7 @@
 package com.datadog.appsec.api.security;
 
 import com.datadog.appsec.gateway.AppSecRequestContext;
+import datadog.trace.api.Config;
 import java.util.concurrent.atomic.AtomicLong;
 
 public class ApiSecurityDownstreamSamplerImpl implements ApiSecurityDownstreamSampler {
@@ -9,12 +10,16 @@ public class ApiSecurityDownstreamSamplerImpl implements ApiSecurityDownstreamSa
   private final AtomicLong globalRequestCount;
   private final double threshold;
 
-  public ApiSecurityDownstreamSamplerImpl(double rate) {
+  public ApiSecurityDownstreamSamplerImpl() {
+    this(Config.get().getApiSecurityDownstreamRequestAnalysisSampleRate());
+  }
+
+  public ApiSecurityDownstreamSamplerImpl(final double rate) {
     threshold = samplingCutoff(rate < 0.0 ? 0 : (rate > 1.0 ? 1 : rate));
     globalRequestCount = new AtomicLong(0);
   }
 
-  private static double samplingCutoff(double rate) {
+  private static double samplingCutoff(final double rate) {
     final double max = Math.pow(2, 64) - 1;
     if (rate < 0.5) {
       return (long) (rate * max) + Long.MIN_VALUE;

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -11,6 +11,7 @@
 
 import com.datadog.appsec.AppSecSystem;
 import com.datadog.appsec.api.security.ApiSecurityDownstreamSampler;
+import com.datadog.appsec.api.security.ApiSecurityDownstreamSamplerImpl;
 import com.datadog.appsec.api.security.ApiSecuritySampler;
 import com.datadog.appsec.config.TraceSegmentPostProcessor;
 import com.datadog.appsec.event.EventProducerService;
@@ -100,8 +101,8 @@ public class GatewayBridge {
   private final SubscriptionService subscriptionService;
   private final EventProducerService producerService;
   private final Supplier<ApiSecuritySampler> requestSamplerSupplier;
-  private final Supplier<ApiSecurityDownstreamSampler> downstreamSamplerSupplier;
   private final List<TraceSegmentPostProcessor> traceSegmentPostProcessors;
+  private volatile ApiSecurityDownstreamSampler downstreamSampler;
 
   // subscriber cache
   private volatile DataSubscriberInfo initialReqDataSubInfo;
@@ -129,12 +130,10 @@ public GatewayBridge(
       SubscriptionService subscriptionService,
       EventProducerService producerService,
       @Nonnull Supplier<ApiSecuritySampler> requestSamplerSupplier,
-      @Nonnull Supplier<ApiSecurityDownstreamSampler> downstreamSamplerSupplier,
       List<TraceSegmentPostProcessor> traceSegmentPostProcessors) {
     this.subscriptionService = subscriptionService;
     this.producerService = producerService;
     this.requestSamplerSupplier = requestSamplerSupplier;
-    this.downstreamSamplerSupplier = downstreamSamplerSupplier;
     this.traceSegmentPostProcessors = traceSegmentPostProcessors;
   }
 
@@ -331,8 +330,7 @@ private Flow<Boolean> onHttpClientSampling(RequestContext ctx_, final long reque
     if (ctx == null) {
       return new Flow.ResultFlow<>(false);
     }
-    final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
-    return new Flow.ResultFlow<>(sampler.sampleHttpClientRequest(ctx, requestId));
+    return new Flow.ResultFlow<>(downstreamSampler().sampleHttpClientRequest(ctx, requestId));
   }
 
   private Flow<Void> onHttpClientRequest(RequestContext ctx_, HttpClientRequest request) {
@@ -347,8 +345,7 @@ private Flow<Void> onHttpClientRequest(RequestContext ctx_, HttpClientRequest re
             .add(KnownAddresses.IO_NET_REQUEST_METHOD, request.getMethod())
             .add(KnownAddresses.IO_NET_REQUEST_HEADERS, request.getHeaders());
 
-    final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
-    if (sampler.isSampled(ctx, request.getRequestId())) {
+    if (downstreamSampler().isSampled(ctx, request.getRequestId())) {
       final Object body = parseHttpClientBody(ctx, request);
       if (body != null) {
         bundleBuilder.add(KnownAddresses.IO_NET_REQUEST_BODY, body);
@@ -388,8 +385,7 @@ private Flow<Void> onHttpClientResponse(RequestContext ctx_, HttpClientResponse
             .add(KnownAddresses.IO_NET_RESPONSE_STATUS, Integer.toString(response.getStatus()))
             .add(KnownAddresses.IO_NET_RESPONSE_HEADERS, response.getHeaders());
     // ignore the response if not sampled
-    final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
-    if (sampler.isSampled(ctx, response.getRequestId())) {
+    if (downstreamSampler().isSampled(ctx, response.getRequestId())) {
       final Object body = parseHttpClientBody(ctx, response);
       if (body != null) {
         bundleBuilder.add(KnownAddresses.IO_NET_RESPONSE_BODY, body);
@@ -1209,6 +1205,14 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
     }
   }
 
+  private ApiSecurityDownstreamSampler downstreamSampler() {
+    if (downstreamSampler == null) {
+      // we don't care about concurrency too much
+      downstreamSampler = new ApiSecurityDownstreamSamplerImpl();
+    }
+    return downstreamSampler;
+  }
+
   private static Map<String, List<String>> parseQueryStringParams(
       String queryString, Charset uriEncoding) {
     if (queryString == null) {

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeIGRegistrationSpecification.groovy

@@ -10,7 +10,7 @@ class GatewayBridgeIGRegistrationSpecification extends DDSpecification {
   SubscriptionService ig = Mock()
   EventDispatcher eventDispatcher = Mock()
 
-  GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, null, [])
+  GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, [])
 
   void 'request_body_start and request_body_done are registered'() {
     given:

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -92,7 +92,9 @@ class GatewayBridgeSpecification extends DDSpecification {
   TraceSegmentPostProcessor pp = Mock()
   ApiSecuritySamplerImpl requestSampler = Mock(ApiSecuritySamplerImpl)
   ApiSecurityDownstreamSampler downstreamSampler = Mock(ApiSecurityDownstreamSampler)
-  GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, () -> requestSampler, () -> downstreamSampler, [pp])
+  GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, () -> requestSampler, [pp]).tap {
+    it.downstreamSampler = downstreamSampler
+  }
 
   Supplier<Flow<AppSecRequestContext>> requestStartedCB
   BiFunction<RequestContext, AgentSpan, Flow<Void>> requestEndedCB

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy

@@ -847,7 +847,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
   }
 
 
-  @IgnoreIf({ !instance.testAppSecAnalysis() })
+  @IgnoreIf({ !instance.testAppSecClientRequest() })
   void 'test appsec client request analysis'() {
     given:
     final url = server.address.resolve(endpoint)
@@ -992,7 +992,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
     true
   }
 
-  boolean testAppSecAnalysis() {
+  boolean testAppSecClientRequest() {
     false
   }