Fix startup perf issues

Author: manuel-alvarez-alvarez

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

@@ -46,6 +46,8 @@ 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 {
@@ -82,14 +84,12 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     }
     sco.createRemaining(config);
 
-    final double maxDownstreamRequestsRate =
-        config.getApiSecurityDownstreamRequestAnalysisSampleRate();
     GatewayBridge gatewayBridge =
         new GatewayBridge(
             gw,
             REPLACEABLE_EVENT_PRODUCER,
             () -> API_SECURITY_SAMPLER,
-            ApiSecurityDownstreamSampler.build(maxDownstreamRequestsRate),
+            () -> API_SECURITY_DOWNSTREAM_SAMPLER,
             APP_SEC_CONFIG_SERVICE.getTraceSegmentPostProcessors());
 
     loadModules(
@@ -216,6 +216,9 @@ 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 = ApiSecurityDownstreamSampler.build(rate);
       }
     }
   }

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

@@ -8,35 +8,25 @@ public interface ApiSecurityDownstreamSampler {
 
   boolean isSampled(AppSecRequestContext ctx, long requestId);
 
-  ApiSecurityDownstreamSampler INCLUDE_ALL =
-      new ApiSecurityDownstreamSampler() {
-        @Override
-        public boolean sampleHttpClientRequest(AppSecRequestContext ctx, long requestId) {
-          return true;
-        }
-
-        @Override
-        public boolean isSampled(AppSecRequestContext ctx, long requestId) {
-          return true;
-        }
-      };
-
-  ApiSecurityDownstreamSampler INCLUDE_NONE =
-      new ApiSecurityDownstreamSampler() {
-        @Override
-        public boolean sampleHttpClientRequest(AppSecRequestContext ctx, long requestId) {
-          return false;
-        }
-
-        @Override
-        public boolean isSampled(AppSecRequestContext ctx, long requestId) {
-          return false;
-        }
-      };
+  class NoOp implements ApiSecurityDownstreamSampler {
+
+    @Override
+    public boolean sampleHttpClientRequest(AppSecRequestContext ctx, long requestId) {
+      return false;
+    }
+
+    @Override
+    public boolean isSampled(AppSecRequestContext ctx, long requestId) {
+      return false;
+    }
+  }
 
   static ApiSecurityDownstreamSampler build(double rate) {
-    return rate <= 0D
-        ? INCLUDE_NONE
-        : (rate >= 1D ? INCLUDE_ALL : new ApiSecurityDownstreamSamplerImpl(rate));
+    if (rate < 0.0D) {
+      rate = 0.D;
+    } else if (rate > 1.0D) {
+      rate = 1.0D;
+    }
+    return new ApiSecurityDownstreamSamplerImpl(rate);
   }
 }

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

@@ -100,7 +100,7 @@ public class GatewayBridge {
   private final SubscriptionService subscriptionService;
   private final EventProducerService producerService;
   private final Supplier<ApiSecuritySampler> requestSamplerSupplier;
-  private final ApiSecurityDownstreamSampler downstreamSampler;
+  private final Supplier<ApiSecurityDownstreamSampler> downstreamSamplerSupplier;
   private final List<TraceSegmentPostProcessor> traceSegmentPostProcessors;
 
   // subscriber cache
@@ -129,12 +129,12 @@ public GatewayBridge(
       SubscriptionService subscriptionService,
       EventProducerService producerService,
       @Nonnull Supplier<ApiSecuritySampler> requestSamplerSupplier,
-      ApiSecurityDownstreamSampler downstreamSampler,
+      @Nonnull Supplier<ApiSecurityDownstreamSampler> downstreamSamplerSupplier,
       List<TraceSegmentPostProcessor> traceSegmentPostProcessors) {
     this.subscriptionService = subscriptionService;
     this.producerService = producerService;
     this.requestSamplerSupplier = requestSamplerSupplier;
-    this.downstreamSampler = downstreamSampler;
+    this.downstreamSamplerSupplier = downstreamSamplerSupplier;
     this.traceSegmentPostProcessors = traceSegmentPostProcessors;
   }
 
@@ -332,7 +332,8 @@ private Flow<Boolean> onHttpClientSampling(RequestContext ctx_, final long reque
       return new Flow.ResultFlow<>(null);
     }
     ctx.increaseHttpClientRequestCount();
-    return new Flow.ResultFlow<>(downstreamSampler.sampleHttpClientRequest(ctx, requestId));
+    final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
+    return new Flow.ResultFlow<>(sampler.sampleHttpClientRequest(ctx, requestId));
   }
 
   private Flow<Void> onHttpClientRequest(RequestContext ctx_, HttpClientRequest request) {
@@ -346,8 +347,9 @@ private Flow<Void> onHttpClientRequest(RequestContext ctx_, HttpClientRequest re
             .add(KnownAddresses.IO_NET_URL, request.getUrl())
             .add(KnownAddresses.IO_NET_REQUEST_METHOD, request.getMethod())
             .add(KnownAddresses.IO_NET_REQUEST_HEADERS, request.getHeaders());
-    ;
-    if (downstreamSampler.isSampled(ctx, request.getRequestId())) {
+
+    final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
+    if (sampler.isSampled(ctx, request.getRequestId())) {
       final Object body = parseHttpClientBody(ctx, request);
       if (body != null) {
         bundleBuilder.add(KnownAddresses.IO_NET_REQUEST_BODY, body);
@@ -387,7 +389,8 @@ private Flow<Void> onHttpClientResponse(RequestContext ctx_, HttpClientResponse
             .add(KnownAddresses.IO_NET_RESPONSE_STATUS, response.getStatus())
             .add(KnownAddresses.IO_NET_RESPONSE_HEADERS, response.getHeaders());
     // ignore the response if not sampled
-    if (downstreamSampler.isSampled(ctx, response.getRequestId())) {
+    final ApiSecurityDownstreamSampler sampler = downstreamSamplerSupplier.get();
+    if (sampler.isSampled(ctx, response.getRequestId())) {
       final Object body = parseHttpClientBody(ctx, response);
       if (body != null) {
         bundleBuilder.add(KnownAddresses.IO_NET_RESPONSE_BODY, body);

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

@@ -92,7 +92,7 @@ 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, () -> downstreamSampler, [pp])
 
   Supplier<Flow<AppSecRequestContext>> requestStartedCB
   BiFunction<RequestContext, AgentSpan, Flow<Void>> requestEndedCB