Fix UpdateWithStart workflow args (#2286)

temporalio · Oct 22, 2024 · 7bcade2 · 7bcade2
1 parent 0ce1d6e
commit 7bcade2
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/temporal-sdk/src/main/java/io/temporal/client/UpdateWithStartWorkflowOperation.java b/temporal-sdk/src/main/java/io/temporal/client/UpdateWithStartWorkflowOperation.java
@@ -276,8 +276,12 @@ public static <R> Builder<R> newBuilder(String updateName, Class<R> resultClass,
 
   private UpdateOptions<R> options;
 
+  // set by constructor (untyped) or `prepareUpdate` (typed)
   private Object[] updateArgs;
 
+  // set by `prepareStart`
+  private Object[] workflowArgs;
+
   private final CompletableFuture<WorkflowUpdateHandle<R>> handle;
 
   private final Functions.Proc request;
@@ -294,9 +298,13 @@ WorkflowUpdateHandle<R> invoke(Functions.Proc workflow) {
     WorkflowInvocationHandler.initAsyncInvocation(
         WorkflowInvocationHandler.InvocationType.UPDATE_WITH_START, this);
     try {
+      // invokes `prepareUpdate` via WorkflowInvocationHandler.UpdateWithStartInvocationHandler
       request.apply();
+
+      // invokes `prepareStart` via WorkflowInvocationHandler.UpdateWithStartInvocationHandler
       workflow.apply();
-      stub.updateWithStart(this, this.updateArgs);
+
+      stub.updateWithStart(this, this.workflowArgs);
       return this.handle.get();
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
@@ -311,6 +319,7 @@ WorkflowUpdateHandle<R> invoke(Functions.Proc workflow) {
     }
   }
 
+  /** Invoked by {@link WorkflowInvocationHandler.UpdateWithStartInvocationHandler}. */
   void prepareUpdate(
       WorkflowStub stub, String updateName, Class resultClass, Type resultType, Object[] args) {
     setStub(stub);
@@ -323,8 +332,10 @@ void prepareUpdate(
             .build();
   }
 
-  void prepareStart(WorkflowStub stub) {
+  /** Invoked by {@link WorkflowInvocationHandler.UpdateWithStartInvocationHandler}. */
+  void prepareStart(WorkflowStub stub, Object[] args) {
     setStub(stub);
+    this.workflowArgs = args;
   }
 
   /** Returns the result of the update request. */

diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowInvocationHandler.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowInvocationHandler.java
@@ -490,7 +490,7 @@ public void invoke(
           throw new IllegalArgumentException(
               "Method '" + method.getName() + "' is not a WorkflowMethod");
         }
-        this.operation.prepareStart(untyped);
+        this.operation.prepareStart(untyped, args);
         state = State.UPDATE_RECEIVED;
       } else {
         throw new IllegalArgumentException(

diff --git a/temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateWithStartTest.java b/temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateWithStartTest.java
@@ -218,12 +218,19 @@ public void startVariousFuncs() throws ExecutionException, InterruptedException
         WorkflowClient.updateWithStart(stubF6::func6, "1", 2, 3, 4, 5, 6, updateOp6);
 
     assertEquals("0", handle0.getResultAsync().get());
+    assertEquals("func", WorkflowStub.fromTyped(stubF).getResult(String.class));
     assertEquals("1", handle1.getResultAsync().get());
+    assertEquals("1", WorkflowStub.fromTyped(stubF1).getResult(String.class));
     assertEquals("2", handle2.getResultAsync().get());
+    assertEquals("12", WorkflowStub.fromTyped(stubF2).getResult(String.class));
     assertEquals("3", handle3.getResultAsync().get());
+    assertEquals("123", WorkflowStub.fromTyped(stubF3).getResult(String.class));
     assertEquals("4", handle4.getResultAsync().get());
+    assertEquals("1234", WorkflowStub.fromTyped(stubF4).getResult(String.class));
     assertEquals("5", handle5.getResultAsync().get());
+    assertEquals("12345", WorkflowStub.fromTyped(stubF5).getResult(String.class));
     assertEquals("6", handle6.getResultAsync().get());
+    assertEquals("123456", WorkflowStub.fromTyped(stubF6).getResult(String.class));
   }
 
   @Test