Support conditional complete on signal or internal empty (#178)

Author: longquanzheng

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

@@ -0,0 +1,15 @@
+package io.iworkflow.core;
+
+import io.iworkflow.gen.models.WorkflowConditionalCloseType;
+import org.immutables.value.Value;
+
+import java.util.Optional;
+
+@Value.Immutable
+public abstract class InternalConditionalClose {
+    public abstract WorkflowConditionalCloseType getWorkflowConditionalCloseType();
+
+    public abstract String getChannelName();
+
+    public abstract Optional<Object> getCloseInput();
+}

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

@@ -1,18 +1,22 @@
 package io.iworkflow.core;
 
+import io.iworkflow.gen.models.WorkflowConditionalCloseType;
 import org.immutables.value.Value;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Optional;
 
 @Value.Immutable
 public abstract class StateDecision {
 
+    public abstract Optional<InternalConditionalClose> getWorkflowConditionalClose();
+
     public abstract List<StateMovement> getNextStates();
 
     // a dead end will just complete its thread, without triggering any closing workflow
-    public static StateDecision deadEnd(){
+    public static StateDecision deadEnd() {
         return ImmutableStateDecision.builder()
                 .nextStates(Arrays.asList(StateMovement.DEAD_END_WORKFLOW_MOVEMENT))
                 .build();
@@ -22,9 +26,9 @@ public static ImmutableStateDecision.Builder builder() {
         return ImmutableStateDecision.builder();
     }
 
-    public static StateDecision gracefulCompleteWorkflow(final Object output) {
+    public static StateDecision gracefulCompleteWorkflow(final Object completionOutput) {
         return ImmutableStateDecision.builder().nextStates(Arrays.asList(
-                StateMovement.gracefulCompleteWorkflow(output)
+                StateMovement.gracefulCompleteWorkflow(completionOutput)
         )).build();
     }
 
@@ -34,9 +38,9 @@ public static StateDecision gracefulCompleteWorkflow() {
         )).build();
     }
 
-    public static StateDecision forceCompleteWorkflow(final Object output) {
+    public static StateDecision forceCompleteWorkflow(final Object completionOutput) {
         return ImmutableStateDecision.builder().nextStates(Arrays.asList(
-                StateMovement.forceCompleteWorkflow(output)
+                StateMovement.forceCompleteWorkflow(completionOutput)
         )).build();
     }
 
@@ -46,9 +50,9 @@ public static StateDecision forceCompleteWorkflow() {
         )).build();
     }
 
-    public static StateDecision forceFailWorkflow(final Object output) {
+    public static StateDecision forceFailWorkflow(final Object completionOutput) {
         return ImmutableStateDecision.builder().nextStates(Arrays.asList(
-                StateMovement.forceFailWorkflow(output)
+                StateMovement.forceFailWorkflow(completionOutput)
         )).build();
     }
 
@@ -58,6 +62,89 @@ public static StateDecision forceFailWorkflow() {
                 .build();
     }
 
+
+    public static StateDecision forceCompleteIfInternalChannelEmptyOrElse(final String internalChannelName, final Class<? extends WorkflowState> orElseStateClass) {
+        return forceCompleteIfInternalChannelEmptyOrElse(internalChannelName, orElseStateClass, null);
+    }
+
+    public static StateDecision forceCompleteIfInternalChannelEmptyOrElse(final String internalChannelName, final Class<? extends WorkflowState> orElseStateClass, final Object stateInput) {
+        return forceCompleteIfInternalChannelEmptyOrElse(null, internalChannelName, StateMovement.create(orElseStateClass, stateInput));
+    }
+
+    public static StateDecision forceCompleteIfInternalChannelEmptyOrElse(final Object completionOutput, final String internalChannelName, final Class<? extends WorkflowState> orElseStateClass) {
+        return forceCompleteIfInternalChannelEmptyOrElse(completionOutput, internalChannelName, orElseStateClass, null);
+    }
+
+    public static StateDecision forceCompleteIfInternalChannelEmptyOrElse(final Object completionOutput, final String internalChannelName, final Class<? extends WorkflowState> orElseStateClass, final Object stateInput) {
+        return forceCompleteIfInternalChannelEmptyOrElse(completionOutput, internalChannelName, StateMovement.create(orElseStateClass, stateInput));
+    }
+
+    /**
+     * Atomically force complete the workflow if internal channel is empty, otherwise trigger the state movements from the current thread
+     * This is important for use case that needs to ensure all the messages in the channel are processed before completing the workflow, otherwise messages will be lost.
+     * Without this atomic API, if user just check the channel emptiness in the State APIs, the channel may receive new messages during the execution of state APIs
+     * <br>
+     * Note that today this doesn't cover the case that internal messages are published from other State APIs yet. It's only for internal messages published from RPCs.
+     * If you do want to use other State APIs to publish messages to the channel at the same time, you can use persistence locking to ensure only the State APIs are not executed
+     * in parallel. See more in TODO https://github.com/indeedeng/iwf/issues/289
+     *
+     * @param completionOutput     the output of workflow completion
+     * @param internalChannelName  the internal channel name for checking emptiness
+     * @param orElseStateMovements the state movements if channel is not empty
+     * @return the decision
+     */
+    public static StateDecision forceCompleteIfInternalChannelEmptyOrElse(final Object completionOutput, final String internalChannelName, final StateMovement... orElseStateMovements) {
+        return ImmutableStateDecision.builder()
+                .workflowConditionalClose(
+                        ImmutableInternalConditionalClose.builder()
+                                .workflowConditionalCloseType(WorkflowConditionalCloseType.FORCE_COMPLETE_ON_INTERNAL_CHANNEL_EMPTY)
+                                .channelName(internalChannelName)
+                                .closeInput(Optional.ofNullable(completionOutput))
+                                .build()
+                )
+                .nextStates(Arrays.asList(orElseStateMovements))
+                .build();
+    }
+
+    public static StateDecision forceCompleteIfSignalChannelEmptyOrElse(final String signalChannelName, final Class<? extends WorkflowState> orElseStateClass) {
+        return forceCompleteIfSignalChannelEmptyOrElse(signalChannelName, orElseStateClass, null);
+    }
+
+    public static StateDecision forceCompleteIfSignalChannelEmptyOrElse(final String signalChannelName, final Class<? extends WorkflowState> orElseStateClass, final Object stateInput) {
+        return forceCompleteIfSignalChannelEmptyOrElse(null, signalChannelName, StateMovement.create(orElseStateClass, stateInput));
+    }
+
+    public static StateDecision forceCompleteIfSignalChannelEmptyOrElse(final Object completionOutput, final String signalChannelName, final Class<? extends WorkflowState> orElseStateClass) {
+        return forceCompleteIfSignalChannelEmptyOrElse(completionOutput, signalChannelName, orElseStateClass, null);
+    }
+
+    public static StateDecision forceCompleteIfSignalChannelEmptyOrElse(final Object completionOutput, final String signalChannelName, final Class<? extends WorkflowState> orElseStateClass, final Object stateInput) {
+        return forceCompleteIfSignalChannelEmptyOrElse(completionOutput, signalChannelName, StateMovement.create(orElseStateClass, stateInput));
+    }
+
+    /**
+     * Atomically force complete the workflow if signal channel is empty, otherwise trigger the state movements from the current thread
+     * This is important for use case that needs to ensure all the messages in the channel are processed before completing the workflow, otherwise messages will be lost.
+     * Without this atomic API, if user just check the channel emptiness in the State APIs, the channel may receive new messages during the execution of state APIs
+     *
+     * @param completionOutput     the output of workflow completion
+     * @param signalChannelName    the signal channel name for checking emptiness
+     * @param orElseStateMovements the state movements if channel is not empty
+     * @return the decision
+     */
+    public static StateDecision forceCompleteIfSignalChannelEmptyOrElse(final Object completionOutput, final String signalChannelName, final StateMovement... orElseStateMovements) {
+        return ImmutableStateDecision.builder()
+                .workflowConditionalClose(
+                        ImmutableInternalConditionalClose.builder()
+                                .workflowConditionalCloseType(WorkflowConditionalCloseType.FORCE_COMPLETE_ON_SIGNAL_CHANNEL_EMPTY)
+                                .channelName(signalChannelName)
+                                .closeInput(Optional.ofNullable(completionOutput))
+                                .build()
+                )
+                .nextStates(Arrays.asList(orElseStateMovements))
+                .build();
+    }
+
     public static StateDecision singleNextState(final Class<? extends WorkflowState> stateClass) {
         return singleNextState(stateClass.getSimpleName());
     }

src/main/java/io/iworkflow/core/mapper/StateDecisionMapper.java

@@ -1,22 +1,43 @@
 package io.iworkflow.core.mapper;
 
+import io.iworkflow.core.InternalConditionalClose;
 import io.iworkflow.core.ObjectEncoder;
 import io.iworkflow.core.Registry;
+import io.iworkflow.gen.models.EncodedObject;
 import io.iworkflow.gen.models.StateDecision;
+import io.iworkflow.gen.models.WorkflowConditionalClose;
 
 import java.util.stream.Collectors;
 
 public class StateDecisionMapper {
     public static StateDecision toGenerated(io.iworkflow.core.StateDecision stateDecision, final String workflowType, final Registry registry, final ObjectEncoder objectEncoder) {
-        if (stateDecision.getNextStates() == null) {
+        if (stateDecision.getNextStates() == null && !stateDecision.getWorkflowConditionalClose().isPresent()) {
             return null;
         }
-        return new StateDecision()
-                .nextStates(stateDecision.getNextStates()
-                        .stream()
-                        .map(movement -> {
-                            return StateMovementMapper.toGenerated(movement, workflowType, registry, objectEncoder);
-                        })
-                        .collect(Collectors.toList()));
+
+        StateDecision decision = new StateDecision();
+
+        if (stateDecision.getNextStates() != null) {
+            decision.nextStates(
+                    stateDecision.getNextStates()
+                            .stream()
+                            .map(movement -> StateMovementMapper.toGenerated(movement, workflowType, registry, objectEncoder))
+                            .collect(Collectors.toList())
+            );
+        }
+
+        if (!stateDecision.getWorkflowConditionalClose().isPresent()) {
+            return decision;
+        }
+
+        InternalConditionalClose conditionalClose = stateDecision.getWorkflowConditionalClose().get();
+        EncodedObject closeInput = objectEncoder.encode(conditionalClose.getCloseInput());
+        decision.conditionalClose(
+                new WorkflowConditionalClose()
+                        .conditionalCloseType(conditionalClose.getWorkflowConditionalCloseType())
+                        .closeInput(closeInput)
+                        .channelName(conditionalClose.getChannelName())
+        );
+        return decision;
     }
 }

src/test/java/io/iworkflow/integ/ConditionalCompleteTest.java

@@ -0,0 +1,62 @@
+package io.iworkflow.integ;
+
+import io.iworkflow.core.Client;
+import io.iworkflow.core.ClientOptions;
+import io.iworkflow.integ.conditional.ConditionalCompleteWorkflow;
+import io.iworkflow.spring.TestSingletonWorkerService;
+import io.iworkflow.spring.controller.WorkflowRegistry;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.ExecutionException;
+
+public class ConditionalCompleteTest {
+
+    @BeforeEach
+    public void setup() throws ExecutionException, InterruptedException {
+        TestSingletonWorkerService.startWorkerIfNotUp();
+    }
+
+    @Test
+    public void testCompleteIfInternalChannelEmpty() throws InterruptedException {
+        testCompleteIfChannelEmpty(false);
+    }
+
+    @Test
+    public void testCompleteIfSignalChannelEmpty() throws InterruptedException {
+        testCompleteIfChannelEmpty(true);
+    }
+
+    public void testCompleteIfChannelEmpty(boolean useSignal) throws InterruptedException {
+        final Client client = new Client(WorkflowRegistry.registry, ClientOptions.localDefault);
+        String namePart;
+        if (useSignal) {
+            namePart = "Signal";
+        } else {
+            namePart = "Internal";
+        }
+        final String wfId = "testCompleteIf" + namePart + "ChannelEmpty" + System.currentTimeMillis() / 1000;
+        final String runId = client.startWorkflow(
+                ConditionalCompleteWorkflow.class, wfId, 10, useSignal);
+
+        Thread.sleep(1000);
+
+        for (int i = 0; i < 3; i++) {
+            if (useSignal) {
+                client.signalWorkflow(ConditionalCompleteWorkflow.class, wfId, "", ConditionalCompleteWorkflow.SIGNAL_CHANNEL_NAME, null);
+            } else {
+                final ConditionalCompleteWorkflow rpcStub = client.newRpcStub(ConditionalCompleteWorkflow.class, wfId, "");
+                client.invokeRPC(rpcStub::publishToInternalChannel);
+            }
+            if (i == 0) {
+                // wait for a second so that the workflow is in execute state
+                Thread.sleep(1000);
+            }
+        }
+
+        final Integer output = client.getSimpleWorkflowResultWithWait(Integer.class, wfId);
+        Assertions.assertEquals(3, output);
+
+    }
+}

src/test/java/io/iworkflow/integ/conditional/ConditionalCompleteWorkflow.java

@@ -0,0 +1,107 @@
+package io.iworkflow.integ.conditional;
+
+import io.iworkflow.core.*;
+import io.iworkflow.core.command.CommandRequest;
+import io.iworkflow.core.command.CommandResults;
+import io.iworkflow.core.communication.*;
+import io.iworkflow.core.persistence.DataAttributeDef;
+import io.iworkflow.core.persistence.Persistence;
+import io.iworkflow.core.persistence.PersistenceFieldDef;
+import org.springframework.stereotype.Component;
+
+import java.util.Arrays;
+import java.util.List;
+
+
+@Component
+public class ConditionalCompleteWorkflow implements ObjectWorkflow {
+
+    public static final String SIGNAL_CHANNEL_NAME = "test-signal-channel";
+
+    public static final String INTERNAL_CHANNEL_NAME = "test-internal-channel";
+    public static final String DA_COUNTER = "counter";
+
+
+    @Override
+    public List<CommunicationMethodDef> getCommunicationSchema() {
+        return Arrays.asList(
+                SignalChannelDef.create(Void.class, SIGNAL_CHANNEL_NAME),
+                InternalChannelDef.create(Void.class, INTERNAL_CHANNEL_NAME)
+        );
+    }
+
+    @Override
+    public List<PersistenceFieldDef> getPersistenceSchema() {
+        return Arrays.asList(
+                DataAttributeDef.create(Integer.class, DA_COUNTER)
+        );
+    }
+
+    @Override
+    public List<StateDef> getWorkflowStates() {
+        return Arrays.asList(
+                StateDef.startingState(new WorkflowState1())
+        );
+    }
+
+    @RPC
+    public void publishToInternalChannel(Context context, Persistence persistence, Communication communication) {
+        communication.publishInternalChannel(INTERNAL_CHANNEL_NAME, null);
+    }
+}
+
+
+class WorkflowState1 implements WorkflowState<Boolean> {
+
+    @Override
+    public Class<Boolean> getInputType() {
+        return Boolean.class;
+    }
+
+    @Override
+    public CommandRequest waitUntil(
+            Context context,
+            Boolean useSignal,
+            Persistence persistence,
+            final Communication communication) {
+
+        if (useSignal) {
+            return CommandRequest.forAnyCommandCompleted(
+                    SignalCommand.create(ConditionalCompleteWorkflow.SIGNAL_CHANNEL_NAME)
+            );
+        } else {
+            return CommandRequest.forAnyCommandCompleted(
+                    InternalChannelCommand.create(ConditionalCompleteWorkflow.INTERNAL_CHANNEL_NAME)
+            );
+        }
+    }
+
+    @Override
+    public StateDecision execute(
+            Context context,
+            Boolean useSignal,
+            CommandResults commandResults,
+            Persistence persistence,
+            final Communication communication) {
+        Integer counter = persistence.getDataAttribute(ConditionalCompleteWorkflow.DA_COUNTER, Integer.class);
+        if (counter == null) {
+            counter = 0;
+        }
+        counter++;
+        persistence.setDataAttribute(ConditionalCompleteWorkflow.DA_COUNTER, counter);
+
+        if (context.getStateExecutionId().get().equals("WorkflowState1-1")) {
+            // wait for 3 seconds so that the channel can have a new message
+            try {
+                Thread.sleep(3000);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+        if (useSignal) {
+            return StateDecision.forceCompleteIfSignalChannelEmptyOrElse(counter, ConditionalCompleteWorkflow.SIGNAL_CHANNEL_NAME, WorkflowState1.class, useSignal);
+        } else {
+            return StateDecision.forceCompleteIfInternalChannelEmptyOrElse(counter, ConditionalCompleteWorkflow.INTERNAL_CHANNEL_NAME, WorkflowState1.class, useSignal);
+        }
+    }
+}