WaitForConfirms with timeouts

Author: Emile Joubert

src/com/rabbitmq/client/Channel.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.concurrent.TimeoutException;
 
 import com.rabbitmq.client.AMQP.BasicProperties;
 import com.rabbitmq.client.AMQP.Exchange;
@@ -726,13 +727,31 @@ void basicNack(long deliveryTag, boolean multiple, boolean requeue)
      */
     boolean waitForConfirms() throws InterruptedException;
 
+    /**
+     * Wait until all messages published since the last call have been
+     * either ack'd or nack'd by the broker; or until timeout elapses.
+     * If the timeout expires a TimeoutException is thrown.  When
+     * called on a non-Confirm channel, waitForConfirms returns true
+     * immediately.
+     * @return whether all the messages were ack'd (and none were nack'd)
+     */
+    boolean waitForConfirms(long timeout) throws InterruptedException, TimeoutException;
+
     /** Wait until all messages published since the last call have
      * been either ack'd or nack'd by the broker.  If any of the
      * messages were nack'd, waitForConfirmsOrDie will throw an
      * IOException.  When called on a non-Confirm channel, it will
      * return immediately. */
     void waitForConfirmsOrDie() throws IOException, InterruptedException;
 
+    /** Wait until all messages published since the last call have
+     * been either ack'd or nack'd by the broker; or until timeout elapses.
+     * If the timeout expires a TimeoutException is thrown.  If any of the
+     * messages were nack'd, waitForConfirmsOrDie will throw an
+     * IOException.  When called on a non-Confirm channel, it will
+     * return immediately. */
+    void waitForConfirmsOrDie(long timeout) throws IOException, InterruptedException, TimeoutException;
+
     /**
      * Asynchronously send a method over this channel.
      * @param method method to transmit over this channel.

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

@@ -167,6 +167,17 @@ public void clearConfirmListeners() {
     public boolean waitForConfirms()
         throws InterruptedException
     {
+        boolean confirms = false;
+        try {
+            confirms = waitForConfirms(0L);
+        } catch (TimeoutException e) { }
+        return confirms;
+    }
+
+    /** {@inheritDoc} */
+    public boolean waitForConfirms(long timeout)
+            throws InterruptedException, TimeoutException {
+        long startTime = System.currentTimeMillis();
         synchronized (unconfirmedSet) {
             while (true) {
                 if (getCloseReason() != null) {
@@ -177,7 +188,16 @@ public boolean waitForConfirms()
                     onlyAcksReceived = true;
                     return aux;
                 }
-                unconfirmedSet.wait();
+                if (timeout == 0L) {
+                    unconfirmedSet.wait();
+                } else {
+                    long elapsed = System.currentTimeMillis() - startTime;
+                    if (timeout > elapsed) {
+                        unconfirmedSet.wait(timeout - elapsed);
+                    } else {
+                        throw new TimeoutException();
+                    }
+                }
             }
         }
     }
@@ -186,7 +206,16 @@ public boolean waitForConfirms()
     public void waitForConfirmsOrDie()
         throws IOException, InterruptedException
     {
-        if (!waitForConfirms()) {
+        try {
+            waitForConfirmsOrDie(0L);
+        } catch (TimeoutException e) { }
+    }
+
+    /** {@inheritDoc} */
+    public void waitForConfirmsOrDie(long timeout)
+        throws IOException, InterruptedException, TimeoutException
+    {
+        if (!waitForConfirms(timeout)) {
             close(AMQP.REPLY_SUCCESS, "NACKS RECEIVED", true, null, false);
             throw new IOException("nacks received");
         }