Use ConcurrentHashMap for batching tasks per executor in TaskBatcher …

[amkhar]

Description

Master tasks execution is based on batchingKey. This is stored in a HashMap and every time there is a task to be enqueued, we take a lock on whole code block where map is being used. This is an optimization to use ConcurrentHashMap and improve the overall enqueue time. This will help when master is overloaded with pending tasks. Using synchronized block on LinkedHashSet operations as it's not concurrent naturally and we're removing synchronized block on overall map.

Update : Using internal methods(merge and computeIfPresent) of ConcurrentHashMap instead of taking lock on LinkedHashSet as there are more race conditions if we don't lock the bucket of CHM while operating over LinkedHashSet. This is much cleaner and correct implementation.

As we know that execution time would still be same for this unit test, enqueuing time has improved more than 43%.

Final comparison based on load run

After changing to CHM and using synchronized block on all LHS operations, we’re able to see tasks are being enqueued and executed faster than before.

Approach	Threads & Tasks	Avg time taken in enqueue & execution of all tasks (seconds)	Comments
Current production	5K & 2.5K	309
Concurrent HashMap without SynchronizedSet	5K & 2.5K	177	Race conditions while implementing the logic
CHM internal methods	5K & 2.5K	177
Concurrent HashMap with SynchronizedSet	5K & 2.5K	161
Current production	2K & 5K	392
Concurrent HashMap without SynchronizedSet	2K & 5K	103	Race conditions while implementing the logic
CHM internal methods	2K & 5K	104
Concurrent HashMap with SynchronizedSet	2K & 5K	N/A	Tests timed out due to CyclicBarrier wait

Load setup

Ran

OpenSearch/server/src/test/java/org/opensearch/cluster/service/TaskBatcherTests.java

Line 214 in 082f059

public void testTasksAreExecutedInOrder() throws BrokenBarrierException, InterruptedException {

7-8 times with 5000 executors and 2500 tasks per executor. So in short, there were 5000 entries in CHM and 2500 entries in each key’s value(LHS).
In the second run switched executors to 2000 and tasks to 5000.

Issues Resolved

• [Optimization] TaskBatcher to use ConcurrentHashMap for batching tasks per executor #5099

Check List

• New functionality includes testing.
  • All tests pass
• New functionality has been documented.
  • New functionality has javadoc added
• Commits are signed per the DCO using --signoff
• Commit changes are listed out in CHANGELOG.md file (See: Changelog)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/9289/
• CommitID: 91d3182
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: UNSTABLE ❕
• TEST FAILURES:

      1 org.opensearch.http.netty4.Netty4HttpRequestSizeLimitIT.testLimitsInFlightRequests
      1 org.opensearch.cluster.service.MasterServiceTests.testClusterStateBatchedUpdates
      1 org.opensearch.cluster.service.MasterServiceTests.testClusterStateBatchedUpdates
      1 org.opensearch.cluster.service.MasterServiceTests.classMethod
      1 org.opensearch.cluster.service.MasterServiceTests.classMethod

• URL: https://build.ci.opensearch.org/job/gradle-check/9290/
• CommitID: 57b96f6
  Please review all flaky tests that succeeded after retry and create an issue if one does not already exist to track the flaky failure.

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/9286/
• CommitID: e8dde1e
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[amkhar]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/9286/
• CommitID: e8dde1e
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

Flaky tests issues link

SimpleBlocksIT - #2472
IndexingIT, RecoveryIT - #5766
SegmentReplication - #5669
BulkProcessorRetryIT - #5830

[andrross]

I'm concerned about a potential race with the following scenario:

Thread 1: executes line 98 and adds an empty LinkedHashSet to the map
Thread 2: executes lines 145-152, which will find an empty LinkedHashSet and remove it from the map
Thread 1: executes lines 103-118, which adds tasks to the LinkedHashSet but the set has been removed from the map

Is that scenario possible?

[shwetathareja]

+1 the race can possibly occur. The synchronize block on the existingTasks object is not sufficient to prevent any operation on that key in the main map tasksPerBatchingKey.
@amkhar it is better to use compute or preferablymerge ( as it provides capability to provide default value as well) on tasksPerBatchingKey and perform the duplicate check, adding tasks etc. (Line 104-117) in the remappingFunction as argument which will prevent any update to the same key from any other code path executing in different thread. For the other code paths as well to remove the key etc. should also use merge

[amkhar]

@andrross @shwetathareja
Sorry for delayed response.
Yes, race condition is possible. A little different than what Andrew explained.

Thread 1 &2 both working on same batchingKey, and note that timeout method can only be called once thread has already executed line number 126 (only after adding the tasks).

Thread 1: executes till line number 126, so it has already added some list in the value of batchingKey
Thread 2: executes line number 98 and get a value from map which has a list of 1 task
Thread 1's async call back due to timeout: executes line 147-151 and removes the entry from map.
Thread 2: reaches on 103 and waits
Thread 2: executes the lines 104-117 and adds another entry in LinkedHashSet but it has already been removed from previous callback.

So, yes a fix is needed for this condition. I've found a way to handle this. It requires adding one line after 117

tasksPerBatchingKey.putIfAbsent(firstTask.batchingKey, existingTasks); to make sure that if key-value was removed, we put it again while adding new tasks in the queue. And we also need to add a check to remove already processed tasks like below at line number 116, because runIfNotProcessed method can also update and process some tasks in another thread.

if (existing.processed.get()) {
     tasksIterator.remove();
}

Something like this

synchronized (existingTasks) {
                Iterator<BatchedTask> tasksIterator = existingTasks.iterator();
                while (tasksIterator.hasNext()) {
                    BatchedTask existing = tasksIterator.next();
                    // check that there won't be two tasks with the same identity for the same batching key
                    BatchedTask duplicateTask = tasksIdentity.get(existing.getTask());
                    if (duplicateTask != null) {
                        throw new IllegalStateException(
                            "task ["
                                + duplicateTask.describeTasks(Collections.singletonList(existing))
                                + "] with source ["
                                + duplicateTask.source
                                + "] is already queued"
                        );
                    }
                    if (existing.processed.get()){
                        tasksIterator.remove();
                    }
                }
                existingTasks.addAll(tasks);
                tasksPerBatchingKey.putIfAbsent(firstTask.batchingKey, existingTasks);
            }

But thanks Shweta for suggesting merge, let me try that also. Looks similar to what we're trying to achieve.

[shwetathareja]

Thanks @amkhar I have one concern with the above proposed fix. The synchronized block is on existingTasks and putIfAbsent will add the key back again (in case it was removed during the timeout code execution form line 151). But lets says another Thread 3 came and at that time there was no entry in the map, it would have taken lock on a diff existingTasks object altogether. So now, 2 threads are executing the code which is in synchronized block from line 103.

[amkhar]

@andrross
@shwetathareja Agreed. That was another case. I realized that, taking separate locks on LinkedHashSet is problematic. If we just maintain the CHM bucket level lock, by using its library method, that will be clean and correct.

This is how code will look using merge or computeIfPresent (only writing the concerning code block)

submitTasks - create new or update existing list then update Map

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());
                    if (duplicateTask != null) {
                        throw new IllegalStateException(
                            "task ["
                                + duplicateTask.describeTasks(Collections.singletonList(existing))
                                + "] with source ["
                                + duplicateTask.source
                                + "] is already queued"
                        );
                    }
                }
                existingTasks.addAll(updatedTasks);
                return existingTasks;
            });

