Don't return update handles until desired stage reached

.github/workflows/features.yml

@@ -3,8 +3,9 @@ on: [push, pull_request]
 
 jobs:
   features-test:
-    uses: temporalio/features/.github/workflows/java.yaml@main
+    uses: temporalio/features/.github/workflows/java.yaml@java-breaking-update
     with:
       java-repo-path: ${{github.event.pull_request.head.repo.full_name}}
       version: ${{github.event.pull_request.head.ref}}
       version-is-repo-ref: true
+      features-repo-ref: java-breaking-update

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> waitCompletedPollCall;
 
   LazyUpdateHandleImpl(
       WorkflowClientCallsInterceptor workflowClientInvoker,
@@ -72,12 +73,23 @@ 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));
 
-    return output
+    WorkflowClientCallsInterceptor.PollWorkflowUpdateOutput<T> pollCall = null;
+    boolean setFromWaitCompleted = false;
+
+    // If waitCompleted was called, use the result from that call.
+    synchronized (this) {
+      if (waitCompletedPollCall != null) {
+        pollCall = waitCompletedPollCall;
+        waitCompletedPollCall = null;
+      }
+    }
+
+    if (!setFromWaitCompleted) {
+      pollCall = pollUntilComplete(timeout, unit);
+    }
+
+    return pollCall
         .getResult()
         .exceptionally(
             failure -> {
@@ -109,4 +121,17 @@ public CompletableFuture<T> getResultAsync(long timeout, TimeUnit unit) {
   public CompletableFuture<T> getResultAsync() {
     return this.getResultAsync(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
   }
+
+  // Can be called immediately after initialization to wait for the update to be completed, but
+  // still have the result be returned by getResultAsync.
+  void waitCompleted() {
+    waitCompletedPollCall = pollUntilComplete(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
+  }
+
+  private WorkflowClientCallsInterceptor.PollWorkflowUpdateOutput<T> pollUntilComplete(
+      long timeout, TimeUnit unit) {
+    return workflowClientInvoker.pollWorkflowUpdate(
+        new WorkflowClientCallsInterceptor.PollWorkflowUpdateInput<>(
+            execution, updateName, id, resultClass, resultType, timeout, unit));
+  }
 }

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

@@ -50,15 +50,15 @@ public static UpdateOptions getDefaultInstance() {
   private final String updateName;
   private final String updateId;
   private final String firstExecutionRunId;
-  private final UpdateWaitPolicy waitPolicy;
+  private final WorkflowUpdateStage waitPolicy;
   private final Class<T> resultClass;
   private final Type resultType;
 
   private UpdateOptions(
       String updateName,
       String updateId,
       String firstExecutionRunId,
-      UpdateWaitPolicy waitPolicy,
+      WorkflowUpdateStage waitPolicy,
       Class<T> resultClass,
       Type resultType) {
     this.updateName = updateName;
@@ -81,7 +81,7 @@ public String getFirstExecutionRunId() {
     return firstExecutionRunId;
   }
 
-  public UpdateWaitPolicy getWaitPolicy() {
+  public WorkflowUpdateStage getWaitPolicy() {
     return waitPolicy;
   }
 
@@ -152,7 +152,7 @@ public static final class Builder<T> {
     private String updateName;
     private String updateId;
     private String firstExecutionRunId;
-    private UpdateWaitPolicy waitPolicy;
+    private WorkflowUpdateStage waitPolicy;
     private Class<T> resultClass;
     private Type resultType;
 
@@ -208,7 +208,7 @@ public Builder<T> setFirstExecutionRunId(String firstExecutionRunId) {
      *   <li><b>Completed</b> Wait for the update to be completed by the workflow.
      * </ul>
      */
-    public Builder<T> setWaitPolicy(UpdateWaitPolicy waitPolicy) {
+    public Builder<T> setWaitPolicy(WorkflowUpdateStage waitPolicy) {
       this.waitPolicy = waitPolicy;
       return this;
     }
@@ -239,7 +239,7 @@ public UpdateOptions<T> build() {
           updateName,
           updateId,
           firstExecutionRunId == null ? "" : firstExecutionRunId,
-          waitPolicy == null ? UpdateWaitPolicy.ACCEPTED : waitPolicy,
+          waitPolicy == null ? WorkflowUpdateStage.ACCEPTED : waitPolicy,
           resultClass,
           resultType == null ? resultClass : resultType);
     }

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 WorkflowUpdateStage#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

@@ -297,7 +297,7 @@ public <R> R update(String updateName, Class<R> resultClass, Object... args) {
       UpdateOptions<R> options =
           UpdateOptions.<R>newBuilder()
               .setUpdateName(updateName)
-              .setWaitPolicy(UpdateWaitPolicy.COMPLETED)
+              .setWaitPolicy(WorkflowUpdateStage.COMPLETED)
               .setResultClass(resultClass)
               .build();
       return startUpdate(options, args).getResultAsync().get();
@@ -316,7 +316,7 @@ public <R> UpdateHandle<R> startUpdate(String updateName, Class<R> resultClass,
     UpdateOptions<R> options =
         UpdateOptions.<R>newBuilder()
             .setUpdateName(updateName)
-            .setWaitPolicy(UpdateWaitPolicy.ACCEPTED)
+            .setWaitPolicy(WorkflowUpdateStage.ACCEPTED)
             .setResultClass(resultClass)
             .setResultType(resultClass)
             .build();
@@ -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() == WorkflowUpdateStage.COMPLETED) {
+          // Don't return the handle until completed, since that's what's been asked for
+          handle.waitCompleted();
+        }
+        return handle;
       }
     } catch (Exception e) {
       Throwable throwable = throwAsWorkflowFailureException(e, targetExecution);

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

@@ -22,8 +22,20 @@
 
 import io.temporal.api.enums.v1.UpdateWorkflowExecutionLifecycleStage;
 
-public enum UpdateWaitPolicy {
-  /** Update request waits for the update to be accepted by the workflow */
+public enum WorkflowUpdateStage {
+  /**
+   * 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),
 
@@ -33,7 +45,7 @@ public enum UpdateWaitPolicy {
 
   private final UpdateWorkflowExecutionLifecycleStage policy;
 
-  UpdateWaitPolicy(UpdateWorkflowExecutionLifecycleStage policy) {
+  WorkflowUpdateStage(UpdateWorkflowExecutionLifecycleStage policy) {
     this.policy = policy;
   }
 

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

@@ -333,9 +333,18 @@ public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
             .setFirstExecutionRunId(input.getFirstExecutionRunId())
             .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 {
+      Deadline pollTimeoutDeadline = Deadline.after(POLL_UPDATE_TIMEOUT_S, TimeUnit.SECONDS);
+      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()) {
@@ -399,20 +408,18 @@ public <R> PollWorkflowUpdateOutput<R> pollWorkflowUpdate(PollWorkflowUpdateInpu
 
     Deadline pollTimeoutDeadline = Deadline.after(input.getTimeout(), input.getTimeoutUnit());
     pollWorkflowUpdateHelper(future, pollUpdateRequest, pollTimeoutDeadline);
-    return new PollWorkflowUpdateOutput(
+    return new PollWorkflowUpdateOutput<>(
         future.thenApply(
             (result) -> {
               if (result.hasOutcome()) {
                 switch (result.getOutcome().getValueCase()) {
                   case SUCCESS:
                     Optional<Payloads> updateResult = Optional.of(result.getOutcome().getSuccess());
-                    R resultValue =
-                        convertResultPayloads(
-                            updateResult,
-                            input.getResultClass(),
-                            input.getResultType(),
-                            dataConverterWithWorkflowContext);
-                    return resultValue;
+                    return convertResultPayloads(
+                        updateResult,
+                        input.getResultClass(),
+                        input.getResultType(),
+                        dataConverterWithWorkflowContext);
                   case FAILURE:
                     throw new WorkflowUpdateException(
                         input.getWorkflowExecution(),
@@ -434,31 +441,26 @@ private void pollWorkflowUpdateHelper(
       CompletableFuture<PollWorkflowExecutionUpdateResponse> resultCF,
       PollWorkflowExecutionUpdateRequest request,
       Deadline deadline) {
-
-    Deadline pollTimeoutDeadline =
-        Deadline.after(POLL_UPDATE_TIMEOUT_S, TimeUnit.SECONDS).minimum(deadline);
     genericClient
-        .pollUpdateAsync(request, pollTimeoutDeadline)
+        .pollUpdateAsync(request, deadline)
         .whenComplete(
             (r, e) -> {
+              if (e == null && !r.hasOutcome()) {
+                pollWorkflowUpdateHelper(resultCF, request, deadline);
+                return;
+              }
               if ((e instanceof StatusRuntimeException
                       && ((StatusRuntimeException) e).getStatus().getCode()
                           == Status.Code.DEADLINE_EXCEEDED)
-                  || pollTimeoutDeadline.isExpired()
-                  || (e == null && !r.hasOutcome())) {
-                // if the request has timed out, stop retrying
-                if (!deadline.isExpired()) {
-                  pollWorkflowUpdateHelper(resultCF, request, deadline);
-                } else {
-                  resultCF.completeExceptionally(
-                      new TimeoutException(
-                          "WorkflowId="
-                              + request.getUpdateRef().getWorkflowExecution().getWorkflowId()
-                              + ", runId="
-                              + request.getUpdateRef().getWorkflowExecution().getRunId()
-                              + ", updateId="
-                              + request.getUpdateRef().getUpdateId()));
-                }
+                  || deadline.isExpired()) {
+                resultCF.completeExceptionally(
+                    new TimeoutException(
+                        "WorkflowId="
+                            + request.getUpdateRef().getWorkflowExecution().getWorkflowId()
+                            + ", runId="
+                            + request.getUpdateRef().getWorkflowExecution().getRunId()
+                            + ", updateId="
+                            + request.getUpdateRef().getUpdateId()));
               } else if (e != null) {
                 resultCF.completeExceptionally(e);
               } else {

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 {