keep track of published message count in java client

Author: Alexandru Scvortov

src/com/rabbitmq/client/Channel.java

@@ -617,4 +617,11 @@ Queue.DeclareOk queueDeclare(String queue, boolean durable, boolean exclusive, b
      */
     Confirm.SelectOk confirmSelect(boolean multiple, boolean nowait)
         throws IOException;
+
+    /**
+     * Returns the number of messages published since the channel was
+     * put in confirm mode.
+     * @return the number of messages published since the first
+     * confirm.select */
+    long getPublishedMessageCount();
 }

src/com/rabbitmq/client/impl/ChannelN.java

@@ -104,6 +104,10 @@ public class ChannelN extends AMQChannel implements com.rabbitmq.client.Channel
      */
     public volatile AckListener ackListener = null;
 
+    /** Current published message count (used by publisher acknowledgements)
+     */
+    private long publishedMessageCount;
+
     /** Reference to the currently-active default consumer, or null if there is
      *  none.
      */
@@ -493,6 +497,10 @@ public void basicPublish(String exchange, String routingKey,
                              BasicProperties props, byte[] body)
         throws IOException
     {
+        synchronized (_channelMutex) {
+            if (publishedMessageCount >= 0)
+                ++publishedMessageCount;
+        }
         BasicProperties useProps = props;
         if (props == null) {
             useProps = MessageProperties.MINIMAL_BASIC;
@@ -823,6 +831,10 @@ public Confirm.SelectOk confirmSelect(boolean multiple)
     public Confirm.SelectOk confirmSelect(boolean multiple, boolean nowait)
         throws IOException
     {
+        synchronized (_channelMutex) {
+            if (publishedMessageCount == -1)
+                publishedMessageCount = 0;
+        }
         return (Confirm.SelectOk)
             exnWrappingRpc(new Confirm.Select(multiple, nowait)).getMethod();
     }
@@ -837,4 +849,8 @@ public Channel.FlowOk getFlow() {
         return new Channel.FlowOk(!_blockContent);
     }
 
+    /** Public API - {@inheritDoc} */
+    public long getPublishedMessageCount() {
+        return publishedMessageCount;
+    }
 }

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

@@ -82,7 +82,7 @@ public static void main(String[] args) {
             int producerTxSize   = intArg(cmd, 'm', 0);
             int consumerTxSize   = intArg(cmd, 'n', 0);
             boolean pubAck       = cmd.hasOption('c');
-            int pubAckCount      = intArg(cmd, 'k', 0);
+            long pubAckCount     = intArg(cmd, 'k', 0);
             boolean autoAck      = cmd.hasOption('a');
             int prefetchCount    = intArg(cmd, 'q', 0);
             int minMsgSize       = intArg(cmd, 's', 0);
@@ -139,9 +139,10 @@ public static void main(String[] args) {
                 if (pubAck) channel.confirmSelect(true);
                 channel.exchangeDeclare(exchangeName, exchangeType);
                 final Producer p = new Producer(channel, exchangeName, id,
-                                                   flags, producerTxSize,
-                                                   1000L * samplingInterval,
-                                                   rateLimit, minMsgSize, timeLimit);
+                                                flags, producerTxSize,
+                                                1000L * samplingInterval,
+                                                rateLimit, minMsgSize, timeLimit,
+                                                pubAckCount);
                 channel.setReturnListener(p);
                 channel.setAckListener(p);
                 Thread t = new Thread(p);
@@ -233,9 +234,13 @@ public static class Producer implements Runnable, ReturnListener, AckListener {
         private int     msgCount;
         private int     basicReturnCount;
 
+        private long    pubAckCount;
+        private long    mostRecentAcked;
+
         public Producer(Channel channel, String exchangeName, String id,
                         List flags, int txSize,
-                        long interval, int rateLimit, int minMsgSize, int timeLimit)
+                        long interval, int rateLimit, int minMsgSize, int timeLimit,
+                        long pubAckCount)
             throws IOException {
 
             this.channel      = channel;
@@ -248,6 +253,7 @@ public Producer(Channel channel, String exchangeName, String id,
             this.interval     = interval;
             this.rateLimit    = rateLimit;
             this.timeLimit    = 1000L * timeLimit;
+            this.pubAckCount  = pubAckCount;
             this.message      = new byte[minMsgSize];
         }
 
@@ -269,7 +275,16 @@ public synchronized void resetBasicReturns() {
         }
 
         public void handleAck(long sequenceNumber, boolean multiple) {
-            System.out.printf("got an ack for %d\n", sequenceNumber);
+            if (multiple) {
+                logAck(sequenceNumber);
+                System.out.printf("got an ack all messages up to %d\n", sequenceNumber);
+            } else {
+                System.out.printf("got an ack for message %d\n", sequenceNumber);
+            }
+        }
+
+        private synchronized void logAck(long seqNum) {
+            mostRecentAcked = seqNum;
         }
 
         public void run() {
@@ -282,6 +297,7 @@ public void run() {
             try {
 
                 while (timeLimit == 0 || now < startTime + timeLimit) {
+                    delayPubAck();
                     delay(now);
                     publish(createMessage(totalMsgCount));
 		    totalMsgCount++;
@@ -314,6 +330,16 @@ private void publish(byte[] msg)
                                  msg);
         }
 
+        private void delayPubAck()
+            throws InterruptedException {
+            if (pubAckCount == 0)
+                return;
+            while (channel.getPublishedMessageCount() - mostRecentAcked
+                   > pubAckCount) {
+                Thread.sleep(100);
+            }
+        }
+
         private void delay(long now)
             throws InterruptedException {