Merge pull request #74 from rabbitmq/rabbitmq-perf-test-72

Enforce producer/consumer message and time limit

Author: michaelklishin

pom.xml

@@ -63,6 +63,7 @@
     <junit.platform.version>1.0.2</junit.platform.version>
     <mockito.version>2.13.0</mockito.version>
     <hamcrest.version>2.0.0.0</hamcrest.version>
+    <awaitility.version>3.0.0</awaitility.version>
 
     <!-- to sign artifacts when releasing -->
     <gpg.keyname>6026DFCA</gpg.keyname>
@@ -159,6 +160,14 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <version>${awaitility.version}</version>
+      <scope>test</scope>
+    </dependency>
+
+
     <dependency>
       <groupId>ch.qos.logback</groupId>
       <artifactId>logback-classic</artifactId>

src/main/java/com/rabbitmq/perf/Consumer.java

@@ -28,7 +28,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiFunction;
 
 public class Consumer extends ProducerConsumerBase implements Runnable {
@@ -42,17 +42,19 @@ public class Consumer extends ProducerConsumerBase implements Runnable {
     private final int                   multiAckEvery;
     private final Stats                 stats;
     private final int                   msgLimit;
-    private final CountDownLatch        latch = new CountDownLatch(1);
     private final Map<String, String>   consumerTagBranchMap = Collections.synchronizedMap(new HashMap<String, String>());
     private final ConsumerLatency       consumerLatency;
     private final BiFunction<BasicProperties, byte[], Long> timestampExtractor;
     private final TimestampProvider     timestampProvider;
+    private final MulticastSet.CompletionHandler completionHandler;
+    private final AtomicBoolean completed = new AtomicBoolean(false);
 
     public Consumer(Channel channel, String id,
                     List<String> queueNames, int txSize, boolean autoAck,
                     int multiAckEvery, Stats stats, float rateLimit, int msgLimit,
                     int consumerLatencyInMicroSeconds,
-                    TimestampProvider timestampProvider) {
+                    TimestampProvider timestampProvider,
+                    MulticastSet.CompletionHandler completionHandler) {
 
         this.channel           = channel;
         this.id                = id;
@@ -64,6 +66,7 @@ public Consumer(Channel channel, String id,
         this.stats             = stats;
         this.msgLimit          = msgLimit;
         this.timestampProvider = timestampProvider;
+        this.completionHandler = completionHandler;
 
         if (consumerLatencyInMicroSeconds <= 0) {
             this.consumerLatency = new NoWaitConsumerLatency();
@@ -147,13 +150,13 @@ public void handleDelivery(String consumerTag, Envelope envelope, BasicPropertie
                 consumerLatency.simulateLatency();
             }
             if (msgLimit != 0 && msgCount >= msgLimit) { // NB: not quite the inverse of above
-                latch.countDown();
+                countDown();
             }
         }
 
         @Override
         public void handleShutdownSignal(String consumerTag, ShutdownSignalException sig) {
-            latch.countDown();
+            countDown();
         }
 
         @Override
@@ -169,6 +172,12 @@ public void handleCancel(String consumerTag) throws IOException {
         }
     }
 
+    private void countDown() {
+        if (completed.compareAndSet(false, true)) {
+            completionHandler.countDown();
+        }
+    }
+
     private interface ConsumerLatency {
 
         void simulateLatency();

src/main/java/com/rabbitmq/perf/MulticastParams.java

@@ -299,7 +299,19 @@ public int getHeartbeatSenderThreads() {
         return heartbeatSenderThreads <= 0 ? producerCount + consumerCount : this.heartbeatSenderThreads;
     }
 
-    public Producer createProducer(Connection connection, Stats stats) throws IOException {
+    public int getTimeLimit() {
+        return timeLimit;
+    }
+
+    public int getProducerMsgCount() {
+        return producerMsgCount;
+    }
+
+    public int getConsumerMsgCount() {
+        return consumerMsgCount;
+    }
+
+    public Producer createProducer(Connection connection, Stats stats, MulticastSet.CompletionHandler completionHandler) throws IOException {
         Channel channel = connection.createChannel();
         if (producerTxSize > 0) channel.txSelect();
         if (confirm >= 0) channel.confirmSelect();
@@ -318,16 +330,15 @@ public Producer createProducer(Connection connection, Stats stats) throws IOExce
         final Producer producer = new Producer(channel, exchangeName, this.topologyHandler.getRoutingKey(),
                                                randomRoutingKey, flags, producerTxSize,
                                                producerRateLimit, producerMsgCount,
-                                               timeLimit,
                                                confirm, confirmTimeout, messageBodySource,
-                                               tsp, stats);
+                                               tsp, stats, completionHandler);
         channel.addReturnListener(producer);
         channel.addConfirmListener(producer);
         this.topologyHandler.next();
         return producer;
     }
 
-    public Consumer createConsumer(Connection connection, Stats stats) throws IOException {
+    public Consumer createConsumer(Connection connection, Stats stats, MulticastSet.CompletionHandler completionHandler) throws IOException {
         Channel channel = connection.createChannel();
         if (consumerTxSize > 0) channel.txSelect();
         List<String> generatedQueueNames = this.topologyHandler.configureQueuesForClient(connection);
@@ -344,7 +355,7 @@ public Consumer createConsumer(Connection connection, Stats stats) throws IOExce
         Consumer consumer = new Consumer(channel, this.topologyHandler.getRoutingKey(), generatedQueueNames,
                                          consumerTxSize, autoAck, multiAckEvery,
                                          stats, consumerRateLimit, consumerMsgCount,
-                                         consumerLatencyInMicroseconds, tsp);
+                                         consumerLatencyInMicroseconds, tsp, completionHandler);
         this.topologyHandler.next();
         return consumer;
     }
@@ -379,6 +390,10 @@ private static boolean queueExists(Connection connection, final String queueName
         return queueName != null && exists(connection, ch -> ch.queueDeclarePassive(queueName));
     }
 
+    public boolean hasLimit() {
+        return this.timeLimit > 0 || this.consumerMsgCount > 0 || this.producerCount > 0;
+    }
+
     private interface Checker {
         void check(Channel ch) throws IOException;
     }

src/main/java/com/rabbitmq/perf/MulticastSet.java

@@ -26,13 +26,12 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Random;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.SynchronousQueue;
-import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Supplier;
@@ -52,19 +51,21 @@ public class MulticastSet {
 
     private ThreadingHandler threadingHandler = new DefaultThreadingHandler();
 
+    private final CompletionHandler completionHandler;
+
     public MulticastSet(Stats stats, ConnectionFactory factory,
-        MulticastParams params, List<String> uris) {
-        this(stats, factory, params, "perftest", uris);
+        MulticastParams params, List<String> uris, CompletionHandler completionHandler) {
+        this(stats, factory, params, "perftest", uris, completionHandler);
     }
 
     public MulticastSet(Stats stats, ConnectionFactory factory,
-        MulticastParams params, String testID, List<String> uris) {
+        MulticastParams params, String testID, List<String> uris, CompletionHandler completionHandler) {
         this.stats = stats;
         this.factory = factory;
         this.params = params;
         this.testID = testID;
         this.uris = uris;
-
+        this.completionHandler = completionHandler;
         this.params.init();
     }
 
@@ -108,7 +109,7 @@ public void run(boolean announceStartup)
                 if (announceStartup) {
                     System.out.println("id: " + testID + ", starting consumer #" + i + ", channel #" + j);
                 }
-                consumerRunnables[(i * params.getConsumerChannelCount()) + j] = params.createConsumer(conn, stats);
+                consumerRunnables[(i * params.getConsumerChannelCount()) + j] = params.createConsumer(conn, stats, this.completionHandler);
             }
         }
 
@@ -134,7 +135,7 @@ public void run(boolean announceStartup)
                     System.out.println("id: " + testID + ", starting producer #" + i + ", channel #" + j);
                 }
                 AgentState agentState = new AgentState();
-                agentState.runnable = params.createProducer(conn, stats);
+                agentState.runnable = params.createProducer(conn, stats, this.completionHandler);
                 producerStates[(i * params.getProducerChannelCount()) + j] = agentState;
             }
         }
