Make CancellationScopeImpl more deterministic (#2512)

Make CancellationScopeImpl more deterministic

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/common/SdkFlag.java

@@ -14,6 +14,10 @@ public enum SdkFlag {
    * Changes behavior of GetVersion to never yield.
    */
   SKIP_YIELD_ON_VERSION(2),
+  /*
+   * Changes behavior of CancellationScope to cancel children in a deterministic order.
+   */
+  DETERMINISTIC_CANCELLATION_SCOPE_ORDER(3),
   UNKNOWN(Integer.MAX_VALUE);
 
   private final int value;

temporal-sdk/src/main/java/io/temporal/internal/sync/CancellationScopeImpl.java

@@ -1,10 +1,7 @@
 package io.temporal.internal.sync;
 
 import io.temporal.workflow.*;
-import java.util.ArrayDeque;
-import java.util.Deque;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.*;
 
 class CancellationScopeImpl implements CancellationScope {
 
@@ -36,7 +33,9 @@ private static void popCurrent(CancellationScopeImpl expected) {
 
   private final Runnable runnable;
   private CancellationScopeImpl parent;
-  private final Set<CancellationScopeImpl> children = new HashSet<>();
+  // We use a LinkedHashSet because we will iterate through the children, so we need to keep a
+  // deterministic order.
+  private final Set<CancellationScopeImpl> children;
 
   /**
    * When disconnected scope has no parent and thus doesn't receive cancellation requests from it.
@@ -45,20 +44,37 @@ private static void popCurrent(CancellationScopeImpl expected) {
 
   private String reason;
 
-  CancellationScopeImpl(boolean ignoreParentCancellation, Runnable runnable) {
-    this(ignoreParentCancellation, runnable, current());
+  CancellationScopeImpl(
+      boolean ignoreParentCancellation, boolean deterministicOrder, Runnable runnable) {
+    this(ignoreParentCancellation, deterministicOrder, runnable, current());
   }
 
-  CancellationScopeImpl(boolean detached, Runnable runnable, CancellationScopeImpl parent) {
+  CancellationScopeImpl(
+      boolean detached,
+      boolean deterministicOrder,
+      Runnable runnable,
+      CancellationScopeImpl parent) {
     this.detached = detached;
     this.runnable = runnable;
+    if (deterministicOrder) {
+      this.children = new LinkedHashSet<>();
+    } else {
+      this.children = new HashSet<>();
+    }
     setParent(parent);
   }
 
   public CancellationScopeImpl(
-      boolean ignoreParentCancellation, Functions.Proc1<CancellationScope> proc) {
+      boolean ignoreParentCancellation,
+      boolean deterministicOrder,
+      Functions.Proc1<CancellationScope> proc) {
     this.detached = ignoreParentCancellation;
     this.runnable = () -> proc.apply(this);
+    if (deterministicOrder) {
+      this.children = new LinkedHashSet<>();
+    } else {
+      this.children = new HashSet<>();
+    }
     setParent(current());
   }
 

temporal-sdk/src/main/java/io/temporal/internal/sync/DeterministicRunnerImpl.java

@@ -4,6 +4,7 @@
 import com.google.common.primitives.Ints;
 import io.temporal.common.context.ContextPropagator;
 import io.temporal.internal.WorkflowThreadMarker;
+import io.temporal.internal.common.SdkFlag;
 import io.temporal.internal.context.ContextThreadLocal;
 import io.temporal.internal.worker.WorkflowExecutorCache;
 import io.temporal.serviceclient.CheckedExceptionWrapper;
@@ -157,7 +158,12 @@ static void setCurrentThreadInternal(WorkflowThread coroutine) {
     // a bad practice
     this.workflowContext.setRunner(this);
     this.cache = cache;
-    this.runnerCancellationScope = new CancellationScopeImpl(true, null, null);
+    boolean deterministicCancellationScopeOrder =
+        workflowContext
+            .getReplayContext()
+            .checkSdkFlag(SdkFlag.DETERMINISTIC_CANCELLATION_SCOPE_ORDER);
+    this.runnerCancellationScope =
+        new CancellationScopeImpl(true, deterministicCancellationScopeOrder, null, null);
     this.rootRunnable = root;
   }
 

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowInternal.java

@@ -24,6 +24,7 @@
 import io.temporal.internal.WorkflowThreadMarker;
 import io.temporal.internal.common.ActivityOptionUtils;
 import io.temporal.internal.common.NonIdempotentHandle;
+import io.temporal.internal.common.SdkFlag;
 import io.temporal.internal.common.SearchAttributesUtil;
 import io.temporal.internal.logging.ReplayAwareLogger;
 import io.temporal.internal.statemachines.UnsupportedContinueAsNewRequest;
@@ -546,12 +547,20 @@ public static Promise<Object> promiseAnyOf(Promise<?>... promises) {
   }
 
   public static CancellationScope newCancellationScope(boolean detached, Runnable runnable) {
-    return new CancellationScopeImpl(detached, runnable);
+    boolean deterministicCancellationScopeOrder =
+        getRootWorkflowContext()
+            .getReplayContext()
+            .checkSdkFlag(SdkFlag.DETERMINISTIC_CANCELLATION_SCOPE_ORDER);
+    return new CancellationScopeImpl(detached, deterministicCancellationScopeOrder, runnable);
   }
 
   public static CancellationScope newCancellationScope(
       boolean detached, Functions.Proc1<CancellationScope> proc) {
-    return new CancellationScopeImpl(detached, proc);
+    boolean deterministicCancellationScopeOrder =
+        getRootWorkflowContext()
+            .getReplayContext()
+            .checkSdkFlag(SdkFlag.DETERMINISTIC_CANCELLATION_SCOPE_ORDER);
+    return new CancellationScopeImpl(detached, deterministicCancellationScopeOrder, proc);
   }
 
   public static CancellationScopeImpl currentCancellationScope() {

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowThreadImpl.java

@@ -4,6 +4,7 @@
 import io.temporal.common.context.ContextPropagator;
 import io.temporal.failure.CanceledFailure;
 import io.temporal.internal.common.NonIdempotentHandle;
+import io.temporal.internal.common.SdkFlag;
 import io.temporal.internal.context.ContextThreadLocal;
 import io.temporal.internal.logging.LoggerTag;
 import io.temporal.internal.replay.ReplayWorkflowContext;
@@ -55,7 +56,11 @@ class RunnableWrapper implements Runnable {
       this.threadContext = threadContext;
       this.replayWorkflowContext = replayWorkflowContext;
       this.name = name;
-      this.cancellationScope = new CancellationScopeImpl(detached, runnable, parent);
+      boolean deterministicCancellationScopeOrder =
+          replayWorkflowContext.checkSdkFlag(SdkFlag.DETERMINISTIC_CANCELLATION_SCOPE_ORDER);
+      this.cancellationScope =
+          new CancellationScopeImpl(
+              detached, deterministicCancellationScopeOrder, runnable, parent);
       Preconditions.checkState(
           context.getStatus() == Status.CREATED, "threadContext not in CREATED state");
       this.contextPropagators = contextPropagators;

temporal-sdk/src/test/java/io/temporal/workflow/cancellationTests/WorkflowCancellationScopeDeterminism.java

@@ -0,0 +1,104 @@
+package io.temporal.workflow.cancellationTests;
+
+import io.temporal.activity.ActivityInterface;
+import io.temporal.activity.ActivityOptions;
+import io.temporal.client.WorkflowClient;
+import io.temporal.client.WorkflowStub;
+import io.temporal.common.WorkflowExecutionHistory;
+import io.temporal.internal.common.SdkFlag;
+import io.temporal.internal.statemachines.WorkflowStateMachines;
+import io.temporal.testing.WorkflowReplayer;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.workflow.*;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+public class WorkflowCancellationScopeDeterminism {
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(TestWorkflowImpl.class)
+          .setActivityImplementations(new TestActivityImpl())
+          .build();
+
+  @Before
+  public void setUp() {
+    WorkflowStateMachines.initialFlags =
+        Collections.unmodifiableList(Arrays.asList(SdkFlag.DETERMINISTIC_CANCELLATION_SCOPE_ORDER));
+  }
+
+  @Test(timeout = 60000)
+  public void replayCanceledWorkflow() throws Exception {
+    for (int i = 0; i < 100; i++) {
+      TestWorkflow testWorkflow = testWorkflowRule.newWorkflowStub(TestWorkflow.class);
+
+      WorkflowClient.start(testWorkflow::start);
+
+      WorkflowStub stub = WorkflowStub.fromTyped(testWorkflow);
+      stub.cancel();
+      try {
+        stub.getResult(Void.class);
+      } catch (Exception e) {
+        // ignore; just blocking to make sure workflow is actually finished
+      }
+
+      WorkflowExecutionHistory history =
+          testWorkflowRule
+              .getWorkflowClient()
+              .fetchHistory(stub.getExecution().getWorkflowId(), stub.getExecution().getRunId());
+      WorkflowReplayer.replayWorkflowExecution(history, testWorkflowRule.getWorker());
+    }
+  }
+
+  @Test
+  public void replayTest() throws Exception {
+    WorkflowReplayer.replayWorkflowExecutionFromResource(
+        "cancellationScopeDeterminism.json", TestWorkflowImpl.class);
+  }
+
+  @WorkflowInterface
+  public interface TestWorkflow {
+    @WorkflowMethod
+    void start();
+  }
+
+  @ActivityInterface
+  public interface TestActivity {
+    void doActivity();
+  }
+
+  public static class TestActivityImpl implements TestActivity {
+    @Override
+    public void doActivity() {
+      try {
+        Thread.sleep(5000);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+
+  public static class TestWorkflowImpl implements TestWorkflow {
+
+    TestActivity activity =
+        Workflow.newActivityStub(
+            TestActivity.class,
+            ActivityOptions.newBuilder().setScheduleToCloseTimeout(Duration.ofSeconds(60)).build());
+
+    @Override
+    public void start() {
+      CancellationScope scope = Workflow.newCancellationScope(() -> activity.doActivity());
+
+      Async.procedure(
+          () -> {
+            Workflow.sleep(Duration.ofMinutes(5));
+          });
+
+      scope.run();
+    }
+  }
+}