Use ConcurrentHashMap internal methods instead of taking lock on Link…

…edHashSet object Signed-off-by: Aman Khare <[email protected]>
opensearch-project · Jan 31, 2023 · e79b69f · e79b69f
1 parent 57b96f6
commit e79b69f
Showing 1 changed file with 18 additions and 27 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/service/TaskBatcher.java b/server/src/main/java/org/opensearch/cluster/service/TaskBatcher.java
@@ -94,13 +94,8 @@ public void submitTasks(List<? extends BatchedTask> tasks, @Nullable TimeValue t
                         IdentityHashMap::new
                     )
                 );
-
-            LinkedHashSet<BatchedTask> existingTasks = tasksPerBatchingKey.computeIfAbsent(
-                firstTask.batchingKey,
-                k -> new LinkedHashSet<>(tasks.size())
-            );
-            // Locking on LinkedHashSet is necessary as it is being modified in concurrent manner.
-            synchronized (existingTasks) {
+            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());
@@ -114,9 +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;
@@ -142,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;
-            LinkedHashSet<BatchedTask> existingTasks = tasksPerBatchingKey.get(batchingKey);
-            if (existingTasks != null) {
-                synchronized (existingTasks) {
-                    existingTasks.removeAll(toRemove);
+            tasksPerBatchingKey.computeIfPresent(batchingKey, (tasksKey, currentTasks) -> {
+                currentTasks.removeAll(toRemove);
+                if (currentTasks.isEmpty()) {
+                    return null;
                 }
-                if (existingTasks.isEmpty()) {
-                    tasksPerBatchingKey.remove(batchingKey);
-                }
-            }
+                return currentTasks;
+            });
             taskBatcherListener.onTimeout(toRemove);
             onTimeout(toRemove, timeout);
         }
@@ -170,15 +163,13 @@ void runIfNotProcessed(BatchedTask updateTask) {
             final Map<String, List<BatchedTask>> processTasksBySource = new HashMap<>();
             LinkedHashSet<BatchedTask> pending = tasksPerBatchingKey.remove(updateTask.batchingKey);
             if (pending != null) {
-                synchronized (pending) {
-                    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);
-                        }
+                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);
                     }
                 }
             }