And runIfNotProcessed - here for iteration lock is not needed, because no other flow will be accessing the same object after it's removed from the Map.

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);
                    }
                }
            }

onTimeoutInterval - Re-compute the value once timed out tasks are removed, and remove it if nothing left in LinkedHashSet.

tasksPerBatchingKey.computeIfPresent(batchingKey, (tasksKey, currentTasks) -> {
                currentTasks.removeAll(toRemove);
                if (currentTasks.isEmpty()) {
                    return null;
                }
                return currentTasks;
            });

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/10294/
• CommitID: e79b69f
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/10296/
• CommitID: fa73fda
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/10295/
• CommitID: 57b96f6
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: SUCCESS ✅
• URL: https://build.ci.opensearch.org/job/gradle-check/10297/
• CommitID: 56cd9d5

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/10293/
• CommitID: 57b96f6
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[shwetathareja]

Changes LGTM. @amkhar can you please ensure we have test to submit tasks from multiple threads else add them.

[amkhar]

Changes LGTM. @amkhar can you please ensure we have test to submit tasks from multiple threads else add them.

Hey @shwetathareja , there are no tests which invokes parallel threads on single batchingKey. Existing test testTasksAreExecutedInOrder hit submitTasks method in for loop so order of execution can be asserted. In case we hit parallel thread on single batchingKey, we can't assert execution order because it's actually non-deterministic that which thread will hit first.
So, only assertion we can do is "no tasks are dropped in execution" (due to the modification or list and CHM with new change), i.e. checking if count of tasks execution per batchingKey is same as the number of parallel threads we run for that batchingKey.

checking if count of tasks execution per batchingKey is same as the number of parallel threads we run for that batchingKey.

This validation is already happening in testTasksAreExecutedInOrder test but there we hit submitTasks serially. So it'll be a similar test but we'll be hitting submitTasks in parallel. Does that make sense ? if it looks good, I can update the PR.

Feel free to suggest other assertions which I may be missing :)

[shwetathareja]

This validation is already happening in testTasksAreExecutedInOrder test but there we hit submitTasks serially. So it'll be a similar test but we'll be hitting submitTasks in parallel. Does that make sense ? if it looks good, I can update the PR.

@amkhar right we need to write a test to ensure none of the tasks are lost due to any race condition in the code when tasks are submitted in parallel.

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/10583/
• CommitID: d771fd4
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[andrross]

Please rebase the latest from main, which should fix up the unrelated test failures.

[andrross]

I think we can probably omit this from the changelog because it is a targeted performance improvement that doesn't change any user-facing behavior. What do you think? (we haven't been super consistent about this, by the way)

If this does stay in the changelog, it should be in the [Unreleased 2.x] section because this will be backported and released in the next 2.x released.

[shwetathareja]

I think if we are keeping changelog strictly to user facing behavior, then I agree this should be removed.

[github-actions]

Gradle Check (Jenkins) Run Completed with:

• RESULT: FAILURE ❌
• URL: https://build.ci.opensearch.org/job/gradle-check/12003/
• CommitID: b96e88a
  Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green.
  Is the failure a flaky test unrelated to your change?

[amkhar]

Opened an issue for flaky test failure : #6565