Move workflow update polling inside of interceptor (#2159)

temporalio · Jul 30, 2024 · bbf2de7 · bbf2de7
1 parent 6b39e44
commit bbf2de7
Show file tree

Hide file tree

Showing 9 changed files with 115 additions and 91 deletions.
diff --git a/.github/workflows/features.yml b/.github/workflows/features.yml
@@ -7,4 +7,5 @@ jobs:
     with:
       java-repo-path: ${{github.event.pull_request.head.repo.full_name}}
       version: ${{github.event.pull_request.head.ref}}
-      version-is-repo-ref: true
+      version-is-repo-ref: true
+      features-repo-ref: java-update-iceptor-change
diff --git a/...main/java/io/temporal/opentracing/internal/OpenTracingWorkflowClientCallsInterceptor.java b/...main/java/io/temporal/opentracing/internal/OpenTracingWorkflowClientCallsInterceptor.java
@@ -23,6 +23,7 @@
 import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.Tracer;
+import io.temporal.client.UpdateHandle;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptor;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptorBase;
 import io.temporal.opentracing.OpenTracingOptions;
@@ -119,7 +120,7 @@ public <R> QueryOutput<R> query(QueryInput<R> input) {
   }
 
   @Override
-  public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
+  public <R> UpdateHandle<R> startUpdate(StartUpdateInput<R> input) {
     Span workflowStartUpdateSpan =
         contextAccessor.writeSpanContextToHeader(
             () ->

diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowStubImpl.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowStubImpl.java
@@ -31,6 +31,7 @@
 import io.temporal.common.interceptors.Header;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptor;
 import io.temporal.failure.CanceledFailure;
+import io.temporal.internal.client.LazyUpdateHandleImpl;
 import io.temporal.serviceclient.CheckedExceptionWrapper;
 import io.temporal.serviceclient.StatusUtils;
 import java.lang.reflect.Type;
@@ -331,42 +332,20 @@ public <R> UpdateHandle<R> startUpdate(UpdateOptions<R> options, Object... args)
     options.validate();
     WorkflowExecution targetExecution = execution.get();
     try {
-      WorkflowClientCallsInterceptor.StartUpdateOutput<R> result =
-          workflowClientInvoker.startUpdate(
-              new WorkflowClientCallsInterceptor.StartUpdateInput<>(
-                  targetExecution,
-                  options.getUpdateName(),
-                  Header.empty(),
-                  options.getUpdateId(),
-                  args,
-                  options.getResultClass(),
-                  options.getResultType(),
-                  options.getFirstExecutionRunId(),
-                  WaitPolicy.newBuilder()
-                      .setLifecycleStage(options.getWaitForStage().getProto())
-                      .build()));
-
-      if (result.hasResult()) {
-        return new CompletedUpdateHandleImpl<>(
-            result.getReference().getUpdateId(),
-            result.getReference().getWorkflowExecution(),
-            result.getResult());
-      } else {
-        LazyUpdateHandleImpl<R> handle =
-            new LazyUpdateHandleImpl<>(
-                workflowClientInvoker,
-                workflowType.orElse(null),
-                options.getUpdateName(),
-                result.getReference().getUpdateId(),
-                result.getReference().getWorkflowExecution(),
-                options.getResultClass(),
-                options.getResultType());
-        if (options.getWaitForStage() == WorkflowUpdateStage.COMPLETED) {
-          // Don't return the handle until completed, since that's what's been asked for
-          handle.waitCompleted();
-        }
-        return handle;
-      }
+      return workflowClientInvoker.startUpdate(
+          new WorkflowClientCallsInterceptor.StartUpdateInput<>(
+              targetExecution,
+              workflowType,
+              options.getUpdateName(),
+              Header.empty(),
+              options.getUpdateId(),
+              args,
+              options.getResultClass(),
+              options.getResultType(),
+              options.getFirstExecutionRunId(),
+              WaitPolicy.newBuilder()
+                  .setLifecycleStage(options.getWaitForStage().getProto())
+                  .build()));
     } catch (Exception e) {
       Throwable throwable = throwAsWorkflowFailureException(e, targetExecution);
       throw new WorkflowServiceException(targetExecution, workflowType.orElse(null), throwable);

diff --git a/...ral-sdk/src/main/java/io/temporal/common/interceptors/WorkflowClientCallsInterceptor.java b/...ral-sdk/src/main/java/io/temporal/common/interceptors/WorkflowClientCallsInterceptor.java
@@ -22,8 +22,8 @@
 
 import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.enums.v1.WorkflowExecutionStatus;
-import io.temporal.api.update.v1.UpdateRef;
 import io.temporal.api.update.v1.WaitPolicy;
+import io.temporal.client.UpdateHandle;
 import io.temporal.client.WorkflowOptions;
 import io.temporal.common.Experimental;
 import java.lang.reflect.Type;
@@ -78,7 +78,7 @@ public interface WorkflowClientCallsInterceptor {
   <R> QueryOutput<R> query(QueryInput<R> input);
 
   @Experimental
-  <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input);
+  <R> UpdateHandle<R> startUpdate(StartUpdateInput<R> input);
 
   @Experimental
   <R> PollWorkflowUpdateOutput<R> pollWorkflowUpdate(PollWorkflowUpdateInput<R> input);
@@ -383,6 +383,7 @@ public WorkflowExecution getWorkflowExecution() {
   @Experimental
   final class StartUpdateInput<R> {
     private final WorkflowExecution workflowExecution;
+    private final Optional<String> workflowType;
     private final String updateName;
     private final Header header;
     private final Object[] arguments;
@@ -394,6 +395,7 @@ final class StartUpdateInput<R> {
 
     public StartUpdateInput(
         WorkflowExecution workflowExecution,
+        Optional<String> workflowType,
         String updateName,
         Header header,
         String updateId,
@@ -403,6 +405,7 @@ public StartUpdateInput(
         String firstExecutionRunId,
         WaitPolicy waitPolicy) {
       this.workflowExecution = workflowExecution;
+      this.workflowType = workflowType;
       this.header = header;
       this.updateId = updateId;
       this.updateName = updateName;
@@ -417,6 +420,10 @@ public WorkflowExecution getWorkflowExecution() {
       return workflowExecution;
     }
 
+    public Optional<String> getWorkflowType() {
+      return workflowType;
+    }
+
     public String getUpdateName() {
       return updateName;
     }
@@ -450,44 +457,6 @@ public WaitPolicy getWaitPolicy() {
     }
   }
 
-  @Experimental
-  final class UpdateOutput<R> {
-    private final R result;
-
-    public UpdateOutput(R result) {
-      this.result = result;
-    }
-
-    public R getResult() {
-      return result;
-    }
-  }
-
-  @Experimental
-  final class StartUpdateOutput<R> {
-    private final UpdateRef reference;
-    private final R result;
-    private final boolean hasResult;
-
-    public StartUpdateOutput(UpdateRef reference, boolean hasResult, R result) {
-      this.reference = reference;
-      this.result = result;
-      this.hasResult = hasResult;
-    }
-
-    public UpdateRef getReference() {
-      return reference;
-    }
-
-    public boolean hasResult() {
-      return hasResult;
-    }
-
-    public R getResult() {
-      return result;
-    }
-  }
-
   @Experimental
   final class PollWorkflowUpdateInput<R> {
     private final WorkflowExecution workflowExecution;

diff --git a/...sdk/src/main/java/io/temporal/common/interceptors/WorkflowClientCallsInterceptorBase.java b/...sdk/src/main/java/io/temporal/common/interceptors/WorkflowClientCallsInterceptorBase.java
@@ -20,6 +20,7 @@
 
 package io.temporal.common.interceptors;
 
+import io.temporal.client.UpdateHandle;
 import java.util.concurrent.TimeoutException;
 
 /** Convenience base class for {@link WorkflowClientCallsInterceptor} implementations. */
@@ -62,7 +63,7 @@ public <R> QueryOutput<R> query(QueryInput<R> input) {
   }
 
   @Override
-  public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
+  public <R> UpdateHandle<R> startUpdate(StartUpdateInput<R> input) {
     return next.startUpdate(input);
   }
 

diff --git a/...ral/client/CompletedUpdateHandleImpl.java → ...nal/client/CompletedUpdateHandleImpl.java b/...ral/client/CompletedUpdateHandleImpl.java → ...nal/client/CompletedUpdateHandleImpl.java
@@ -18,21 +18,22 @@
  * limitations under the License.
  */
 
-package io.temporal.client;
+package io.temporal.internal.client;
 
 import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.client.UpdateHandle;
 import io.temporal.common.Experimental;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 @Experimental
-final class CompletedUpdateHandleImpl<T> implements UpdateHandle<T> {
+public final class CompletedUpdateHandleImpl<T> implements UpdateHandle<T> {
 
   private final String id;
   private final WorkflowExecution execution;
   private final T result;
 
-  CompletedUpdateHandleImpl(String id, WorkflowExecution execution, T result) {
+  public CompletedUpdateHandleImpl(String id, WorkflowExecution execution, T result) {
     this.id = id;
     this.execution = execution;
     this.result = result;

diff --git a/...temporal/client/LazyUpdateHandleImpl.java → ...internal/client/LazyUpdateHandleImpl.java b/...temporal/client/LazyUpdateHandleImpl.java → ...internal/client/LazyUpdateHandleImpl.java
@@ -18,11 +18,14 @@
  * limitations under the License.
  */
 
-package io.temporal.client;
+package io.temporal.internal.client;
 
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.client.UpdateHandle;
+import io.temporal.client.WorkflowException;
+import io.temporal.client.WorkflowServiceException;
 import io.temporal.common.Experimental;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptor;
 import io.temporal.serviceclient.CheckedExceptionWrapper;
@@ -33,7 +36,7 @@
 import java.util.concurrent.TimeoutException;
 
 @Experimental
-final class LazyUpdateHandleImpl<T> implements UpdateHandle<T> {
+public final class LazyUpdateHandleImpl<T> implements UpdateHandle<T> {
 
   private final WorkflowClientCallsInterceptor workflowClientInvoker;
   private final String workflowType;
@@ -44,7 +47,7 @@ final class LazyUpdateHandleImpl<T> implements UpdateHandle<T> {
   private final Type resultType;
   private WorkflowClientCallsInterceptor.PollWorkflowUpdateOutput<T> waitCompletedPollCall;
 
-  LazyUpdateHandleImpl(
+  public LazyUpdateHandleImpl(
       WorkflowClientCallsInterceptor workflowClientInvoker,
       String workflowType,
       String updateName,

diff --git a/temporal-sdk/src/main/java/io/temporal/internal/client/RootWorkflowClientInvoker.java b/temporal-sdk/src/main/java/io/temporal/internal/client/RootWorkflowClientInvoker.java
@@ -31,8 +31,7 @@
 import io.temporal.api.query.v1.WorkflowQuery;
 import io.temporal.api.update.v1.*;
 import io.temporal.api.workflowservice.v1.*;
-import io.temporal.client.WorkflowClientOptions;
-import io.temporal.client.WorkflowUpdateException;
+import io.temporal.client.*;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptor;
 import io.temporal.internal.client.external.GenericWorkflowClient;
@@ -298,7 +297,7 @@ public <R> QueryOutput<R> query(QueryInput<R> input) {
   }
 
   @Override
-  public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
+  public <R> UpdateHandle<R> startUpdate(StartUpdateInput<R> input) {
     DataConverter dataConverterWithWorkflowContext =
         clientOptions
             .getDataConverter()
@@ -337,10 +336,11 @@ public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
     // Re-attempt the update until it is at least accepted, or passes the lifecycle stage specified
     // by the user.
     UpdateWorkflowExecutionResponse result;
+    UpdateWorkflowExecutionLifecycleStage waitForStage = input.getWaitPolicy().getLifecycleStage();
     do {
       Deadline pollTimeoutDeadline = Deadline.after(POLL_UPDATE_TIMEOUT_S, TimeUnit.SECONDS);
       result = genericClient.update(updateRequest, pollTimeoutDeadline);
-    } while (result.getStage().getNumber() < input.getWaitPolicy().getLifecycleStage().getNumber()
+    } while (result.getStage().getNumber() < waitForStage.getNumber()
         && result.getStage().getNumber()
             < UpdateWorkflowExecutionLifecycleStage
                 .UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ACCEPTED
@@ -356,7 +356,10 @@ public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
                   input.getResultClass(),
                   input.getResultType(),
                   dataConverterWithWorkflowContext);
-          return new StartUpdateOutput<R>(result.getUpdateRef(), true, resultValue);
+          return new CompletedUpdateHandleImpl<>(
+              result.getUpdateRef().getUpdateId(),
+              result.getUpdateRef().getWorkflowExecution(),
+              resultValue);
         case FAILURE:
           throw new WorkflowUpdateException(
               result.getUpdateRef().getWorkflowExecution(),
@@ -370,7 +373,20 @@ public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
                   + result.getOutcome().getValueCase());
       }
     } else {
-      return new StartUpdateOutput<R>(result.getUpdateRef(), false, null);
+      LazyUpdateHandleImpl<R> handle =
+          new LazyUpdateHandleImpl<>(
+              this,
+              input.getWorkflowType().orElse(null),
+              input.getUpdateName(),
+              result.getUpdateRef().getUpdateId(),
+              result.getUpdateRef().getWorkflowExecution(),
+              input.getResultClass(),
+              input.getResultType());
+      if (waitForStage == WorkflowUpdateStage.COMPLETED.getProto()) {
+        // Don't return the handle until completed, since that's what's been asked for
+        handle.waitCompleted();
+      }
+      return handle;
     }
   }