@@ -154,9 +155,11 @@ public void run(boolean announceStartup)
             producerState.task = producersExecutorService.submit(producerState.runnable);
         }
 
+        this.completionHandler.waitForCompletion();
+
         int count = 1; // counting the threads
         for (int i = 0; i < producerStates.length; i++) {
-            producerStates[i].task.get();
+            producerStates[i].task.cancel(true);
             if(count % params.getProducerChannelCount() == 0) {
                 // this is the end of a group of threads on the same connection,
                 // closing the connection
@@ -251,4 +254,51 @@ private static class AgentState {
 
     }
 
+    interface CompletionHandler {
+
+        void waitForCompletion() throws InterruptedException;
+
+        void countDown();
+
+    }
+
+    static class DefaultCompletionHandler implements CompletionHandler {
+
+        private final int timeLimit;
+        private final CountDownLatch latch;
+
+        DefaultCompletionHandler(int timeLimit, int countLimit) {
+            this.timeLimit = timeLimit;
+            this.latch = new CountDownLatch(countLimit <= 0 ? 1 : countLimit);
+        }
+
+        @Override
+        public void waitForCompletion() throws InterruptedException {
+            if (timeLimit <= 0) {
+                this.latch.await();
+            } else {
+                this.latch.await(timeLimit, TimeUnit.SECONDS);
+            }
+        }
+
+        @Override
+        public void countDown() {
+            latch.countDown();
+        }
+    }
+
+    static class NoLimitCompletionHandler implements CompletionHandler {
+
+        private final CountDownLatch latch = new CountDownLatch(1);
+
+        @Override
+        public void waitForCompletion() throws InterruptedException {
+            latch.await();
+        }
+
+        @Override
+        public void countDown() { }
+    }
+
+
 }

src/main/java/com/rabbitmq/perf/PerfTest.java

@@ -209,7 +209,9 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             p.setQueueSequenceTo(to);
             p.setHeartbeatSenderThreads(heartbeatSenderThreads);
 
-            MulticastSet set = new MulticastSet(stats, factory, p, testID, uris);
+            MulticastSet.CompletionHandler completionHandler = getCompletionHandler(p);
+
+            MulticastSet set = new MulticastSet(stats, factory, p, testID, uris, completionHandler);
             set.run(true);
 
             stats.printFinal();
@@ -224,6 +226,26 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
         }
     }
 
+    static MulticastSet.CompletionHandler getCompletionHandler(MulticastParams p) {
+        MulticastSet.CompletionHandler completionHandler;
+        if (p.hasLimit()) {
+            int countLimit = 0;
+            if (p.getProducerMsgCount() > 0) {
+                countLimit += p.getProducerThreadCount();
+            }
+            if (p.getConsumerMsgCount() > 0) {
+                countLimit += p.getProducerThreadCount();
+            }
+            completionHandler = new MulticastSet.DefaultCompletionHandler(
+                p.getTimeLimit(),
+                 countLimit
+            );
+        } else {
+            completionHandler = new MulticastSet.NoLimitCompletionHandler();
+        }
+        return completionHandler;
+    }
+
     public static void main(String[] args) {
         main(args, new PerfTestOptions().setSystemExiter(new JvmSystemExiter()).setSkipSslContextConfiguration(false));
     }