Fix ConcurrentModificationException in LocalActivityMeteringHelper (#2108)

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandler.java

@@ -422,30 +422,32 @@ public void cancel(HistoryEvent cancelEvent) {
     }
   }
 
-  private static class LocalActivityMeteringHelper {
+  @VisibleForTesting
+  static class LocalActivityMeteringHelper {
     private final Map<String, AtomicInteger> firstWftActivities = new HashMap<>();
     private final Map<String, AtomicInteger> nonFirstWftActivities = new HashMap<>();
     private final Set<String> completed = new HashSet<>();
 
-    private void newWFTStarting() {
+    void newWFTStarting() {
       for (String activityId : firstWftActivities.keySet()) {
-        AtomicInteger removed = firstWftActivities.remove(activityId);
-        removed.set(0);
-        nonFirstWftActivities.put(activityId, removed);
+        AtomicInteger attemptCount = firstWftActivities.get(activityId);
+        attemptCount.set(0);
+        nonFirstWftActivities.put(activityId, attemptCount);
       }
+      firstWftActivities.clear();
     }
 
-    private void addNewLocalActivity(ExecuteLocalActivityParameters params) {
+    void addNewLocalActivity(ExecuteLocalActivityParameters params) {
       AtomicInteger attemptsDuringWFTCounter = new AtomicInteger(0);
       params.setOnNewAttemptCallback(attemptsDuringWFTCounter::incrementAndGet);
       firstWftActivities.put(params.getActivityId(), attemptsDuringWFTCounter);
     }
 
-    private void markLocalActivityComplete(String activityId) {
+    void markLocalActivityComplete(String activityId) {
       completed.add(activityId);
     }
 
-    private int getNonfirstAttempts() {
+    int getNonfirstAttempts() {
       int result =
           nonFirstWftActivities.values().stream()
               .map(ai -> ai.getAndSet(0))

temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java

@@ -36,11 +36,9 @@
 import io.temporal.api.history.v1.History;
 import io.temporal.api.history.v1.HistoryEvent;
 import io.temporal.api.taskqueue.v1.StickyExecutionAttributes;
-import io.temporal.api.workflowservice.v1.GetSystemInfoResponse;
-import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse;
-import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse;
-import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc;
+import io.temporal.api.workflowservice.v1.*;
 import io.temporal.internal.common.InternalUtils;
+import io.temporal.internal.statemachines.ExecuteLocalActivityParameters;
 import io.temporal.internal.worker.SingleWorkerOptions;
 import io.temporal.internal.worker.WorkflowExecutorCache;
 import io.temporal.internal.worker.WorkflowRunLockManager;
@@ -79,7 +77,6 @@ public void ifStickyExecutionAttributesAreNotSetThenWorkflowsAreNotCached() thro
     // Act
     WorkflowTaskHandler.Result result =
         taskHandler.handleWorkflowTask(HistoryUtils.generateWorkflowTaskWithInitialHistory());
-
     // Assert
     assertEquals(0, cache.size());
     assertNotNull(result.getTaskCompleted());
@@ -142,6 +139,41 @@ public void workflowTaskFailOnIncompleteHistory() throws Throwable {
         result.getTaskFailed().getFailure().getMessage());
   }
 
+  @Test
+  public void localActivityMeteringHelper() {
+    ReplayWorkflowRunTaskHandler.LocalActivityMeteringHelper laMeteringHelper =
+        new ReplayWorkflowRunTaskHandler.LocalActivityMeteringHelper();
+    ExecuteLocalActivityParameters executeLA =
+        new ExecuteLocalActivityParameters(
+            PollActivityTaskQueueResponse.newBuilder().setActivityId("1"),
+            null,
+            0,
+            null,
+            false,
+            null);
+    laMeteringHelper.addNewLocalActivity(executeLA);
+    laMeteringHelper.addNewLocalActivity(
+        new ExecuteLocalActivityParameters(
+            PollActivityTaskQueueResponse.newBuilder().setActivityId("2"),
+            null,
+            0,
+            null,
+            false,
+            null));
+    for (int i = 0; i < 5; i++) {
+      executeLA.getOnNewAttemptCallback().apply();
+    }
+    // Verify retries are not counted for the first task
+    assertEquals(0, laMeteringHelper.getNonfirstAttempts());
+    laMeteringHelper.newWFTStarting();
+    assertEquals(0, laMeteringHelper.getNonfirstAttempts());
+    // Verify retries are counted for the non first task
+    for (int i = 0; i < 5; i++) {
+      executeLA.getOnNewAttemptCallback().apply();
+    }
+    assertEquals(5, laMeteringHelper.getNonfirstAttempts());
+  }
+
   @Test
   public void ifStickyExecutionAttributesAreSetThenWorkflowsAreCached() throws Throwable {
     assumeFalse("skipping for docker tests", SDKTestWorkflowRule.useExternalService);