💥 Preserve non-cancellable timed promise waits during cancellation

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

@@ -36,6 +36,17 @@ public enum SdkFlag {
    * Bug: https://github.com/temporalio/sdk-java/issues/2796
    */
   VERSION_WAIT_FOR_MARKER(5),
+  /*
+   * Prevents cancellation of the enclosing CancellationScope from canceling the timeout timer
+   * used by non-cancellable Promise.get(timeout, unit).
+   *
+   * Introduced: 1.39.0
+   *
+   * Enabled: (pending)
+   *
+   * Bug: https://github.com/temporalio/sdk-java/issues/3026
+   */
+  DETACH_NON_CANCELLABLE_PROMISE_GET_TIMER(6),
   UNKNOWN(Integer.MAX_VALUE);
 
   private final int value;

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

@@ -1,6 +1,7 @@
 package io.temporal.internal.sync;
 
 import io.temporal.failure.TemporalFailure;
+import io.temporal.internal.common.SdkFlag;
 import io.temporal.workflow.CancellationScope;
 import io.temporal.workflow.CompletablePromise;
 import io.temporal.workflow.Functions;
@@ -91,15 +92,28 @@ public V cancellableGet(long timeout, TimeUnit unit) throws TimeoutException {
   public V cancellableGetImpl(boolean cancellable, long timeout, TimeUnit unit)
       throws TimeoutException {
     if (!completed) {
-      WorkflowInternal.await(
-          Duration.ofMillis(unit.toMillis(timeout)),
-          "Feature.get",
-          () -> {
-            if (cancellable) {
-              CancellationScope.throwCanceled();
-            }
-            return completed;
-          });
+      Runnable waitForCompletion =
+          () ->
+              WorkflowInternal.await(
+                  Duration.ofMillis(unit.toMillis(timeout)),
+                  "Feature.get",
+                  () -> {
+                    if (cancellable) {
+                      CancellationScope.throwCanceled();
+                    }
+                    return completed;
+                  });
+      if (!cancellable
+          && WorkflowInternal.checkSdkFlag(SdkFlag.DETACH_NON_CANCELLABLE_PROMISE_GET_TIMER)) {
+        CancellationScope timeoutScope = Workflow.newDetachedCancellationScope(waitForCompletion);
+        try {
+          timeoutScope.run();
+        } finally {
+          timeoutScope.cancel("Promise.get completed");
+        }
+      } else {
+        waitForCompletion.run();
+      }
     }
     if (!completed) {
       throw new TimeoutException();

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

@@ -875,6 +875,10 @@ static WorkflowOutboundCallsInterceptor getWorkflowOutboundInterceptor() {
     return getRootWorkflowContext().getWorkflowOutboundInterceptor();
   }
 
+  static boolean checkSdkFlag(SdkFlag flag) {
+    return getRootWorkflowContext().getReplayContext().checkSdkFlag(flag);
+  }
+
   static SyncWorkflowContext getRootWorkflowContext() {
     // If we are in a query handler, we need to get the workflow context from the
     // QueryDispatcher, otherwise we get it from the current thread's internal context.

temporal-sdk/src/test/java/io/temporal/workflow/PromiseTimedGetCancellationReplayTest.java

@@ -0,0 +1,81 @@
+package io.temporal.workflow;
+
+import io.temporal.testing.WorkflowReplayer;
+import java.util.Arrays;
+import java.util.Collection;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class PromiseTimedGetCancellationReplayTest {
+
+  @Parameterized.Parameters(name = "{0}")
+  public static Collection<Object[]> parameters() {
+    return Arrays.asList(
+        new Object[][] {
+          {
+            "promiseGetTimeoutCancellationLegacyAwaitLegacy.json",
+            PromiseTimedGetCancellationTest.ReplayWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationLegacyAwaitCancel.json",
+            PromiseTimedGetCancellationTest.ReplayWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationDetachedAwaitLegacy.json",
+            PromiseTimedGetCancellationTest.ReplayWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationDetachedAwaitCancel.json",
+            PromiseTimedGetCancellationTest.ReplayWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationCompletionAwaitLegacy.json",
+            PromiseTimedGetCancellationTest.ControlledPromiseWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationCompletionAwaitCancel.json",
+            PromiseTimedGetCancellationTest.ControlledPromiseWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationChildAwaitLegacy.json",
+            PromiseTimedGetCancellationTest.ChildStartCancellationWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationChildAwaitCancel.json",
+            PromiseTimedGetCancellationTest.ChildStartCancellationWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationActivityLegacyAwaitLegacy.json",
+            PromiseTimedGetCancellationTest.ActivityCancellationWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationActivityLegacyAwaitCancel.json",
+            PromiseTimedGetCancellationTest.ActivityCancellationWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationActivityDetachedAwaitLegacy.json",
+            PromiseTimedGetCancellationTest.ActivityCancellationWorkflowImpl.class
+          },
+          {
+            "promiseGetTimeoutCancellationActivityDetachedAwaitCancel.json",
+            PromiseTimedGetCancellationTest.ActivityCancellationWorkflowImpl.class
+          }
+        });
+  }
+
+  private final String historyResource;
+  private final Class<?> workflowImplementation;
+
+  public PromiseTimedGetCancellationReplayTest(
+      String historyResource, Class<?> workflowImplementation) {
+    this.historyResource = historyResource;
+    this.workflowImplementation = workflowImplementation;
+  }
+
+  @Test
+  public void replaysHistoryForFlagCombination() throws Exception {
+    WorkflowReplayer.replayWorkflowExecutionFromResource(historyResource, workflowImplementation);
+  }
+}