fix concurrency bug

Signed-off-by: Amit Galitzky <amgalitz@amazon.com>

Author: amitgalitz

src/main/java/org/opensearch/timeseries/ratelimit/BatchWorker.java

@@ -113,8 +113,11 @@ protected void execute(Runnable afterProcessCallback, Runnable emptyQueueCallbac
         if (false == toProcess.isEmpty()) {
             final List<String> inflights = new ArrayList<>();
             for (RequestType request : toProcess) {
-                inflightConfigs.add(request.getConfigId());
-                inflights.add(request.getConfigId());
+                String configId = request.getConfigId();
+                if (configId != null) {
+                    inflightConfigs.add(configId);
+                    inflights.add(configId);
+                }
             }
 
             BatchRequestType batchRequest = toBatchRequest(toProcess);

src/main/java/org/opensearch/timeseries/ratelimit/RateLimitedRequestWorker.java

@@ -26,6 +26,7 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -249,7 +250,7 @@ public RateLimitedRequestWorker(
         this.stateTtl = stateTtl;
         this.nodeStateManager = nodeStateManager;
         this.context = context;
-        this.inflightConfigs = new HashSet<>();
+        this.inflightConfigs = ConcurrentHashMap.newKeySet();
     }
 
     public String getWorkerName() {

src/main/java/org/opensearch/timeseries/ratelimit/SingleRequestWorker.java

@@ -92,7 +92,10 @@ protected void execute(Runnable afterProcessCallback, Runnable emptyQueueCallbac
         if (false == queue.isEmpty()) {
             request = queue.poll();
             if (request != null) {
-                inflightConfigs.add(request.getConfigId());
+                String configId = request.getConfigId();
+                if (configId != null) {
+                    inflightConfigs.add(configId);
+                }
             }
         }
 

src/test/java/org/opensearch/ad/ratelimit/ResultWriteWorkerTests.java

@@ -26,9 +26,12 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.ConcurrentModificationException;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Random;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.opensearch.OpenSearchStatusException;
@@ -209,4 +212,72 @@ public void testOverloaded() {
         verify(resultHandler, times(1)).flush(any(), any());
         verify(nodeStateManager, times(1)).setException(eq(detectorId), any(OpenSearchRejectedExecutionException.class));
     }
+
+    public void testConcurrentModificationException() throws InterruptedException {
+        for (int attempt = 0; attempt < 10; attempt++) {
+            if (runConcurrencyTest()) {
+                fail("ConcurrentModificationException occurred on attempt " + (attempt + 1));
+            }
+            Thread.sleep(5);
+        }
+    }
+
+   private boolean runConcurrencyTest() throws InterruptedException {
+        final int numberOfThreads = 20; final CountDownLatch startLatch = new CountDownLatch(1);
+        final CountDownLatch endLatch = new CountDownLatch(numberOfThreads);
+        final AtomicBoolean exceptionOccurred = new AtomicBoolean(false);
+
+        doAnswer(invocation -> {
+            ActionListener<ResultBulkResponse> listener = invocation.getArgument(1);
+            try {
+                Thread.sleep(5);
+                listener.onResponse(new ResultBulkResponse(Collections.emptyList()));
+            } catch (Exception e) {
+                listener.onFailure(e);
+            }
+            return null;
+        }).when(resultHandler).flush(any(), any());
+
+        doAnswer(invocation -> {
+            Exception e = invocation.getArgument(1);
+            if (e instanceof ConcurrentModificationException) {
+                exceptionOccurred.set(true);
+            }
+            return null;
+        }).when(nodeStateManager).setException(any(), any());
+
+        for (int i = 0; i < numberOfThreads; i++) {
+            Thread thread = new Thread(() -> {
+                try {
+                    startLatch.await();
+                    // Each thread adds requests
+                    for (int j = 0; j < 30; j++) {
+                        try {
+                            String detectorId = "detector-" + (j % 3);
+                            resultWriteQueue.put(
+                                    new ADResultWriteRequest(
+                                            System.currentTimeMillis() + j,
+                                            detectorId,
+                                            RequestPriority.MEDIUM,
+                                            detectResult,
+                                            null,
+                                            null
+                                    )
+                            );
+                        } catch (ConcurrentModificationException e) {
+                            exceptionOccurred.set(true);
+                            return;
+                        }
+                    }
+                } catch (Exception ignored) {
+                } finally {
+                    endLatch.countDown();
+                }
+            });
+            thread.start();
+        }
+        startLatch.countDown();
+        assertTrue("Test should complete", endLatch.await(30, TimeUnit.SECONDS));
+       return exceptionOccurred.get();
+    }
 }

src/test/java/org/opensearch/ad/task/ADTaskManagerTests.java

@@ -653,6 +653,14 @@ public void testUpdateLatestRealtimeTaskOnCoordinatingNode() {
         verify(actionListener, times(1)).onResponse(any());
     }
 
+    public void testTriageStateWithNullRcfTotalUpdates() {
+        Boolean hasResult = true;
+        String error = null;
+        Long rcfTotalUpdates = null;
+        String result = adTaskManager.triageState(hasResult, error, rcfTotalUpdates);
+        assertEquals(TaskState.INIT.name(), result);
+    }
+
     public void testGetLocalADTaskProfilesByDetectorId() {
         doReturn(node1).when(clusterService).localNode();
         when(adTaskCacheManager.isHCTaskRunning(anyString())).thenReturn(true);