IWF-397: Add channel sizes data for RPC (#280)

longquanzheng · web-flow · commit fcc0e66bcdbf · 2024-12-11T10:54:13.000-08:00
diff --git a/iwf-idl b/iwf-idl
@@ -1 +1 @@
-Subproject commit a6fbd8feb427fdf77b7ddc0f7c9d7eb154c385e7
+Subproject commit e5bc6236d36b977c2fec85c2bf5e0742f859e04e
diff --git a/src/main/java/io/iworkflow/core/WorkerService.java b/src/main/java/io/iworkflow/core/WorkerService.java
@@ -70,7 +70,10 @@ public WorkflowWorkerRpcResponse handleWorkflowWorkerRpc(final WorkflowWorkerRpc
         final Map<String, SearchAttributeValueType> searchAttrsTypeMap = registry.getSearchAttributeKeyToTypeMap(req.getWorkflowType());
         final SearchAttributeRWImpl searchAttributeRW = new SearchAttributeRWImpl(searchAttrsTypeMap, req.getSearchAttributes());
         final CommunicationImpl communication = new CommunicationImpl(
+                req.getInternalChannelInfos(),
+                req.getSignalChannelInfos(),
                 registry.getInternalChannelTypeStore(req.getWorkflowType()),
+                registry.getSignalChannelTypeStore(req.getWorkflowType()),
                 workerOptions.getObjectEncoder(),
                 true
         );
@@ -169,7 +172,10 @@ public WorkflowStateWaitUntilResponse handleWorkflowStateWaitUntil(final Workflo
         final Map<String, SearchAttributeValueType> searchAttrsTypeMap = registry.getSearchAttributeKeyToTypeMap(req.getWorkflowType());
         final SearchAttributeRWImpl searchAttributeRW = new SearchAttributeRWImpl(searchAttrsTypeMap, req.getSearchAttributes());
         final CommunicationImpl communication = new CommunicationImpl(
+                new HashMap<>(),
+                new HashMap<>(),
                 registry.getInternalChannelTypeStore(req.getWorkflowType()),
+                registry.getSignalChannelTypeStore(req.getWorkflowType()),
                 workerOptions.getObjectEncoder(),
                 false
         );
@@ -234,7 +240,10 @@ public WorkflowStateExecuteResponse handleWorkflowStateExecute(final WorkflowSta
         final Map<String, SearchAttributeValueType> saTypeMap = registry.getSearchAttributeKeyToTypeMap(req.getWorkflowType());
         final SearchAttributeRWImpl searchAttributeRW = new SearchAttributeRWImpl(saTypeMap, req.getSearchAttributes());
         final CommunicationImpl communication = new CommunicationImpl(
+                new HashMap<>(),
+                new HashMap<>(),
                 registry.getInternalChannelTypeStore(req.getWorkflowType()),
+                registry.getSignalChannelTypeStore(req.getWorkflowType()),
                 workerOptions.getObjectEncoder(),
                 false
         );
diff --git a/src/main/java/io/iworkflow/core/communication/Communication.java b/src/main/java/io/iworkflow/core/communication/Communication.java
@@ -4,6 +4,22 @@
 
 public interface Communication {
 
+    /**
+     * Get the size of the internal channel(including the one being sent in the buffer)
+     * NOTE: currently only supported in RPC
+     * @param channelName the channel name to get size
+     * @return the size of the internal channel
+     */
+    int getInternalChannelSize(final String channelName);
+
+    /**
+     * Get the size of the signal channel(including the one being sent in the buffer)
+     * NOTE: currently only supported in RPC
+     * @param channelName the channel name to get size
+     * @return the size of the signal channel
+     */
+    int getSignalChannelSize(final String channelName);
+
     /**
      * Publish a value to an internal Channel
      *
diff --git a/src/main/java/io/iworkflow/core/communication/CommunicationImpl.java b/src/main/java/io/iworkflow/core/communication/CommunicationImpl.java
@@ -4,6 +4,7 @@
 import io.iworkflow.core.StateMovement;
 import io.iworkflow.core.TypeStore;
 import io.iworkflow.core.WorkflowDefinitionException;
+import io.iworkflow.gen.models.ChannelInfo;
 import io.iworkflow.gen.models.EncodedObject;
 
 import java.util.ArrayList;
@@ -14,7 +15,12 @@
 
 public class CommunicationImpl implements Communication {
 
-    final TypeStore typeStore;
+    final TypeStore internalChannelTypeStore;
+    final TypeStore signalChannelTypeStore;
+
+    final Map<String, ChannelInfo> internalChannelInfos;
+
+    final Map<String, ChannelInfo> signalChannelInfos;
     final Map<String, List<EncodedObject>> toPublish = new HashMap<>();
 
     final List<StateMovement> stateMovements = new ArrayList<>();
@@ -23,23 +29,63 @@ public class CommunicationImpl implements Communication {
     final ObjectEncoder objectEncoder;
 
     public CommunicationImpl(
-            final TypeStore typeStore,
+            final Map<String, ChannelInfo> internalChannelInfos,
+            final Map<String, ChannelInfo> signalChannelInfos,
+            final TypeStore internalChannelTypeStore,
+            final TypeStore signalChannelTypeStore,
             final ObjectEncoder objectEncoder,
             final boolean allowTriggerStateMovements) {
-        this.typeStore = typeStore;
+        this.internalChannelInfos = internalChannelInfos;
+        this.signalChannelInfos = signalChannelInfos;
+        this.internalChannelTypeStore = internalChannelTypeStore;
+        this.signalChannelTypeStore = signalChannelTypeStore;
         this.objectEncoder = objectEncoder;
         this.allowTriggerStateMovements = allowTriggerStateMovements;
     }
 
+    @Override
+    public int getInternalChannelSize(final String channelName) {
+        checkInternalChannelNameValid(channelName, null);
+        int size = 0;
+        if(internalChannelInfos.containsKey(channelName)){
+            size += internalChannelInfos.get(channelName).getSize();
+        }
+        if(toPublish.containsKey(channelName)){
+            size += toPublish.get(channelName).size();
+        }
+        return size;
+    }
+
+    @Override
+    public int getSignalChannelSize(final String channelName) {
+        checkSignalChannelNameValid(channelName, null);
+        if(signalChannelInfos.containsKey(channelName)){
+            return signalChannelInfos.get(channelName).getSize();
+        }
+        return 0;
+    }
+
     @Override
     public void publishInternalChannel(final String channelName, final Object value) {
-        final Class<?> type = typeStore.getType(channelName);
+        checkInternalChannelNameValid(channelName, value);
+        final List<EncodedObject> publish = toPublish.computeIfAbsent(channelName, s -> new ArrayList<>());
+        publish.add(objectEncoder.encode(value));
+    }
+
+    private void checkInternalChannelNameValid(final String channelName, final Object value) {
+        final Class<?> type = internalChannelTypeStore.getType(channelName);
 
         if (value != null && !type.isInstance(value)) {
             throw new WorkflowDefinitionException(String.format("InternalChannel value is not of type %s", type.getName()));
         }
-        final List<EncodedObject> publish = toPublish.computeIfAbsent(channelName, s -> new ArrayList<>());
-        publish.add(objectEncoder.encode(value));
+    }
+
+    private void checkSignalChannelNameValid(final String channelName, final Object value) {
+        final Class<?> type = signalChannelTypeStore.getType(channelName);
+
+        if (value != null && !type.isInstance(value)) {
+            throw new WorkflowDefinitionException(String.format("SignalChannel value is not of type %s", type.getName()));
+        }
     }
 
     @Override
diff --git a/src/test/java/io/iworkflow/integ/RpcTest.java b/src/test/java/io/iworkflow/integ/RpcTest.java
@@ -8,6 +8,7 @@
 import io.iworkflow.core.ImmutableWorkflowOptions;
 import io.iworkflow.gen.models.*;
 import io.iworkflow.integ.persistence.BasicPersistenceWorkflow;
+import io.iworkflow.integ.rpc.DeadEndStateWorkflow;
 import io.iworkflow.integ.rpc.NoStateWorkflow;
 import io.iworkflow.integ.rpc.RpcWorkflow;
 import io.iworkflow.integ.rpc.RpcWorkflowState2;
@@ -313,4 +314,23 @@ public void testRpcError() throws InterruptedException {
         client.stopWorkflow(wfId, null);
     }
 
+
+    @Test
+    public void testSignalChannelSizeInfo(){
+        final Client client = new Client(WorkflowRegistry.registry, ClientOptions.localDefault);
+        final String wfId = "testSignalChannelSizeInfo" + System.currentTimeMillis() / 1000;
+        client.startWorkflow(
+                DeadEndStateWorkflow.class, wfId, 10);
+        final DeadEndStateWorkflow rpcStub = client.newRpcStub(DeadEndStateWorkflow.class, wfId, "");
+        client.invokeRPC(rpcStub::sendAndGetInternalChannelSize);
+        final Integer size1 = client.invokeRPC(rpcStub::sendAndGetInternalChannelSize);
+        Assertions.assertEquals(2, size1);
+
+        client.signalWorkflow(DeadEndStateWorkflow.class, wfId, DeadEndStateWorkflow.IDLE_SIGNAL_CHANNEL, null);
+        client.signalWorkflow(DeadEndStateWorkflow.class, wfId, DeadEndStateWorkflow.IDLE_SIGNAL_CHANNEL, null);
+        client.signalWorkflow(DeadEndStateWorkflow.class, wfId, DeadEndStateWorkflow.IDLE_SIGNAL_CHANNEL, null);
+        final Integer size2 = client.invokeRPC(rpcStub::getSignalChannelSize);
+        Assertions.assertEquals(3, size2);
+
+    }
 }
diff --git a/src/test/java/io/iworkflow/integ/rpc/DeadEndStateWorkflow.java b/src/test/java/io/iworkflow/integ/rpc/DeadEndStateWorkflow.java
@@ -9,6 +9,9 @@
 import io.iworkflow.core.WorkflowState;
 import io.iworkflow.core.command.CommandResults;
 import io.iworkflow.core.communication.Communication;
+import io.iworkflow.core.communication.CommunicationMethodDef;
+import io.iworkflow.core.communication.InternalChannelDef;
+import io.iworkflow.core.communication.SignalChannelDef;
 import io.iworkflow.core.persistence.Persistence;
 import org.springframework.stereotype.Component;
 
@@ -19,6 +22,17 @@
 
 @Component
 public class DeadEndStateWorkflow implements ObjectWorkflow {
+
+    public static final String IDLE_INTERNAL_CHANNEL = "ideal-internal-channel";
+    public static final String IDLE_SIGNAL_CHANNEL = "ideal-signal-channel";
+    @Override
+    public List<CommunicationMethodDef> getCommunicationSchema() {
+        return Arrays.asList(
+                InternalChannelDef.create(Void.class, IDLE_INTERNAL_CHANNEL),
+                SignalChannelDef.create(Void.class, IDLE_SIGNAL_CHANNEL)
+        );
+    }
+
     @Override
     public List<StateDef> getWorkflowStates() {
         return Arrays.asList(
@@ -27,6 +41,16 @@ public List<StateDef> getWorkflowStates() {
         );
     }
 
+    @RPC
+    public int getSignalChannelSize(Context context, Persistence persistence, Communication communication) {
+        return communication.getSignalChannelSize(IDLE_SIGNAL_CHANNEL);
+    }
+
+    @RPC
+    public int sendAndGetInternalChannelSize(Context context,  Persistence persistence, Communication communication) {
+        communication.publishInternalChannel(IDLE_INTERNAL_CHANNEL, null);
+        return communication.getInternalChannelSize(IDLE_INTERNAL_CHANNEL);
+    }
     @RPC
     public Long testRpcFunc1(Context context, String input, Persistence persistence, Communication communication) {
         if (context.getWorkflowId().isEmpty() || context.getWorkflowRunId().isEmpty() ||
diff --git a/src/test/java/io/iworkflow/spring/controller/ApiController.java b/src/test/java/io/iworkflow/spring/controller/ApiController.java
@@ -68,19 +68,21 @@ private ResponseEntity<?> processWorkerException(final String workflowId, final
         PrintWriter pw = new PrintWriter(sw);
         e.printStackTrace(pw);
         String sStackTrace = sw.toString(); // stack trace as a string
-        if(sStackTrace.length()>500){
+        if(sStackTrace.length()>1000){
             // make sure NOT exceeding 4K limit in Temporal https://github.com/indeedeng/iwf/issues/272
-            sStackTrace = sStackTrace.substring(0, 500) + "...(TRUNCATED)";
+            sStackTrace = sStackTrace.substring(0, 1000) + "...(TRUNCATED)";
         }
         String msg = e.getMessage();
         if(msg == null){
             System.out.println("empty message for exception "+sw.toString());
             msg = "empty message for exception "+e;
         }
-        if (msg.length() > 50) {
-            msg = msg.substring(0, 50) + "...(TRUNCATED)";
+        if (msg.length() > 500) {
+            msg = msg.substring(0, 500) + "...(TRUNCATED)";
         }
 
+        System.out.println("Exception in "+methodType+" for workflow "+workflowId+" : "+msg + "\n STACKTRACE: \n"+sStackTrace);
+
         final WorkerErrorResponse errResp = new WorkerErrorResponse()
                 .detail(msg+"\n STACKTRACE: \n"+sStackTrace)
                 .errorType(e.getClass().getName());
