Added post-processing limit

Author: ValentinZakharov

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadContext.java

@@ -2,7 +2,7 @@
 
 import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED;
 
-import com.datadog.iast.util.NonBlockingSemaphore;
+import datadog.trace.util.NonBlockingSemaphore;
 
 public class OverheadContext {
 

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java

@@ -4,7 +4,6 @@
 
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.IastSystem;
-import com.datadog.iast.util.NonBlockingSemaphore;
 import datadog.trace.api.Config;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
@@ -13,6 +12,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.util.AgentTaskScheduler;
+import datadog.trace.util.NonBlockingSemaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import javax.annotation.Nullable;

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

@@ -38,10 +38,11 @@ public ApiAccessTracker(int capacity, long expirationTimeInMs) {
   }
 
   /**
-   * Updates the API access log with the given route, method, and status code. If the record exists
-   * and is outdated, it is updated by moving to the end of the list. If the record does not exist,
-   * a new record is added. If the capacity limit is reached, the oldest record is removed. Returns
-   * true if the record was updated or added, false otherwise.
+   * Updates the API access log with the given route, method, and status code. If the record already
+   * exists and is outdated, it is updated by moving to the end of the list. If the record does not
+   * exist, a new record is added. If the capacity limit is reached, the oldest record is removed.
+   * This method should not be called concurrently by multiple threads, due absence of additional
+   * synchronization for updating data structures is not required.
    *
    * @param route The route of the API endpoint request
    * @param method The method of the API request
@@ -66,7 +67,7 @@ public boolean updateApiAccessIfExpired(String route, String method, int statusC
       isNewOrUpdated = true;
 
       // Remove the oldest hash if capacity is reached
-      while (apiAccessQueue.size() > this.capacity) {
+      while (apiAccessMap.size() > this.capacity) {
         Long oldestHash = apiAccessQueue.pollFirst();
         if (oldestHash != null) {
           apiAccessMap.remove(oldestHash);

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

@@ -41,6 +41,7 @@
 import datadog.trace.api.telemetry.WafMetricCollector;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
+import datadog.trace.util.NonBlockingSemaphore;
 import datadog.trace.util.stacktrace.StackTraceEvent;
 import datadog.trace.util.stacktrace.StackUtils;
 import java.net.URI;
@@ -94,6 +95,10 @@ public class GatewayBridge {
 
   private static final String METASTRUCT_EXPLOIT = "exploit";
 
+  private final int MAX_POST_PROCESSING_TASKS = 16;
+  private final NonBlockingSemaphore postProcessingCounter =
+      NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS);
+
   private final SubscriptionService subscriptionService;
   private final EventProducerService producerService;
   private final ApiSecurityRequestSampler requestSampler;
@@ -833,7 +838,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
     if (route instanceof String) {
       ctx.setRoute((String) route);
     }
-    if (requestSampler.preSampleRequest(ctx)) {
+    if (requestSampler.preSampleRequest(ctx) && postProcessingCounter.acquire()) {
       // The request is pre-sampled - we need to post-process it
       spanInfo.setRequiresPostProcessing(true);
     }
@@ -905,6 +910,8 @@ private void onPostProcessing(RequestContext ctx_) {
 
     maybeExtractSchemas(ctx);
     ctx.close();
+    // Decrease the counter to allow the next request to be post-processed
+    postProcessingCounter.release();
   }
 
   public void stop() {

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java renamed to internal-api/src/main/java/datadog/trace/util/NonBlockingSemaphore.java

@@ -1,4 +1,4 @@
-package com.datadog.iast.util;
+package datadog.trace.util;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy renamed to internal-api/src/test/groovy/datadog/trace/util/NonBlockingSemaphoreTest.groovy

@@ -1,4 +1,4 @@
-package com.datadog.iast.util
+package datadog.trace.util
 
 import datadog.trace.test.util.DDSpecification
 import groovy.transform.CompileDynamic