Add promise-based workflow await

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

@@ -9,6 +9,7 @@
 import java.time.Duration;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 /**
  * Contains support for asynchronous invocations. The basic idea is that any code is invoked in a
@@ -249,6 +250,14 @@ public static <A1, A2, A3, A4, A5, A6> Promise<Void> procedure(
     return procedure(isAsync(procedure), () -> procedure.apply(arg1, arg2, arg3, arg4, arg5, arg6));
   }
 
+  public static Promise<Void> await(Supplier<Boolean> unblockCondition) {
+    return procedure(false, () -> Workflow.await(unblockCondition));
+  }
+
+  public static Promise<Boolean> await(Duration timeout, Supplier<Boolean> unblockCondition) {
+    return execute(false, () -> Workflow.await(timeout, unblockCondition));
+  }
+
   public static <R> Promise<R> retry(
       RetryOptions options, Optional<Duration> expiration, Functions.Func<Promise<R>> fn) {
     return WorkflowRetryerInternal.retryAsync(options, expiration, fn);

temporal-sdk/src/main/java/io/temporal/workflow/Async.java

@@ -4,6 +4,7 @@
 import io.temporal.internal.sync.AsyncInternal;
 import java.time.Duration;
 import java.util.Optional;
+import java.util.function.Supplier;
 
 /** Supports invoking lambdas and activity and child workflow references asynchronously. */
 public final class Async {
@@ -217,6 +218,48 @@ public static <A1, A2, A3, A4, A5, A6> Promise<Void> procedure(
     return AsyncInternal.procedure(procedure, arg1, arg2, arg3, arg4, arg5, arg6);
   }
 
+  /**
+   * Returns a promise that completes with {@code null} when {@code unblockCondition} evaluates to
+   * {@code true}. Unlike {@link Workflow#await(Supplier)}, this method does not block the calling
+   * workflow thread.
+   *
+   * <p>The condition is evaluated on every workflow state transition. It must not call blocking
+   * operations or mutate workflow state. It must also not contain time-based conditions; use {@link
+   * #await(Duration, Supplier)} for those.
+   *
+   * <p>If the {@link CancellationScope} active when this method is invoked is canceled, the promise
+   * completes exceptionally with a {@link io.temporal.failure.CanceledFailure}. An exception thrown
+   * by the condition also completes the promise exceptionally.
+   *
+   * @param unblockCondition condition that completes the promise when satisfied.
+   * @return promise that completes when the condition is satisfied.
+   */
+  public static Promise<Void> await(Supplier<Boolean> unblockCondition) {
+    return AsyncInternal.await(unblockCondition);
+  }
+
+  /**
+   * Returns a promise that completes with {@code true} when {@code unblockCondition} evaluates to
+   * {@code true}, or with {@code false} when {@code timeout} expires. Unlike {@link
+   * Workflow#await(Duration, Supplier)}, this method does not block the calling workflow thread.
+   *
+   * <p>The condition is evaluated on every workflow state transition. It must not call blocking
+   * operations or mutate workflow state. It must also not contain time-based conditions; use the
+   * {@code timeout} parameter for those.
+   *
+   * <p>If the {@link CancellationScope} active when this method is invoked is canceled, the promise
+   * completes exceptionally with a {@link io.temporal.failure.CanceledFailure}. An exception thrown
+   * by the condition also completes the promise exceptionally.
+   *
+   * @param timeout time after which the promise completes with {@code false} if the condition is
+   *     not satisfied.
+   * @param unblockCondition condition that completes the promise with {@code true} when satisfied.
+   * @return promise that contains whether the condition was satisfied before the timeout.
+   */
+  public static Promise<Boolean> await(Duration timeout, Supplier<Boolean> unblockCondition) {
+    return AsyncInternal.await(timeout, unblockCondition);
+  }
+
   /**
    * Invokes function retrying in case of failures according to retry options. Asynchronous variant.
    * Use {@link Workflow#retry(RetryOptions, Optional, Functions.Func)} for synchronous functions.

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

@@ -0,0 +1,209 @@
+package io.temporal.workflow;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import io.temporal.activity.ActivityInterface;
+import io.temporal.activity.ActivityOptions;
+import io.temporal.client.WorkflowClient;
+import io.temporal.client.WorkflowStub;
+import io.temporal.failure.CanceledFailure;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import java.time.Duration;
+import java.util.function.Supplier;
+import org.junit.Rule;
+import org.junit.Test;
+
+public class AsyncAwaitTest {
+
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(TestAsyncAwaitWorkflowImpl.class)
+          .setActivityImplementations(new TestAwaitActivityImpl())
+          .build();
+
+  @Test
+  public void testAlreadySatisfiedConditionCompletesWithTrue() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+
+    assertEquals("true", workflow.execute("alreadySatisfied"));
+  }
+
+  @Test
+  public void testUnsatisfiedConditionCompletesWithFalseAfterTimeout() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+
+    assertEquals("false:true", workflow.execute("timeout"));
+  }
+
+  @Test
+  public void testAwaitDoesNotBlockCallingWorkflowThread() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+
+    assertEquals("true", workflow.execute("nonBlocking"));
+  }
+
+  @Test
+  public void testSignalConditionComposesWithActivityPromise() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+    WorkflowClient.start(workflow::execute, "signalAndActivity");
+
+    workflow.unblock();
+
+    assertEquals("true:activity", WorkflowStub.fromTyped(workflow).getResult(String.class));
+  }
+
+  @Test
+  public void testCancellationFailsPromise() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+
+    assertEquals("CanceledFailure", workflow.execute("cancellation"));
+  }
+
+  @Test
+  public void testPredicateExceptionFailsPromise() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+
+    assertEquals("IllegalStateException:predicate failed", workflow.execute("predicateFailure"));
+  }
+
+  @Test
+  public void testAlreadySatisfiedConditionWithoutTimeoutCompletesWithNull() {
+    assertEquals("null", newWorkflowStub().execute("alreadySatisfiedWithoutTimeout"));
+  }
+
+  @Test
+  public void testAwaitWithoutTimeoutDoesNotBlockCallingWorkflowThread() {
+    assertEquals("null", newWorkflowStub().execute("nonBlockingWithoutTimeout"));
+  }
+
+  @Test
+  public void testSignalConditionWithoutTimeoutComposesWithActivityPromise() {
+    TestAsyncAwaitWorkflow workflow = newWorkflowStub();
+    WorkflowClient.start(workflow::execute, "signalAndActivityWithoutTimeout");
+
+    workflow.unblock();
+
+    assertEquals("null:activity", WorkflowStub.fromTyped(workflow).getResult(String.class));
+  }
+
+  @Test
+  public void testCancellationWithoutTimeoutFailsPromise() {
+    assertEquals("CanceledFailure", newWorkflowStub().execute("cancellationWithoutTimeout"));
+  }
+
+  @Test
+  public void testPredicateExceptionWithoutTimeoutFailsPromise() {
+    assertEquals(
+        "IllegalStateException:predicate failed",
+        newWorkflowStub().execute("predicateFailureWithoutTimeout"));
+  }
+
+  private TestAsyncAwaitWorkflow newWorkflowStub() {
+    return testWorkflowRule.newWorkflowStubTimeoutOptions(TestAsyncAwaitWorkflow.class);
+  }
+
+  @WorkflowInterface
+  public interface TestAsyncAwaitWorkflow {
+
+    @WorkflowMethod
+    String execute(String testCase);
+
+    @SignalMethod
+    void unblock();
+  }
+
+  @ActivityInterface
+  public interface TestAwaitActivity {
+    String execute();
+  }
+
+  public static class TestAwaitActivityImpl implements TestAwaitActivity {
+
+    @Override
+    public String execute() {
+      return "activity";
+    }
+  }
+
+  public static class TestAsyncAwaitWorkflowImpl implements TestAsyncAwaitWorkflow {
+
+    private boolean unblocked;
+    private Promise<?> cancellationPromise;
+
+    @Override
+    public String execute(String testCase) {
+      switch (testCase) {
+        case "alreadySatisfied":
+          return Async.await(Duration.ofHours(1), () -> true).get().toString();
+        case "alreadySatisfiedWithoutTimeout":
+          Promise<Void> satisfied = Async.await(() -> true);
+          return String.valueOf(satisfied.get());
+        case "timeout":
+          long timeoutStart = Workflow.currentTimeMillis();
+          boolean result = Async.await(Duration.ofMinutes(1), () -> false).get();
+          return result
+              + ":"
+              + (Workflow.currentTimeMillis() - timeoutStart >= Duration.ofMinutes(1).toMillis());
+        case "nonBlocking":
+          long start = Workflow.currentTimeMillis();
+          Async.await(Duration.ofHours(1), () -> false);
+          return Boolean.toString(
+              Workflow.currentTimeMillis() - start < Duration.ofHours(1).toMillis());
+        case "nonBlockingWithoutTimeout":
+          Promise<Void> pending = Async.await(() -> unblocked);
+          unblocked = true;
+          return String.valueOf(pending.get());
+        case "signalAndActivity":
+        case "signalAndActivityWithoutTimeout":
+          Promise<?> condition =
+              testCase.equals("signalAndActivity")
+                  ? Async.await(Duration.ofHours(1), () -> unblocked)
+                  : Async.await(() -> unblocked);
+          TestAwaitActivity activity =
+              Workflow.newActivityStub(
+                  TestAwaitActivity.class,
+                  ActivityOptions.newBuilder()
+                      .setStartToCloseTimeout(Duration.ofSeconds(10))
+                      .build());
+          Promise<String> activityResult = Async.function(activity::execute);
+          Promise.allOf(condition, activityResult).get();
+          return condition.get() + ":" + activityResult.get();
+        case "cancellation":
+        case "cancellationWithoutTimeout":
+          CancellationScope scope =
+              Workflow.newCancellationScope(
+                  () ->
+                      cancellationPromise =
+                          testCase.equals("cancellation")
+                              ? Async.await(Duration.ofHours(1), () -> false)
+                              : Async.await(() -> false));
+          scope.run();
+          scope.cancel();
+          RuntimeException cancellationFailure = cancellationPromise.getFailure();
+          assertTrue(cancellationFailure instanceof CanceledFailure);
+          return cancellationFailure.getClass().getSimpleName();
+        case "predicateFailure":
+        case "predicateFailureWithoutTimeout":
+          Supplier<Boolean> predicate =
+              () -> {
+                throw new IllegalStateException("predicate failed");
+              };
+          Promise<?> failed =
+              testCase.equals("predicateFailure")
+                  ? Async.await(Duration.ofHours(1), predicate)
+                  : Async.await(predicate);
+          RuntimeException predicateFailure = failed.getFailure();
+          return predicateFailure.getClass().getSimpleName() + ":" + predicateFailure.getMessage();
+        default:
+          throw new IllegalArgumentException("Unknown test case: " + testCase);
+      }
+    }
+
+    @Override
+    public void unblock() {
+      unblocked = true;
+    }
+  }
+}