[fix]: fix ConcurrentModificationException when completing invocations (#362)

* fix ConcurrentModificationException when completing invocations

* fix non-existent branches in checkpointed result

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -5,14 +5,14 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
@@ -59,7 +59,7 @@ public class ExecutionManager implements AutoCloseable {
     private final DurableConfig durableConfig;
 
     // ===== Thread Coordination =====
-    private final Map<String, BaseDurableOperation> registeredOperations = Collections.synchronizedMap(new HashMap<>());
+    private final Map<String, BaseDurableOperation> registeredOperations = new ConcurrentHashMap<>();
     private final Set<String> activeThreads = Collections.synchronizedSet(new HashSet<>());
     private static final ThreadLocal<ThreadContext> currentThreadContext = new ThreadLocal<>();
     private final CompletableFuture<Void> executionExceptionFuture = new CompletableFuture<>();

sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.operation;
 
+import java.util.List;
 import java.util.function.Function;
 import software.amazon.awssdk.services.lambda.model.ContextOptions;
 import software.amazon.awssdk.services.lambda.model.Operation;
@@ -66,7 +67,8 @@ public ParallelOperation(
 
     @Override
     protected void handleCompletion(ConcurrencyCompletionStatus concurrencyCompletionStatus) {
-        var items = getBranches();
+
+        var items = List.copyOf(getBranches());
         var statuses = items.stream().map(this::getParallelItemStatus).toList();
         int succeededCount = Math.toIntExact(statuses.stream()
                 .filter(s -> s == ParallelResult.Status.SUCCEEDED)
@@ -76,6 +78,9 @@ protected void handleCompletion(ConcurrencyCompletionStatus concurrencyCompletio
         int skippedCount = items.size() - succeededCount - failedCount;
         cachedResult = new ParallelResult(
                 items.size(), succeededCount, failedCount, skippedCount, concurrencyCompletionStatus, statuses);
+
+        // Branches added after checkpoint will not exist in the checkpointed result, but they'll be in the returned
+        // value from get() method.
         sendOperationUpdate(OperationUpdate.builder()
                 .action(OperationAction.SUCCEED)
                 .subType(getSubType().getValue())
@@ -157,9 +162,14 @@ public <T> DurableFuture<T> branch(
             throw new IllegalStateException("Cannot add branches after join() has been called");
         }
 
-        // ConcurrencyOperation will skip this branch if skip=true
+        var nextBranchIndex = getBranches().size();
+
+        // ConcurrencyOperation will skip this branch if skip=true:
+        // 1. if the parallel operation is already completed (partialResult is not null)
+        // 2. if the branch is already skipped in the partialResult or nonexistent in the partialResult
         var skip = partialResult != null
-                && partialResult.statuses().get(getBranches().size()) == ParallelResult.Status.SKIPPED;
+                && (partialResult.statuses().size() <= nextBranchIndex
+                        || partialResult.statuses().get(nextBranchIndex) == ParallelResult.Status.SKIPPED);
         var serDes = config.serDes() == null ? getContext().getDurableConfig().getSerDes() : config.serDes();
         return enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip);
     }

sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java

@@ -270,7 +270,7 @@ void minSuccessful_notExecuteSkippedBranchWhenReplay() {
                         .status(OperationStatus.SUCCEEDED)
                         .contextDetails(ContextDetails.builder()
                                 .result(
-                                        "{\"succeeded\": 1, \"completionStatus\": \"MIN_SUCCESSFUL_REACHED\", \"statuses\":[\"SKIPPED\", \"SUCCEEDED\"]}")
+                                        "{\"size\": 2, \"skipped\": 1, \"succeeded\": 1, \"completionStatus\": \"MIN_SUCCESSFUL_REACHED\", \"statuses\":[\"SKIPPED\", \"SUCCEEDED\"]}")
                                 .build())
                         .build());
         when(executionManager.getOperationAndUpdateReplayState(CHILD_OP_2))