Refactor transactional counting to a producer consumer pattern

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/transaction/TransactionCountConfig.java

@@ -0,0 +1,15 @@
+package org.wso2.carbon.apimgt.gateway.handlers.transaction;
+
+import org.wso2.carbon.apimgt.impl.APIManagerConfiguration;
+import org.wso2.carbon.apimgt.impl.internal.ServiceReferenceHolder;
+
+public class TransactionCountConfig {
+
+    private static APIManagerConfiguration apiManagerConfiguration;
+    public TransactionCountConfig() {
+        apiManagerConfiguration = ServiceReferenceHolder.getInstance().getAPIManagerConfigurationService()
+                .getAPIManagerConfiguration();
+    }
+
+
+}

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/transaction/TransactionCountHandler.java

@@ -6,52 +6,46 @@
 import org.apache.synapse.MessageContext;
 import org.apache.synapse.core.axis2.Axis2MessageContext;
 import org.wso2.carbon.apimgt.gateway.APIMgtGatewayConstants;
-import org.wso2.carbon.apimgt.gateway.handlers.transaction.store.TransactionCountStore;
+import org.wso2.carbon.apimgt.gateway.handlers.transaction.queue.TransactionRecordQueue;
+import org.wso2.carbon.apimgt.gateway.handlers.transaction.store.TransactionRecordStore;
 
 import java.lang.reflect.Constructor;
