Upgrade waitForCompletionState to wait for all executions (#234)

src/main/java/io/iworkflow/core/Client.java

@@ -138,6 +138,7 @@ public String startWorkflow(
             unregisterWorkflowOptions.workflowRetryPolicy(options.getWorkflowRetryPolicy());
             unregisterWorkflowOptions.workflowConfigOverride(options.getWorkflowConfigOverride());
             unregisterWorkflowOptions.waitForCompletionStateExecutionIds(options.getWaitForCompletionStateExecutionIds());
+            unregisterWorkflowOptions.waitForCompletionStateIds(options.getWaitForCompletionStateIds());
 
             final Map<String, SearchAttributeValueType> saTypes = registry.getSearchAttributeKeyToTypeMap(wfType);
             final List<SearchAttribute> convertedSAs = convertToSearchAttributeList(saTypes, options.getInitialSearchAttribute());

src/main/java/io/iworkflow/core/UnregisteredClient.java

@@ -174,6 +174,7 @@ public String startWorkflow(
             request.workflowStartOptions(startOptions);
 
             request.waitForCompletionStateExecutionIds(options.getWaitForCompletionStateExecutionIds());
+            request.waitForCompletionStateIds(options.getWaitForCompletionStateIds());
         }
 
         try {

src/main/java/io/iworkflow/core/UnregisteredWorkflowOptions.java

@@ -29,4 +29,6 @@ public abstract class UnregisteredWorkflowOptions {
     public abstract Optional<Boolean> getUsingMemoForDataAttributes();
 
     public abstract List<String> getWaitForCompletionStateExecutionIds();
+
+    public abstract List<String> getWaitForCompletionStateIds();
 }

src/main/java/io/iworkflow/core/WorkflowOptionBuilderExtension.java

@@ -10,7 +10,7 @@ public class WorkflowOptionBuilderExtension {
     private ImmutableWorkflowOptions.Builder builder = ImmutableWorkflowOptions.builder();
 
     /**
-     * Add a state to wait for completion. This only waiting for the first completion of the state
+     * Add a state to wait for completion. This only waiting for all the completion of the state executions
      * NOTE: this will not be needed/required once server implements <a href="https://github.com/indeedeng/iwf/issues/349">this</a>
      * @param state The state to wait for completion.
      * @return The builder.
@@ -21,22 +21,22 @@ public WorkflowOptionBuilderExtension waitForCompletionState(Class<? extends Wor
     }
 
     /**
-     * Add states to wait for completion. This only waiting for the first completion of the state
+     * Add states to wait for completion. This only waiting for all the completion of the state executions
      * NOTE: this will not be needed/required once server implements <a href="https://github.com/indeedeng/iwf/issues/349">this</a>
      * @param states The states to wait for completion.
      * @return The builder.
      */
     @SafeVarargs
     public final WorkflowOptionBuilderExtension waitForCompletionStates(Class<? extends WorkflowState>... states) {
         Arrays.stream(states).forEach(
-                state -> builder.addWaitForCompletionStateExecutionIds(
-                        WorkflowState.getStateExecutionId(state,1)
+                state -> builder.addWaitForCompletionStateIds(
+                        WorkflowState.getDefaultStateId(state)
                 ));
         return this;
     }
 
     /**
-     * Add a state to wait for completion. This can wait for any times completion of the state
+     * Add a state to wait for completion. This can wait for a certain completion of the state execution
      * @param state The state to wait for completion.
      * @param number The number of the state completion to wait for. E.g. when it's 2, it's waiting for the second completion of the state.
      * @return The builder.

src/main/java/io/iworkflow/core/WorkflowOptions.java

@@ -23,6 +23,7 @@ public abstract class WorkflowOptions {
 
     public abstract Optional<WorkflowConfig> getWorkflowConfigOverride();
 
+    public abstract List<String> getWaitForCompletionStateIds();
     public abstract List<String> getWaitForCompletionStateExecutionIds();
 
     public static WorkflowOptionBuilderExtension extendedBuilder() {

src/main/java/io/iworkflow/core/WorkflowState.java

@@ -78,7 +78,7 @@ StateDecision execute(
      * @return the StateId of the state
      */
     default String getStateId() {
-        return this.getClass().getSimpleName();
+        return getDefaultStateId(this.getClass());
     }
 
     /**
@@ -116,6 +116,10 @@ static boolean shouldSkipWaitUntil(final WorkflowState state) {
     static String getStateExecutionId(Class<? extends WorkflowState> state, int number) {
         return String.format("%s-%d", state.getSimpleName(), number);
     }
+
+    static String getDefaultStateId(Class<? extends WorkflowState> state) {
+        return state.getSimpleName();
+    }
 }