Merge pull request #204 from rabbitmq/rabbitmq-perf-test-198-nack

Add option to nack messages

Author: gerhard

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

@@ -35,6 +35,12 @@ public class Consumer extends AgentBase implements Runnable {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(Consumer.class);
 
+    private static final AckNackOperation ACK_OPERATION =
+            (ch, envelope, multiple) -> ch.basicAck(envelope.getDeliveryTag(), multiple);
+
+    private static final AckNackOperation NACK_OPERATION =
+            (ch, envelope, multiple) -> ch.basicNack(envelope.getDeliveryTag(), multiple, true);
+
     private volatile ConsumerImpl       q;
     private final Channel               channel;
     private final String                id;
@@ -64,6 +70,8 @@ public class Consumer extends AgentBase implements Runnable {
 
     private final int pollingInterval;
 
+    private final AckNackOperation ackNackOperation;
+
     public Consumer(ConsumerParameters parameters) {
         this.channel           = parameters.getChannel();
         this.id                = parameters.getId();
@@ -108,6 +116,13 @@ public Consumer(ConsumerParameters parameters) {
             };
         }
 
+        if (parameters.isNack()) {
+            this.ackNackOperation = NACK_OPERATION;
+        } else {
+            this.ackNackOperation = ACK_OPERATION;
+        }
+
+
         this.state = new ConsumerState(parameters.getRateLimit());
         this.recoveryProcess = parameters.getRecoveryProcess();
         this.recoveryProcess.init(this);
@@ -243,9 +258,9 @@ private void ackIfNecessary(Envelope envelope, int currentMessageCount, final Ch
             if (!autoAck) {
                 dealWithWriteOperation(() -> {
                     if (multiAckEvery == 0) {
-                        ch.basicAck(envelope.getDeliveryTag(), false);
+                        ackNackOperation.apply(ch, envelope, false);
                     } else if (currentMessageCount % multiAckEvery == 0) {
-                        ch.basicAck(envelope.getDeliveryTag(), true);
+                        ackNackOperation.apply(ch, envelope, true);
                     }
                 }, recoveryProcess);
             }
@@ -399,4 +414,11 @@ public void simulateLatency() {
         }
     }
 
+    @FunctionalInterface
+    private interface AckNackOperation {
+
+        void apply(Channel channel, Envelope envelope, boolean multiple) throws IOException;
+
+    }
+
 }

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

@@ -41,6 +41,7 @@ public class ConsumerParameters {
     private Recovery.RecoveryProcess recoveryProcess;
     private ExecutorService executorService;
     private boolean polling;
+    private boolean nack = false;
 
     private int pollingInterval;
 
@@ -187,4 +188,13 @@ public ConsumerParameters setPollingInterval(int pollingInterval) {
         this.pollingInterval = pollingInterval;
         return this;
     }
+
+    public boolean isNack() {
+        return nack;
+    }
+
+    public ConsumerParameters setNack(boolean nack) {
+        this.nack = nack;
+        return this;
+    }
 }

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

@@ -107,6 +107,8 @@ public class MulticastParams {
 
     private int pollingInterval = -1;
 
+    private boolean nack = false;
+
     public void setExchangeType(String exchangeType) {
         this.exchangeType = exchangeType;
     }
@@ -396,6 +398,10 @@ public void setPollingInterval(int pollingInterval) {
         this.pollingInterval = pollingInterval;
     }
 
+    public void setNack(boolean nack) {
+        this.nack = nack;
+    }
+
     public Producer createProducer(Connection connection, Stats stats, MulticastSet.CompletionHandler completionHandler,
                                    ValueIndicator<Float> rateIndicator, ValueIndicator<Integer> messageSizeIndicator) throws IOException {
         Channel channel = connection.createChannel(); //NOSONAR
@@ -471,6 +477,7 @@ public Consumer createConsumer(Connection connection, Stats stats, MulticastSet.
                 .setExecutorService(executorService)
                 .setPolling(this.polling)
                 .setPollingInterval(this.pollingInterval)
+                .setNack(this.nack)
         );
         this.topologyHandler.next();
         return consumer;

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

@@ -167,6 +167,8 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             boolean polling = hasOption(cmd, "po");
             int pollingInterval = intArg(cmd, "pi", -1);
 
+            boolean nack = hasOption(cmd, "na");
+
             String uri               = strArg(cmd, 'h', "amqp://localhost");
             String urisParameter     = strArg(cmd, 'H', null);
             String outputFile        = strArg(cmd, 'o', null);
@@ -298,6 +300,7 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             p.setMessageSizes(variableSizes);
             p.setPolling(polling);
             p.setPollingInterval(pollingInterval);
+            p.setNack(nack);
 
             MulticastSet.CompletionHandler completionHandler = getCompletionHandler(p);
 
@@ -567,6 +570,8 @@ public static Options getOptions() {
         options.addOption(new Option("po", "polling",false,"use basic.get to consume messages"));
         options.addOption(new Option("pi", "polling-interval",true, "time to wait before polling with basic.get, " +
                 "in millisecond, default is 0."));
+
+        options.addOption(new Option("na", "nack",false,"nack and requeue messages"));
         return options;
     }
 

src/test/java/com/rabbitmq/perf/MessageCountTimeLimitAndPublishingIntervalRateTest.java

@@ -24,7 +24,9 @@
 import org.junit.jupiter.api.TestInfo;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.ArgumentsSource;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -666,6 +668,75 @@ public void pollingOnlyDoesNotStop() throws Exception {
         waitAtMost(20, () -> testIsDone.get());
     }
 
+    @ParameterizedTest
+    @ValueSource(strings = {"false", "true"})
+    public void ackNack(String nackParameter) throws Exception {
+        boolean nack = Boolean.valueOf(nackParameter);
+        int messagesCount = 100;
+        countsAndTimeLimit(0, messagesCount, 0);
+        params.setConsumerCount(1);
+        params.setProducerCount(0);
+        params.setQueueNames(asList("queue"));
+        params.setNack(nack);
+
+        CountDownLatch consumersLatch = new CountDownLatch(1);
+        AtomicInteger consumerTagCounter = new AtomicInteger(0);
+        List<Consumer> consumers = new CopyOnWriteArrayList<>();
+        AtomicInteger acks = new AtomicInteger(0);
+        AtomicInteger nacks = new AtomicInteger(0);
+        Channel channel = proxy(Channel.class,
+                callback("basicConsume", (proxy, method, args) -> {
+                    consumers.add((Consumer) args[2]);
+                    consumersLatch.countDown();
+                    return consumerTagCounter.getAndIncrement() + "";
+                }),
+                callback("basicAck", (proxy, method, args) -> {
+                    acks.incrementAndGet();
+                    return null;
+                }),
+                callback("basicNack", (proxy, method, args) -> {
+                    nacks.incrementAndGet();
+                    return null;
+                })
+        );
+
+        Connection connection = proxy(Connection.class,
+                callback("createChannel", (proxy, method, args) -> channel)
+        );
+
+        MulticastSet multicastSet = getMulticastSet(connectionFactoryThatReturns(connection));
+        run(multicastSet);
+
+        waitForRunToStart();
+
+        assertThat("consumer should have been registered by now",
+                consumersLatch.await(5, TimeUnit.SECONDS), is(true));
+
+        waitAtMost(20, () -> consumers.size() == 1);
+
+        Collection<Future<?>> sendTasks = new ArrayList<>(consumers.size());
+        for (Consumer consumer : consumers) {
+            Collection<Future<?>> tasks = sendMessagesToConsumer(messagesCount, consumer);
+            sendTasks.addAll(tasks);
+        }
+
+        for (Future<?> latch : sendTasks) {
+            latch.get(10, TimeUnit.SECONDS);
+        }
+
+        waitAtMost(20, () -> testIsDone.get());
+
+        if (nack) {
+            Assertions.assertThat(acks).hasValue(0);
+            Assertions.assertThat(nacks).hasValue(messagesCount);
+        } else {
+            Assertions.assertThat(acks).hasValue(messagesCount);
+            Assertions.assertThat(nacks).hasValue(0);
+        }
+
+
+    }
+
     private Collection<Future<?>> sendMessagesToConsumer(int messagesCount, Consumer consumer) {
         final Collection<Future<?>> tasks = new ArrayList<>(messagesCount);
         IntStream.range(0, messagesCount).forEach(i -> {