Use ConcurrentHashMap for batching tasks per executor in TaskBatcher …

CHANGELOG.md

@@ -46,6 +46,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Changed http code on create index API with bad input raising NotXContentException from 500 to 400 ([#4773](https://github.com/opensearch-project/OpenSearch/pull/4773))
 - Change http code for DecommissioningFailedException from 500 to 400 ([#5283](https://github.com/opensearch-project/OpenSearch/pull/5283))
 - Require MediaType in Strings.toString API ([#6009](https://github.com/opensearch-project/OpenSearch/pull/6009))
+- Use ConcurrentHashMap for batching tasks per executor in TaskBatcher ([#5827](https://github.com/opensearch-project/OpenSearch/pull/5827))
 
 ### Deprecated
 

server/src/main/java/org/opensearch/cluster/service/TaskBatcher.java

@@ -46,6 +46,7 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -61,7 +62,7 @@ public abstract class TaskBatcher {
     private final Logger logger;
     private final PrioritizedOpenSearchThreadPoolExecutor threadExecutor;
     // package visible for tests
-    final Map<Object, LinkedHashSet<BatchedTask>> tasksPerBatchingKey = new HashMap<>();
+    final Map<Object, LinkedHashSet<BatchedTask>> tasksPerBatchingKey = new ConcurrentHashMap<>();
     private final TaskBatcherListener taskBatcherListener;
 
     public TaskBatcher(Logger logger, PrioritizedOpenSearchThreadPoolExecutor threadExecutor, TaskBatcherListener taskBatcherListener) {
@@ -93,12 +94,8 @@ public void submitTasks(List<? extends BatchedTask> tasks, @Nullable TimeValue t
                         IdentityHashMap::new
                     )
                 );
-
-            synchronized (tasksPerBatchingKey) {
-                LinkedHashSet<BatchedTask> existingTasks = tasksPerBatchingKey.computeIfAbsent(
-                    firstTask.batchingKey,
-                    k -> new LinkedHashSet<>(tasks.size())
-                );
+            LinkedHashSet<BatchedTask> newTasks = new LinkedHashSet<>(tasks);
+            tasksPerBatchingKey.merge(firstTask.batchingKey, newTasks, (existingTasks, updatedTasks) -> {
                 for (BatchedTask existing : existingTasks) {
                     // check that there won't be two tasks with the same identity for the same batching key
                     BatchedTask duplicateTask = tasksIdentity.get(existing.getTask());
@@ -112,8 +109,9 @@ public void submitTasks(List<? extends BatchedTask> tasks, @Nullable TimeValue t
                         );
                     }
                 }
-                existingTasks.addAll(tasks);
-            }
+                existingTasks.addAll(updatedTasks);
+                return existingTasks;
+            });
         } catch (Exception e) {
             taskBatcherListener.onSubmitFailure(tasks);
             throw e;
@@ -139,15 +137,13 @@ private void onTimeoutInternal(List<? extends BatchedTask> tasks, TimeValue time
             Object batchingKey = firstTask.batchingKey;
             assert tasks.stream().allMatch(t -> t.batchingKey == batchingKey)
                 : "tasks submitted in a batch should share the same batching key: " + tasks;
-            synchronized (tasksPerBatchingKey) {
-                LinkedHashSet<BatchedTask> existingTasks = tasksPerBatchingKey.get(batchingKey);
-                if (existingTasks != null) {
-                    existingTasks.removeAll(toRemove);
-                    if (existingTasks.isEmpty()) {
-                        tasksPerBatchingKey.remove(batchingKey);
-                    }
+            tasksPerBatchingKey.computeIfPresent(batchingKey, (tasksKey, existingTasks) -> {
+                existingTasks.removeAll(toRemove);
+                if (existingTasks.isEmpty()) {
+                    return null;
                 }
-            }
+                return existingTasks;
+            });
             taskBatcherListener.onTimeout(toRemove);
             onTimeout(toRemove, timeout);
         }
@@ -165,17 +161,15 @@ void runIfNotProcessed(BatchedTask updateTask) {
         if (updateTask.processed.get() == false) {
             final List<BatchedTask> toExecute = new ArrayList<>();
             final Map<String, List<BatchedTask>> processTasksBySource = new HashMap<>();
-            synchronized (tasksPerBatchingKey) {
-                LinkedHashSet<BatchedTask> pending = tasksPerBatchingKey.remove(updateTask.batchingKey);
-                if (pending != null) {
-                    for (BatchedTask task : pending) {
-                        if (task.processed.getAndSet(true) == false) {
-                            logger.trace("will process {}", task);
-                            toExecute.add(task);
-                            processTasksBySource.computeIfAbsent(task.source, s -> new ArrayList<>()).add(task);
-                        } else {
-                            logger.trace("skipping {}, already processed", task);
-                        }
+            LinkedHashSet<BatchedTask> pending = tasksPerBatchingKey.remove(updateTask.batchingKey);
+            if (pending != null) {
+                for (BatchedTask task : pending) {
+                    if (task.processed.getAndSet(true) == false) {
+                        logger.trace("will process {}", task);
+                        toExecute.add(task);
+                        processTasksBySource.computeIfAbsent(task.source, s -> new ArrayList<>()).add(task);
+                    } else {
+                        logger.trace("skipping {}, already processed", task);
                     }
                 }
             }