Reworked ApiAccessTracker to avoid global bottleneck

Author: ValentinZakharov

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

@@ -1,8 +1,9 @@
 package com.datadog.appsec.api.security;
 
-import java.util.Collections;
-import java.util.LinkedHashMap;
+import java.util.Deque;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
 
 /**
  * The ApiAccessTracker class provides a mechanism to track API access events, managing them within
@@ -20,23 +21,20 @@
 public class ApiAccessTracker {
   private static final int INTERVAL_SECONDS = 30;
   private static final int MAX_SIZE = 4096;
-  private final Map<Long, Long> apiAccessLog; // Map<hash, timestamp>
+  private final Map<Long, Long> apiAccessMap; // Map<hash, timestamp>
+  private final Deque<Long> apiAccessQueue;   // hashes ordered by access time
   private final long expirationTimeInMs;
+  private final int capacity;
 
   public ApiAccessTracker() {
     this(MAX_SIZE, INTERVAL_SECONDS * 1000);
   }
 
   public ApiAccessTracker(int capacity, long expirationTimeInMs) {
+    this.capacity = capacity;
     this.expirationTimeInMs = expirationTimeInMs;
-    this.apiAccessLog =
-        Collections.synchronizedMap(
-            new LinkedHashMap<Long, Long>() {
-              @Override
-              protected boolean removeEldestEntry(Map.Entry<Long, Long> eldest) {
-                return size() > capacity;
-              }
-            });
+    this.apiAccessMap = new ConcurrentHashMap<>();
+    this.apiAccessQueue = new ConcurrentLinkedDeque<>();
   }
 
   /**
@@ -54,17 +52,41 @@ public boolean updateApiAccessIfExpired(String route, String method, int statusC
     long currentTime = System.currentTimeMillis();
     long hash = computeApiHash(route, method, statusCode);
 
-    synchronized (apiAccessLog) {
-      if (apiAccessLog.containsKey(hash)) {
-        long lastAccessTime = apiAccessLog.get(hash);
-        if (currentTime - lastAccessTime > expirationTimeInMs) {
-          apiAccessLog.put(hash, currentTime);
-          return true;
-        }
-        return false;
+    cleanupExpiredEntries(currentTime);
+
+    // New or updated record
+    boolean isNewOrUpdated = false;
+    if (!apiAccessMap.containsKey(hash)
+        || currentTime - apiAccessMap.get(hash) > expirationTimeInMs) {
+      apiAccessMap.put(hash, currentTime); // Update timestamp
+      // move hash to the end of the queue
+      apiAccessQueue.remove(hash);
+      apiAccessQueue.addLast(hash);
+      isNewOrUpdated = true;
+    }
+
+    // Remove the oldest hash if capacity is reached
+    while (apiAccessQueue.size() > this.capacity) {
+      Long oldestHash = apiAccessQueue.pollFirst();
+      if (oldestHash != null) {
+        apiAccessMap.remove(oldestHash);
+      }
+    }
+
+    return isNewOrUpdated;
+  }
+
+  private void cleanupExpiredEntries(long currentTime) {
+    while (!apiAccessQueue.isEmpty()) {
+      Long oldestHash = apiAccessQueue.peekFirst();
+      if (oldestHash == null) break;
+
+      Long lastAccessTime = apiAccessMap.get(oldestHash);
+      if (lastAccessTime == null || currentTime - lastAccessTime > expirationTimeInMs) {
+        apiAccessQueue.pollFirst(); // remove from head
+        apiAccessMap.remove(oldestHash);
       } else {
-        apiAccessLog.put(hash, currentTime);
-        return true;
+        break; // is up-to-date
       }
     }
   }

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

@@ -9,18 +9,18 @@ class ApiAccessTrackerTest extends DDSpecification {
 
     when: "Adding new api access"
     tracker.updateApiAccessIfExpired("route1", "GET", 200)
-    def firstAccessTime = tracker.apiAccessLog.values().iterator().next()
+    def firstAccessTime = tracker.apiAccessMap.values().iterator().next()
 
     then: "The access is added"
-    tracker.apiAccessLog.size() == 1
+    tracker.apiAccessMap.size() == 1
 
     when: "Waiting more than expiration time and adding another access with the same key"
     Thread.sleep(1100) // Waiting more than 1 second to ensure expiration
     tracker.updateApiAccessIfExpired("route1", "GET", 200)
-    def secondAccessTime = tracker.apiAccessLog.values().iterator().next()
+    def secondAccessTime = tracker.apiAccessMap.values().iterator().next()
 
     then: "The access is updated and moved to the end"
-    tracker.apiAccessLog.size() == 1
+    tracker.apiAccessMap.size() == 1
     secondAccessTime > firstAccessTime
   }
 
@@ -34,9 +34,9 @@ class ApiAccessTrackerTest extends DDSpecification {
     tracker.updateApiAccessIfExpired("route2", "POST", 404)
 
     then: "The oldest access is removed"
-    tracker.apiAccessLog.size() == 1
-    !tracker.apiAccessLog.containsKey(tracker.computeApiHash("route1", "GET", 200))
-    tracker.apiAccessLog.containsKey(tracker.computeApiHash("route2", "POST", 404))
+    tracker.apiAccessMap.size() == 1
+    !tracker.apiAccessMap.containsKey(tracker.computeApiHash("route1", "GET", 200))
+    tracker.apiAccessMap.containsKey(tracker.computeApiHash("route2", "POST", 404))
   }
 
   def "should not update access if not expired"() {
@@ -50,6 +50,6 @@ class ApiAccessTrackerTest extends DDSpecification {
 
     then: "The access is not updated"
     !updatedBeforeExpiration
-    tracker.apiAccessLog.get(tracker.computeApiHash("route1", "GET", 200)) == updateTime
+    tracker.apiAccessMap.get(tracker.computeApiHash("route1", "GET", 200)) == updateTime
   }
 }