Support x-prefetch

Author: Simon MacMullen

test/src/com/rabbitmq/examples/PerfTest.java

@@ -58,7 +58,8 @@ public static void main(String[] args) {
             long confirm         = intArg(cmd, 'c', -1);
             boolean autoAck      = cmd.hasOption('a');
             int multiAckEvery    = intArg(cmd, 'A', 0);
-            int prefetchCount    = intArg(cmd, 'q', 0);
+            int channelPrefetch  = intArg(cmd, 'q', 0);
+            int consumerPrefetch = intArg(cmd, 'P', 0);
             int minMsgSize       = intArg(cmd, 's', 0);
             int timeLimit        = intArg(cmd, 'z', 0);
             int producerMsgCount = intArg(cmd, 'C', 0);
@@ -100,7 +101,8 @@ public static void main(String[] args) {
             p.setMultiAckEvery(    multiAckEvery);
             p.setMinMsgSize(       minMsgSize);
             p.setPredeclared(      predeclared);
-            p.setPrefetchCount(    prefetchCount);
+            p.setConsumerPrefetch( consumerPrefetch);
+            p.setChannelPrefetch(  channelPrefetch);
             p.setProducerCount(    producerCount);
             p.setProducerMsgCount( producerMsgCount);
             p.setProducerTxSize(   producerTxSize);
@@ -146,7 +148,8 @@ private static Options getOptions() {
         options.addOption(new Option("c", "confirm",       true, "max unconfirmed publishes"));
         options.addOption(new Option("a", "autoack",       false,"auto ack"));
         options.addOption(new Option("A", "multiAckEvery", true, "multi ack every"));
-        options.addOption(new Option("q", "qos",           true, "qos prefetch count"));
+        options.addOption(new Option("q", "qos",           true, "qos (channel) prefetch count"));
+        options.addOption(new Option("P", "prefetch",      true, "consumer prefetch count"));
         options.addOption(new Option("s", "size",          true, "message size"));
         options.addOption(new Option("z", "time",          true, "time limit"));
         options.addOption(new Option("C", "pmessages", true, "producer message count"));

test/src/com/rabbitmq/examples/perf/Consumer.java

@@ -25,13 +25,16 @@
 import java.io.ByteArrayInputStream;
 import java.io.DataInputStream;
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
 
 public class Consumer implements Runnable {
 
     private QueueingConsumer q;
     private Channel channel;
     private String           id;
     private String           queueName;
+    private int              prefetch;
     private int              txSize;
     private boolean          autoAck;
     private int              multiAckEvery;
@@ -40,11 +43,12 @@ public class Consumer implements Runnable {
     private long             timeLimit;
 
     public Consumer(Channel channel, String id,
-                    String queueName, int txSize, boolean autoAck,
+                    String queueName, int prefetch, int txSize, boolean autoAck,
                     int multiAckEvery, Stats stats, int msgLimit, int timeLimit) {
 
         this.channel       = channel;
         this.id            = id;
+        this.prefetch      = prefetch;
         this.queueName     = queueName;
         this.txSize        = txSize;
         this.autoAck       = autoAck;
@@ -77,7 +81,7 @@ public void run() {
                 } catch (ConsumerCancelledException e) {
                     System.out.println("Consumer cancelled by broker. Re-consuming.");
                     q = new QueueingConsumer(channel);
-                    channel.basicConsume(queueName, autoAck, q);
+                    channel.basicConsume(queueName, autoAck, "consumer-tag", false, false, args(prefetch), q);
                     continue;
                 }
                 totalMsgCount++;
@@ -114,4 +118,12 @@ public void run() {
             throw new RuntimeException(e);
         }
     }
+
+    private Map args(int prefetch) {
+        Map a = new HashMap();
+        if (prefetch != 0) {
+            a.put("x-prefetch", prefetch);
+        }
+        return a;
+    }
 }

test/src/com/rabbitmq/examples/perf/MulticastParams.java

@@ -32,7 +32,8 @@ public class MulticastParams {
     private int producerCount = 1;
     private int consumerTxSize = 0;
     private int producerTxSize = 0;
-    private int prefetchCount = 0;
+    private int channelPrefetch = 0;
+    private int consumerPrefetch = 0;
     private int minMsgSize = 0;
 
     private int timeLimit = 0;
@@ -102,8 +103,12 @@ public void setMultiAckEvery(int multiAckEvery) {
         this.multiAckEvery = multiAckEvery;
     }
 
-    public void setPrefetchCount(int prefetchCount) {
-        this.prefetchCount = prefetchCount;
+    public void setChannelPrefetch(int channelPrefetch) {
+        this.channelPrefetch = channelPrefetch;
+    }
+
+    public void setConsumerPrefetch(int consumerPrefetch) {
+        this.consumerPrefetch = consumerPrefetch;
     }
 
     public void setMinMsgSize(int minMsgSize) {
@@ -180,8 +185,8 @@ public Consumer createConsumer(Connection connection, Stats stats, String id) th
         Channel channel = connection.createChannel();
         if (consumerTxSize > 0) channel.txSelect();
         String qName = configureQueue(connection, id);
-        if (prefetchCount > 0) channel.basicQos(prefetchCount);
-        return new Consumer(channel, id, qName,
+        if (channelPrefetch > 0) channel.basicQos(channelPrefetch);
+        return new Consumer(channel, id, qName, consumerPrefetch,
                                          consumerTxSize, autoAck, multiAckEvery,
                                          stats, consumerMsgCount, timeLimit);
     }