Replace ConcurrentHashMap with AtomicLongArray for raspErrorCodeCounter to improve performance and memory efficiency

Author: jandro996

internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java

@@ -8,8 +8,6 @@
 import java.util.Map;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicLongArray;
@@ -53,8 +51,8 @@ private WafMetricCollector() {
       new AtomicLongArray(RuleType.getNumValues());
   private static final AtomicLongArray raspTimeoutCounter =
       new AtomicLongArray(RuleType.getNumValues());
-  private static final ConcurrentMap<Integer, AtomicLongArray> raspErrorCodeCounter =
-      new ConcurrentSkipListMap<>();
+  private static final AtomicLongArray raspErrorCodeCounter =
+      new AtomicLongArray(WafErrorCode.values().length * RuleType.getNumValues());
   private static final AtomicLongArray wafErrorCodeCounter =
       new AtomicLongArray(WafErrorCode.values().length);
   private static final AtomicLongArray missingUserLoginQueue =
@@ -143,10 +141,14 @@ public void raspTimeout(final RuleType ruleType) {
     raspTimeoutCounter.incrementAndGet(ruleType.ordinal());
   }
 
-  public void raspErrorCode(final RuleType ruleType, final int errorCode) {
-    raspErrorCodeCounter
-        .computeIfAbsent(errorCode, code -> new AtomicLongArray(RuleType.getNumValues()))
-        .incrementAndGet(ruleType.ordinal());
+  public void raspErrorCode(RuleType ruleType, final int errorCode) {
+    WafErrorCode wafErrorCode = WafErrorCode.fromCode(errorCode);
+    // Unsupported waf error code
+    if (wafErrorCode == null) {
+      return;
+    }
+    int index = wafErrorCode.ordinal() * RuleType.getNumValues() + ruleType.ordinal();
+    raspErrorCodeCounter.incrementAndGet(index);
   }
 
   public void wafErrorCode(final int errorCode) {
@@ -319,15 +321,13 @@ public void prepareMetrics() {
     }
 
     // RASP rule type for each possible error code
-    for (Map.Entry<Integer, AtomicLongArray> entry : raspErrorCodeCounter.entrySet()) {
-      int errorCode = entry.getKey();
-      AtomicLongArray counters = entry.getValue();
-
+    for (WafErrorCode errorCode : WafErrorCode.values()) {
       for (RuleType ruleType : RuleType.values()) {
-        long count = counters.getAndSet(ruleType.ordinal(), 0);
+        int index = errorCode.ordinal() * RuleType.getNumValues() + ruleType.ordinal();
+        long count = raspErrorCodeCounter.getAndSet(index, 0);
         if (count > 0) {
           if (!rawMetricsQueue.offer(
-              new RaspError(count, ruleType, WafMetricCollector.wafVersion, errorCode))) {
+              new RaspError(count, ruleType, WafMetricCollector.wafVersion, errorCode.getCode()))) {
             return;
           }
         }

internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy

@@ -182,7 +182,19 @@ class WafMetricCollectorTest extends DDSpecification {
     raspTimeout.metricName == 'rasp.timeout'
     raspTimeout.tags.toSet() == ['rule_type:sql_injection', 'waf_version:waf_ver1'].toSet()
 
-    def raspInvalidCode = (WafMetricCollector.RaspError) metrics[11]
+    def raspInvalidObjectCode = (WafMetricCollector.RaspError)metrics[11]
+    raspInvalidObjectCode.type == 'count'
+    raspInvalidObjectCode.value == 1
+    raspInvalidObjectCode.namespace == 'appsec'
+    raspInvalidObjectCode.metricName == 'rasp.error'
+    raspInvalidObjectCode.tags.toSet() == [
+      'rule_type:sql_injection',
+      'waf_version:waf_ver1',
+      'waf_error:' + DD_WAF_RUN_INVALID_OBJECT_ERROR
+    ]
+    .toSet()
+
+    def raspInvalidCode = (WafMetricCollector.RaspError)metrics[12]
     raspInvalidCode.type == 'count'
     raspInvalidCode.value == 1
     raspInvalidCode.namespace == 'appsec'
@@ -195,18 +207,6 @@ class WafMetricCollectorTest extends DDSpecification {
       'waf_error:' + DD_WAF_RUN_INTERNAL_ERROR
     ].toSet()
 
-    def raspInvalidObjectCode = (WafMetricCollector.RaspError)metrics[12]
-    raspInvalidObjectCode.type == 'count'
-    raspInvalidObjectCode.value == 1
-    raspInvalidObjectCode.namespace == 'appsec'
-    raspInvalidObjectCode.metricName == 'rasp.error'
-    raspInvalidObjectCode.tags.toSet() == [
-      'rule_type:sql_injection',
-      'waf_version:waf_ver1',
-      'waf_error:' + DD_WAF_RUN_INVALID_OBJECT_ERROR
-    ]
-    .toSet()
-
     def wafInvalidCode = (WafMetricCollector.WafError)metrics[13]
     wafInvalidCode.type == 'count'
     wafInvalidCode.value == 1