-import java.util.ArrayList;
 import java.util.concurrent.*;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.locks.ReentrantLock;
 
 public class TransactionCountHandler extends AbstractExtendedSynapseHandler {
 
     // Todo: Make these parameters configurable via deployment.toml
-    private static final double MAX_TRANSACTION_COUNT = Integer.MAX_VALUE * 0.9;
-    private int MAX_RETRY_COUNT = 3;
-    private int TRANSACTION_COUNT_COMMIT_INTERVAL = 10;
-    private int THREAD_POOL_SIZE = 5;
-    private int TRANSACTION_COUNT_QUEUE_SIZE = 1000;
-    private String TRANSACTION_COUNT_STORE_CLASS = "org.wso2.carbon.apimgt.gateway.handlers.transaction.store.TransactionCountStoreImpl";
+    private final int PRODUCER_THREAD_POOL_SIZE = 5;
+    private final int CONSUMER_THREAD_POOL_SIZE = 5;
+    private final int TRANSACTION_RECORD_QUEUE_SIZE = 1000;
+    private final String TRANSACTION_COUNT_STORE_CLASS = "org.wso2.carbon.apimgt.gateway.handlers.transaction.store.TransactionRecordStoreImpl";
 
     private static final Log LOG = LogFactory.getLog(TransactionCountHandler.class);
-    private static ReentrantLock lock = new ReentrantLock();
-    private static AtomicInteger transactionCount = new AtomicInteger(0);
-
-    private BlockingQueue<TransactionCountRecord> transactionCountRecordQueue;
-    private ExecutorService transactionCountExecutor;
-    private ScheduledExecutorService transactionCountScheduledExecutor;
-    private TransactionCountStore trasactionCountStore;
+    private TransactionRecordQueue transactionRecordQueue;
+    private TransactionRecordProducer transactionRecordProducer;
+    private TransactionRecordConsumer transactionRecordConsumer;
+    private TransactionRecordStore trasactionCountStore;
 
 
     public TransactionCountHandler() {
 
-        this.transactionCountRecordQueue = new LinkedBlockingDeque<>(TRANSACTION_COUNT_QUEUE_SIZE);
-        this.transactionCountExecutor = Executors.newFixedThreadPool(THREAD_POOL_SIZE);
-        this.transactionCountScheduledExecutor = Executors.newSingleThreadScheduledExecutor();
-
+        this.transactionRecordQueue = TransactionRecordQueue.getInstance(TRANSACTION_RECORD_QUEUE_SIZE);
         // Load the transaction count store
         try {
             Class<?> clazz = Class.forName(TRANSACTION_COUNT_STORE_CLASS);
             Constructor<?> constructor = clazz.getConstructor();
-            this.trasactionCountStore = (TransactionCountStore) constructor.newInstance();
+            this.trasactionCountStore = (TransactionRecordStore) constructor.newInstance();
         } catch (Exception e) {
             LOG.error("Error while loading the transaction count store.", e);
         }
 
-        // Start the transaction count record scheduler
-        transactionCountScheduledExecutor.scheduleAtFixedRate(this::handleScheduledTransactionCountCommit,
-                0, TRANSACTION_COUNT_COMMIT_INTERVAL, TimeUnit.SECONDS);
+        this.transactionRecordProducer = TransactionRecordProducer.getInstance(transactionRecordQueue,
+                PRODUCER_THREAD_POOL_SIZE);
+        this.transactionRecordConsumer = TransactionRecordConsumer.getInstance(transactionRecordQueue,
+                trasactionCountStore, CONSUMER_THREAD_POOL_SIZE);
+
+        this.transactionRecordProducer.start();
+        this.transactionRecordConsumer.start();
     }
 
     @Override
@@ -67,7 +61,7 @@ public boolean handleRequestInFlow(MessageContext messageContext) {
             String transport = axis2MessageContext.getIncomingTransportName();
             if (transport.equals(APIMgtGatewayConstants.TRANSPORT_WS) || transport.equals(APIMgtGatewayConstants.TRANSPORT_WSS)){
                 LOG.info("Counting WebSocket message");
-                transactionCountExecutor.execute(this::handleTransactionCountCommit);
+                this.transactionRecordProducer.addTransaction();
             }
         } catch (RejectedExecutionException e) {
             LOG.error("Transaction could not be counted.", e);
@@ -86,7 +80,7 @@ public boolean handleRequestOutFlow(MessageContext messageContext) {
             // Counting outgoing messages that are not related to any request-response pair
             if (isThereAnAssociatedIncomingRequest == null) {
                 LOG.info("Counting async outgoing message");
-                transactionCountExecutor.execute(this::handleTransactionCountCommit);
+                this.transactionRecordProducer.addTransaction();
             }
         } catch (RejectedExecutionException e) {
             LOG.error("Transaction could not be counted.", e);
@@ -109,67 +103,11 @@ public boolean handleResponseOutFlow(MessageContext messageContext) {
         // Counting request-response pairs
         if (isThereAnAssociatedIncomingRequest instanceof Boolean) {
             LOG.info("Counting request-response pair");
-            transactionCountExecutor.execute(this::handleTransactionCountCommit);
+            this.transactionRecordProducer.addTransaction();
         }
         return true;
     }
 
-    private void handleTransactionCountCommit() {
-        lock.lock();
-        try {
-            if (transactionCount.incrementAndGet() >= MAX_TRANSACTION_COUNT) {
-                TransactionCountRecord transactionCountRecord = new TransactionCountRecord(transactionCount.get());
-                transactionCountRecordQueue.add(transactionCountRecord);
-                transactionCount.set(0);
-            }
-        } catch (Exception e) {
-            LOG.error("Error while handling transaction count.", e);
-        } finally {
-            lock.unlock();
-        }
-        this.commitWithRetries();
-    }
-
-    private void handleScheduledTransactionCountCommit() {
-        lock.lock();
-        try {
-            int transactionCountValue = transactionCount.get();
-            if (transactionCountValue != 0) {
-                TransactionCountRecord transactionCountRecord = new TransactionCountRecord(transactionCountValue);
-                transactionCountRecordQueue.add(transactionCountRecord);
-                transactionCount.set(0);
-            }
-        } catch (Exception e) {
-            LOG.error("Error while handling transaction count.", e);
-        } finally {
-            lock.unlock();
-        }
-        this.commitWithRetries();
-    }
-
-    private void commitWithRetries() {
-        LOG.info("Local transaction count: " + transactionCount.get());
-        // Arraylist of transaction count records will be committed to the store
-        ArrayList<TransactionCountRecord> transactionCountRecordList = new ArrayList<>();
-        transactionCountRecordQueue.drainTo(transactionCountRecordList);
-
-        if (transactionCountRecordList.isEmpty()) {
-            return;
-        }
-
-        // Committing the transaction count records to the store with retries
-        // If failed to commit after MAX_RETRY_COUNT, the transaction count records will be added to the queue again
-        boolean commited = false;
-        int retryCount = 0;
-        while (!commited && retryCount < MAX_RETRY_COUNT) {
-            commited = trasactionCountStore.commit(transactionCountRecordList);
-            retryCount++;
-        }
-        if (!commited) {
-            transactionCountRecordQueue.addAll(transactionCountRecordList);
-        }
-    }
-
     @Override
     public boolean handleServerInit() {
         // Nothing to implement
@@ -179,8 +117,9 @@ public boolean handleServerInit() {
     @Override
     public boolean handleServerShutDown() {
         // Clen up resources
-        transactionCountExecutor.shutdown();
-        transactionCountScheduledExecutor.shutdown();
+        transactionRecordProducer.shutdown();
+        transactionRecordConsumer.shutdown();
+        transactionRecordQueue.clenUp();
         trasactionCountStore.clenUp();
         return true;
     }

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/transaction/TransactionCounter.java

@@ -0,0 +1,31 @@
+package org.wso2.carbon.apimgt.gateway.handlers.transaction;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+
+public class TransactionCounter {
+
+    private static TransactionCounter instance = null;
+    private static AtomicInteger transactionCount = new AtomicInteger(0);
+    private static ReentrantLock lock = new ReentrantLock();
+    private static int MAX_TRANSACTION_COUNT;
+    private TransactionCounter(int maxTransactionCount) {
+        this.MAX_TRANSACTION_COUNT = maxTransactionCount;
+    }
+
+    public static TransactionCounter getInstance(int maxTransactionCount) {
+        if(instance == null) {
+            instance = new TransactionCounter(maxTransactionCount);
+        }
+        return instance;
+    }
+    public static void increment() {
+        lock.lock();
+        try {
+            transactionCount.incrementAndGet();
+        } finally {
+            lock.unlock();
+        }
+    }
+
+}

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/transaction/TransactionCountRecord.java → components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/transaction/TransactionRecord.java

@@ -3,13 +3,13 @@
 import java.time.Instant;
 import java.util.UUID;
 
-public class TransactionCountRecord {
+public class TransactionRecord {
 
-    private final String id = UUID.randomUUID().toString();
+    private static final String id = UUID.randomUUID().toString();
     private Integer count;
     private final String recordedTime;
 
-    public TransactionCountRecord(Integer count) {
+    public TransactionRecord(Integer count) {
         this.count = count;
         this.recordedTime = Instant.now().toString();
     }

components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/transaction/TransactionRecordConsumer.java

@@ -0,0 +1,88 @@
+package org.wso2.carbon.apimgt.gateway.handlers.transaction;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.wso2.carbon.apimgt.gateway.handlers.transaction.queue.TransactionRecordQueue;
+import org.wso2.carbon.apimgt.gateway.handlers.transaction.store.TransactionRecordStore;
+
+import java.util.ArrayList;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class TransactionRecordConsumer {
+
+    // Todo: Make these parameters configurable via deployment.toml
+    private int MAX_RETRY_COUNT = 3;
+    private int TRANSACTION_COUNT_COMMIT_INTERVAL = 10;
+    private int MAX_TRANSACTION_RECORDS_PER_COMMIT = 10;
+
+    private static final Log LOG = LogFactory.getLog(TransactionRecordConsumer.class);
+    private static TransactionRecordConsumer instance = null;
+    private TransactionRecordStore transactionRecordStore;
+    private TransactionRecordQueue transactionRecordQueue;
+    private ExecutorService executorService;
+    private final int threadPoolSize;
+
+    private TransactionRecordConsumer(TransactionRecordStore transactionRecordStore,
+                                      TransactionRecordQueue transactionRecordQueue, int threadPoolSize) {
+        this.transactionRecordStore = transactionRecordStore;
+        this.transactionRecordQueue = transactionRecordQueue;
+        this.executorService = Executors.newFixedThreadPool(threadPoolSize);
+        this.threadPoolSize = threadPoolSize;
+    }
+
+    public static TransactionRecordConsumer getInstance(TransactionRecordQueue transactionRecordQueue,
+                                                        TransactionRecordStore transactionRecordStore,
+                                                        int threadPoolSize) {
+        if(instance == null) {
+            instance = new TransactionRecordConsumer(transactionRecordStore, transactionRecordQueue, threadPoolSize);
+        }
+        return instance;
+    }
+
+    public void start() {
+        LOG.info("Transaction record consumer started");
+        // execute the startCommitting method in all the threads
+        for (int i = 0; i < threadPoolSize; i++) {
+            executorService.execute(this::startCommitting);
+        }
+    }
+
+    private void startCommitting() {
+        try {
+            while (true) {
+                commitWithRetries();
+            }
+        } catch (InterruptedException ex) {
+            LOG.debug("Transaction record consumer interrupted");
+        }
+    }
+
+    private void commitWithRetries() throws InterruptedException {
+
+        // Arraylist of transaction count records will be committed to the store
+        ArrayList<TransactionRecord> transactionRecordList = new ArrayList<>();
+        TransactionRecord transactionRecord = null;
+        transactionRecord = transactionRecordQueue.take();
+
+        transactionRecordList.add(transactionRecord);
+        transactionRecordQueue.drain(transactionRecordList, MAX_TRANSACTION_RECORDS_PER_COMMIT);
+
+        // Committing the transaction count records to the store with retries
+        // If failed to commit after MAX_RETRY_COUNT, the transaction count records will be added to the queue again
+        boolean commited = false;
+        int retryCount = 0;
+        while (!commited && retryCount < MAX_RETRY_COUNT) {
+            commited = this.transactionRecordStore.commit(transactionRecordList);
+            retryCount++;
+        }
+        if (!commited) {
+            transactionRecordQueue.addAll(transactionRecordList);
+        }
+    }
+
+    public void shutdown() {
+        this.executorService.shutdownNow();
+    }
+
+}