Don't return update handles until desired stage reached

temporal-sdk/src/main/java/io/temporal/client/LazyUpdateHandleImpl.java

@@ -42,6 +42,7 @@ final class LazyUpdateHandleImpl<T> implements UpdateHandle<T> {
   private final WorkflowExecution execution;
   private final Class<T> resultClass;
   private final Type resultType;
+  private WorkflowClientCallsInterceptor.PollWorkflowUpdateOutput<T> previousPollCall;
 
   LazyUpdateHandleImpl(
       WorkflowClientCallsInterceptor workflowClientInvoker,
@@ -72,12 +73,15 @@ public String getId() {
 
   @Override
   public CompletableFuture<T> getResultAsync(long timeout, TimeUnit unit) {
-    WorkflowClientCallsInterceptor.PollWorkflowUpdateOutput output =
-        workflowClientInvoker.pollWorkflowUpdate(
-            new WorkflowClientCallsInterceptor.PollWorkflowUpdateInput<>(
-                execution, updateName, id, resultClass, resultType, timeout, unit));
+    WorkflowClientCallsInterceptor.PollWorkflowUpdateOutput<T> pollCall;
+    if (previousPollCall == null) {
+      pollUntilComplete(timeout, unit);
+    }
+    pollCall = previousPollCall;
+    // Get result may be called more than once for non-complete wait policies, so reset it here.
+    previousPollCall = null;
 
-    return output
+    return pollCall
         .getResult()
         .exceptionally(
             failure -> {
@@ -109,4 +113,11 @@ public CompletableFuture<T> getResultAsync(long timeout, TimeUnit unit) {
   public CompletableFuture<T> getResultAsync() {
     return this.getResultAsync(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
   }
+
+  void pollUntilComplete(long timeout, TimeUnit unit) {
+    previousPollCall =
+        workflowClientInvoker.pollWorkflowUpdate(
+            new WorkflowClientCallsInterceptor.PollWorkflowUpdateInput<>(
+                execution, updateName, id, resultClass, resultType, timeout, unit));
+  }
 }

temporal-sdk/src/main/java/io/temporal/client/UpdateWaitPolicy.java

@@ -23,7 +23,19 @@
 import io.temporal.api.enums.v1.UpdateWorkflowExecutionLifecycleStage;
 
 public enum UpdateWaitPolicy {
-  /** Update request waits for the update to be accepted by the workflow */
+  /**
+   * Update request waits for the update to be until the update request has been admitted by the
+   * server - it may be the case that due to a considerations like load or resource limits that an
+   * update is made to wait before the server will indicate that it has been received and will be
+   * processed. This value does not wait for any sort of acknowledgement from a worker.
+   */
+  ADMITTED(
+      UpdateWorkflowExecutionLifecycleStage.UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ADMITTED),
+
+  /**
+   * Update request waits for the update to be accepted (and validated, if there is a validator) by
+   * the workflow
+   */
   ACCEPTED(
       UpdateWorkflowExecutionLifecycleStage.UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ACCEPTED),
 

temporal-sdk/src/main/java/io/temporal/client/WorkflowStub.java

@@ -107,7 +107,8 @@ static <T> WorkflowStub fromTyped(T typed) {
 
   /**
    * Asynchronously update a workflow execution by invoking its update handler and returning a
-   * handle to the update request.
+   * handle to the update request. If {@link UpdateWaitPolicy#COMPLETED} is specified, in the
+   * options, the handle will not be returned until the update is completed.
    *
    * @param options options that will be used to configure and start a new update request.
    * @param args update method arguments

temporal-sdk/src/main/java/io/temporal/client/WorkflowStubImpl.java

@@ -351,14 +351,20 @@ public <R> UpdateHandle<R> startUpdate(UpdateOptions<R> options, Object... args)
             result.getReference().getWorkflowExecution(),
             result.getResult());
       } else {
-        return new LazyUpdateHandleImpl<>(
-            workflowClientInvoker,
-            workflowType.orElse(null),
-            options.getUpdateName(),
-            result.getReference().getUpdateId(),
-            result.getReference().getWorkflowExecution(),
-            options.getResultClass(),
-            options.getResultType());
+        LazyUpdateHandleImpl<R> handle =
+            new LazyUpdateHandleImpl<>(
+                workflowClientInvoker,
+                workflowType.orElse(null),
+                options.getUpdateName(),
+                result.getReference().getUpdateId(),
+                result.getReference().getWorkflowExecution(),
+                options.getResultClass(),
+                options.getResultType());
+        if (options.getWaitPolicy() == UpdateWaitPolicy.COMPLETED) {
+          // Don't return the handle until completed, since that's what's been asked for
+          handle.pollUntilComplete(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
+        }
+        return handle;
       }
     } catch (Exception e) {
       Throwable throwable = throwAsWorkflowFailureException(e, targetExecution);

temporal-sdk/src/main/java/io/temporal/internal/client/RootWorkflowClientInvoker.java

@@ -334,8 +334,17 @@ public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
             .setRequest(request)
             .build();
     Deadline pollTimeoutDeadline = Deadline.after(POLL_UPDATE_TIMEOUT_S, TimeUnit.SECONDS);
-    UpdateWorkflowExecutionResponse result =
-        genericClient.update(updateRequest, pollTimeoutDeadline);
+
+    // Re-attempt the update until it is at least accepted, or passes the lifecycle stage specified
+    // by the user.
+    UpdateWorkflowExecutionResponse result;
+    do {
+      result = genericClient.update(updateRequest, pollTimeoutDeadline);
+    } while (result.getStage().getNumber() < input.getWaitPolicy().getLifecycleStage().getNumber()
+        && result.getStage().getNumber()
+            < UpdateWorkflowExecutionLifecycleStage
+                .UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ACCEPTED
+                .getNumber());
 
     if (result.hasOutcome()) {
       switch (result.getOutcome().getValueCase()) {

temporal-sdk/src/test/java/io/temporal/workflow/shared/TestWorkflows.java

@@ -177,16 +177,6 @@ public interface QueryableWorkflow {
     void mySignal(String value);
   }
 
-  @WorkflowInterface
-  public interface SimpleWorkflowWithUpdate {
-
-    @WorkflowMethod
-    String execute();
-
-    @UpdateMethod
-    String update(String value);
-  }
-
   @WorkflowInterface
   public interface WorkflowWithUpdate {
 

temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateTest.java

@@ -40,6 +40,9 @@
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.Rule;
 import org.junit.Test;
 import org.slf4j.Logger;
@@ -53,7 +56,7 @@ public class UpdateTest {
   public SDKTestWorkflowRule testWorkflowRule =
       SDKTestWorkflowRule.newBuilder()
           .setWorkerOptions(WorkerOptions.newBuilder().build())
-          .setWorkflowTypes(TestUpdateWorkflowImpl.class)
+          .setWorkflowTypes(TestUpdateWorkflowImpl.class, TestWaitingUpdate.class)
           .setActivityImplementations(new ActivityImpl())
           .build();
 
@@ -107,15 +110,15 @@ public void testUpdateUntyped() throws ExecutionException, InterruptedException
             workflowType,
             SDKTestOptions.newWorkflowOptionsWithTimeouts(testWorkflowRule.getTaskQueue()));
 
-    WorkflowExecution execution = workflowStub.start();
+    workflowStub.start();
 
     SDKTestWorkflowRule.waitForOKQuery(workflowStub);
     assertEquals("initial", workflowStub.query("getState", String.class));
 
     // send an update through the sync path
     assertEquals("Execute-Hello", workflowStub.update("update", String.class, 0, "Hello"));
     // send an update through the async path
-    UpdateHandle updateRef = workflowStub.startUpdate("update", String.class, 0, "World");
+    UpdateHandle<String> updateRef = workflowStub.startUpdate("update", String.class, 0, "World");
     assertEquals("Execute-World", updateRef.getResultAsync().get());
     // send a bad update that will be rejected through the sync path
     assertThrows(
@@ -137,6 +140,48 @@ public void testUpdateUntyped() throws ExecutionException, InterruptedException
     assertEquals("Execute-Hello Execute-World", workflowStub.getResult(String.class));
   }
 
+  @Test
+  public void testUpdateHandleNotReturnedUntilCompleteWhenAsked()
+      throws ExecutionException, InterruptedException {
+    WorkflowClient workflowClient = testWorkflowRule.getWorkflowClient();
+    String workflowType = TestWorkflows.WorkflowWithUpdateAndSignal.class.getSimpleName();
+    WorkflowStub workflowStub =
+        workflowClient.newUntypedWorkflowStub(
+            workflowType,
+            SDKTestOptions.newWorkflowOptionsWithTimeouts(testWorkflowRule.getTaskQueue()));
+
+    workflowStub.start();
+
+    SDKTestWorkflowRule.waitForOKQuery(workflowStub);
+    assertEquals("initial", workflowStub.query("getState", String.class));
+
+    // Start the update but verify it does not return the handle until the update is complete
+    AtomicBoolean updateCompletedLast = new AtomicBoolean(false);
+    Future<?> asyncUpdate =
+        Executors.newSingleThreadExecutor()
+            .submit(
+                () -> {
+                  UpdateHandle<String> handle =
+                      workflowStub.startUpdate(
+                          UpdateOptions.newBuilder(String.class).setUpdateName("update").build(),
+                          "Enchi");
+                  updateCompletedLast.set(true);
+                  try {
+                    assertEquals("Enchi", handle.getResultAsync().get());
+                  } catch (Exception e) {
+                    throw new RuntimeException(e);
+                  }
+                });
+
+    workflowStub.signal("signal", "whatever");
+    updateCompletedLast.set(false);
+
+    asyncUpdate.get();
+    assertTrue(updateCompletedLast.get());
+    workflowStub.update("complete", void.class);
+    workflowStub.getResult(List.class);
+  }
+
   public static class TestUpdateWorkflowImpl implements WorkflowWithUpdate {
     String state = "initial";
     List<String> updates = new ArrayList<>();
@@ -203,4 +248,45 @@ public String execute(String input) {
       return Activity.getExecutionContext().getInfo().getActivityType() + "-" + input;
     }
   }
+
+  public static class TestWaitingUpdate implements TestWorkflows.WorkflowWithUpdateAndSignal {
+    String state = "initial";
+    List<String> updates = new ArrayList<>();
+    CompletablePromise<Void> signalled = Workflow.newPromise();
+    CompletablePromise<Void> promise = Workflow.newPromise();
+    private final TestActivities.TestActivity1 activity =
+        Workflow.newActivityStub(
+            TestActivities.TestActivity1.class,
+            ActivityOptions.newBuilder().setScheduleToCloseTimeout(Duration.ofHours(1)).build());
+
+    @Override
+    public List<String> execute() {
+      promise.get();
+      return updates;
+    }
+
+    @Override
+    public String getState() {
+      return state;
+    }
+
+    @Override
+    public void signal(String value) {
+      signalled.complete(null);
+    }
+
+    @Override
+    public String update(String value) {
+      Workflow.await(() -> signalled.isCompleted());
+      return value;
+    }
+
+    @Override
+    public void validator(String value) {}
+
+    @Override
+    public void complete() {
+      promise.complete(null);
+    }
+  }
 }

temporal-test-server/src/main/java/io/temporal/internal/testservice/StateMachines.java

@@ -66,12 +66,7 @@
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.common.v1.RetryPolicy;
 import io.temporal.api.common.v1.WorkflowExecution;
-import io.temporal.api.enums.v1.CancelExternalWorkflowExecutionFailedCause;
-import io.temporal.api.enums.v1.EventType;
-import io.temporal.api.enums.v1.RetryState;
-import io.temporal.api.enums.v1.SignalExternalWorkflowExecutionFailedCause;
-import io.temporal.api.enums.v1.StartChildWorkflowExecutionFailedCause;
-import io.temporal.api.enums.v1.TimeoutType;
+import io.temporal.api.enums.v1.*;
 import io.temporal.api.errordetails.v1.QueryFailedFailure;
 import io.temporal.api.failure.v1.ApplicationFailureInfo;
 import io.temporal.api.failure.v1.Failure;
@@ -1810,6 +1805,9 @@ private static void acceptUpdate(
                   UpdateRef.newBuilder()
                       .setWorkflowExecution(ctx.getExecution())
                       .setUpdateId(data.id))
+              .setStage(
+                  UpdateWorkflowExecutionLifecycleStage
+                      .UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ACCEPTED)
               .build();
 
       data.acceptance.complete(response);
@@ -1852,6 +1850,9 @@ private static void completeUpdate(
                       .setWorkflowExecution(ctx.getExecution())
                       .setUpdateId(data.id))
               .setOutcome(response.getOutcome())
+              .setStage(
+                  UpdateWorkflowExecutionLifecycleStage
+                      .UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_COMPLETED)
               .build();
 
       data.complete.complete(updateResponse);

temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java

@@ -1625,6 +1625,9 @@ private void processRejectionMessage(
                     .setUpdateId(rejection.getRejectedRequest().getMeta().getUpdateId())
                     .setWorkflowExecution(ctx.getExecution()))
             .setOutcome(Outcome.newBuilder().setFailure(rejection.getFailure()).build())
+            .setStage(
+                UpdateWorkflowExecutionLifecycleStage
+                    .UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_COMPLETED)
             .build();
     u.getAcceptance().complete(response);
   }
@@ -2157,7 +2160,13 @@ public PollWorkflowExecutionUpdateResponse pollUpdateWorkflowExecution(
           .setOutcome(completionResponse.getOutcome())
           .build();
     } catch (TimeoutException e) {
-      return PollWorkflowExecutionUpdateResponse.getDefaultInstance();
+      PollWorkflowExecutionUpdateResponse resp =
+          PollWorkflowExecutionUpdateResponse.getDefaultInstance();
+      return resp.toBuilder()
+          .setStage(
+              UpdateWorkflowExecutionLifecycleStage
+                  .UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ADMITTED)
+          .build();
     } catch (ExecutionException e) {
       Throwable cause = e.getCause();
       if (cause instanceof StatusRuntimeException) {