Add timeout when waiting for publisher confirms

The --confirm options is for the outstanding publisher confirms,
Producer doesn't publish when the limit is reached, waiting on a Semaphore.
There was no timeout, so the producer waits for ever even if the connection is dead.
This could block the JVM process when the broker would die abruptly.

The --confirmTimeout option now allows to specify a timeout in seconds.
The default is 30 seconds, a negative value means no timeout (i.e.
same behaviour as before).

Fixes #53

Author: acogoluegnes

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

@@ -27,6 +27,7 @@
 
 public class MulticastParams {
     private long confirm = -1;
+    private int confirmTimeout = 30;
     private int consumerCount = 1;
     private int producerCount = 1;
     private int consumerChannelCount = 1;
@@ -124,6 +125,10 @@ public void setConfirm(long confirm) {
         this.confirm = confirm;
     }
 
+    public void setConfirmTimeout(int confirmTimeout) {
+        this.confirmTimeout = confirmTimeout;
+    }
+
     public void setAutoAck(boolean autoAck) {
         this.autoAck = autoAck;
     }
@@ -246,7 +251,7 @@ public Producer createProducer(Connection connection, Stats stats, String id) th
                                                randomRoutingKey, flags, producerTxSize,
                                                producerRateLimit, producerMsgCount,
                                                timeLimit,
-                                               confirm, messageBodySource, stats);
+                                               confirm, confirmTimeout, messageBodySource, stats);
         channel.addReturnListener(producer);
         channel.addConfirmListener(producer);
         return producer;

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

@@ -72,6 +72,7 @@ public static void main(String[] args) {
             int producerTxSize       = intArg(cmd, 'm', 0);
             int consumerTxSize       = intArg(cmd, 'n', 0);
             long confirm             = intArg(cmd, 'c', -1);
+            int confirmTimeout       = intArg(cmd, "ct", 30);
             boolean autoAck          = cmd.hasOption('a');
             int multiAckEvery        = intArg(cmd, 'A', 0);
             int channelPrefetch      = intArg(cmd, 'Q', 0);
@@ -127,6 +128,7 @@ public static void main(String[] args) {
             p.setAutoAck(               autoAck);
             p.setAutoDelete(            autoDelete);
             p.setConfirm(               confirm);
+            p.setConfirmTimeout(        confirmTimeout);
             p.setConsumerCount(         consumerCount);
             p.setConsumerChannelCount(  consumerChannelCount);
             p.setConsumerMsgCount(      consumerMsgCount);
@@ -221,6 +223,7 @@ private static Options getOptions() {
         options.addOption(new Option("m", "ptxsize",                true, "producer tx size"));
         options.addOption(new Option("n", "ctxsize",                true, "consumer tx size"));
         options.addOption(new Option("c", "confirm",                true, "max unconfirmed publishes"));
+        options.addOption(new Option("ct", "confirmTimeout",        true, "waiting timeout for unconfirmed publishes before failing (in seconds)"));
         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, "consumer prefetch count"));
@@ -256,6 +259,10 @@ private static int intArg(CommandLine cmd, char opt, int def) {
         return Integer.parseInt(cmd.getOptionValue(opt, Integer.toString(def)));
     }
 
+    private static int intArg(CommandLine cmd, String opt, int def) {
+        return Integer.parseInt(cmd.getOptionValue(opt, Integer.toString(def)));
+    }
+
     private static float floatArg(CommandLine cmd, char opt, float def) {
         return Float.parseFloat(cmd.getOptionValue(opt, Float.toString(def)));
     }

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

@@ -27,6 +27,7 @@
 import java.util.TreeSet;
 import java.util.UUID;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
 
 public class Producer extends ProducerConsumerBase implements Runnable, ReturnListener,
         ConfirmListener
@@ -46,13 +47,14 @@ public class Producer extends ProducerConsumerBase implements Runnable, ReturnLi
     private final MessageBodySource messageBodySource;
 
     private Semaphore confirmPool;
+    private int confirmTimeout;
     private final SortedSet<Long> unconfirmedSet =
         Collections.synchronizedSortedSet(new TreeSet<Long>());
 
     public Producer(Channel channel, String exchangeName, String id, boolean randomRoutingKey,
                     List<?> flags, int txSize,
                     float rateLimit, int msgLimit, int timeLimit,
-                    long confirm, MessageBodySource messageBodySource, Stats stats)
+                    long confirm, int confirmTimeout, MessageBodySource messageBodySource, Stats stats)
         throws IOException {
 
         this.channel            = channel;
@@ -68,6 +70,7 @@ public Producer(Channel channel, String exchangeName, String id, boolean randomR
         this.messageBodySource = messageBodySource;
         if (confirm > 0) {
             this.confirmPool  = new Semaphore((int)confirm);
+            this.confirmTimeout = confirmTimeout;
         }
         this.stats        = stats;
     }
@@ -129,7 +132,16 @@ public void run() {
                    (msgLimit == 0 || msgCount < msgLimit)) {
                 delay(now);
                 if (confirmPool != null) {
-                    confirmPool.acquire();
+                    if (confirmTimeout < 0) {
+                        confirmPool.acquire();
+                    } else {
+                        boolean acquired = confirmPool.tryAcquire(confirmTimeout, TimeUnit.SECONDS);
+                        if (!acquired) {
+                            // waiting for too long, broker may be gone, stopping thread
+                            throw new RuntimeException("Waiting for publisher confirms for too long");
+                        }
+                    }
+
                 }
                 publish(messageBodySource.create(totalMsgCount));
                 totalMsgCount++;