Update transaction counting

Author: Isuru Wijesiri

components/org.wso2.micro.integrator.initializer/src/main/java/org/wso2/micro/integrator/initializer/handler/transaction/TransactionCountHandler.java

@@ -26,6 +26,7 @@
 import org.apache.synapse.core.axis2.Axis2MessageContext;
 import org.wso2.micro.integrator.initializer.handler.DataHolder;
 import org.wso2.micro.integrator.initializer.handler.transaction.store.TransactionStore;
+import org.wso2.micro.integrator.initializer.utils.Constants;
 
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.RejectedExecutionException;
@@ -45,7 +46,17 @@ public TransactionCountHandler(ExecutorService executor) {
     @Override
     public boolean handleRequestInFlow(MessageContext messageContext) {
         try {
-            transactionCountExecutor.execute(() -> handleTransactionCount(messageContext));
+            org.apache.axis2.context.MessageContext axis2MessageContext =
+                    ((Axis2MessageContext) messageContext).getAxis2MessageContext();
+
+            // Setting this property to identify request-response pairs
+            axis2MessageContext.setProperty(Constants.IS_THERE_ASSOCIATED_INCOMING_REQUEST, true);
+
+            // Counting message received via an open WebSocket
+            String transport = axis2MessageContext.getIncomingTransportName();
+            if (transport.equals(Constants.TRANSPORT_WS) || transport.equals(Constants.TRANSPORT_WSS)){
+                transactionCountExecutor.execute(() -> handleTransactionCount(messageContext));
+            }
         } catch (RejectedExecutionException e) {
             LOG.error("Transaction could not be counted.", e);
         }
@@ -54,6 +65,19 @@ public boolean handleRequestInFlow(MessageContext messageContext) {
 
     @Override
     public boolean handleRequestOutFlow(MessageContext messageContext) {
+        try {
+            org.apache.axis2.context.MessageContext axis2MessageContext =
+                    ((Axis2MessageContext) messageContext).getAxis2MessageContext();
+            Object isThereAnAssociatedIncomingRequest = axis2MessageContext.getProperty(
+                    Constants.IS_THERE_ASSOCIATED_INCOMING_REQUEST);
+
+            // Counting outgoing messages that are not related to any request-response pair
+            if (isThereAnAssociatedIncomingRequest == null) {
+                transactionCountExecutor.execute(() -> handleTransactionCount(messageContext));
+            }
+        } catch (RejectedExecutionException e) {
+            LOG.error("Transaction could not be counted.", e);
+        }
         return true;
     }
 
@@ -64,6 +88,15 @@ public boolean handleResponseInFlow(MessageContext messageContext) {
 
     @Override
     public boolean handleResponseOutFlow(MessageContext messageContext) {
+        org.apache.axis2.context.MessageContext axis2MessageContext =
+                ((Axis2MessageContext) messageContext).getAxis2MessageContext();
+        Object isThereAnAssociatedIncomingRequest = axis2MessageContext.getProperty(
+                Constants.IS_THERE_ASSOCIATED_INCOMING_REQUEST);
+
+        // Counting request-response pairs
+        if (isThereAnAssociatedIncomingRequest instanceof Boolean) {
+            transactionCountExecutor.execute(() -> handleTransactionCount(messageContext));
+        }
         return true;
     }
 

components/org.wso2.micro.integrator.initializer/src/main/java/org/wso2/micro/integrator/initializer/utils/Constants.java

@@ -76,5 +76,7 @@ public class Constants {
     public static final String HTTPS_LISTENER_PORT = "transport.https.listener.parameter.port";
     public static final String HTTP_LISTENER_PORT = "transport.http.listener.parameter.port";
     public static final String SERVER_PORT_OFFSET = "portOffset";
-
+    public static final String IS_THERE_ASSOCIATED_INCOMING_REQUEST = "is_there_incoming_request";
+    public static final String TRANSPORT_WS = "ws";
+    public static final String TRANSPORT_WSS = "wss";
 }