Merge pull request #56 from pmackowski/retry-ack

Retry ack

Author: acogoluegnes

src/docs/asciidoc/api-guide.adoc

@@ -385,7 +385,11 @@ include::{test-examples}/ApiGuideReceiver.java[tag=auto-ack-retry-settings]
 <1> Retry acknowledgment for 20 seconds every 500 milliseconds on connection failure
 
 When using `Receiver#consumeManualAck`, acknowledgment is handled by the developer, who
-can do pretty anything they want on acknowledgment failure. It is possible to benefit from Reactor RabbitMQ
+can do pretty anything they want on acknowledgment failure.
+
+`AcknowledgableDelivery#ack` and `AcknowledgableDelivery#nack` methods handle retry internally
+based on `BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler` in the `ConsumeOptions`.
+Developer does not have to execute retry explicitly on acknowledgment failure and benefits from Reactor RabbitMQ
 retry support when acknowledging a message:
 
 [source,java,indent=0]
@@ -395,7 +399,6 @@ include::{test-examples}/ApiGuideReceiver.java[tag=manual-ack-retry]
 <1> Configure retry logic when exception occurs
 <2> Process message
 <3> Send acknowledgment after business processing
-<4> Execute retry on acknowledgment failure
 
 Note the exception handler is a `BiConsumer<Receiver.AcknowledgmentContext, Exception>`. This means
 acknowledgment failure can be handled in any way, here we choose to retry the acknowledgment.

src/main/java/reactor/rabbitmq/AcknowledgableDelivery.java

@@ -21,13 +21,16 @@
 
 import java.io.IOException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 
 /**
  * A RabbitMQ {@link Delivery} that can be manually acknowledged or rejected.
  */
 public class AcknowledgableDelivery extends Delivery {
 
     private final Channel channel;
+    private final BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler;
 
     private final AtomicBoolean notAckedOrNacked = new AtomicBoolean(true);
 
@@ -37,29 +40,28 @@ public class AcknowledgableDelivery extends Delivery {
      *
      * @param delivery
      * @param channel
+     * @param exceptionHandler
      */
-    public AcknowledgableDelivery(Delivery delivery, Channel channel) {
+    public AcknowledgableDelivery(Delivery delivery, Channel channel, BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler) {
         super(delivery.getEnvelope(), delivery.getProperties(), delivery.getBody());
         this.channel = channel;
+        this.exceptionHandler = exceptionHandler;
     }
 
     /**
      * Acknowledges this message if it has not been previously acked or nacked.
      * Subsequent calls to the method for previously acknowledged message will not produce errors and will simply
      * return instantly.
+     * In case of connection failure, {@link AcknowledgableDelivery#exceptionHandler} is executed.
      * @param multiple Defines whether all messages up to and including the supplied delivery tag should be
      * acknowledged or not.
      */
     public void ack(boolean multiple) {
         if (notAckedOrNacked.getAndSet(false)) {
             try {
-                channel.basicAck(getEnvelope().getDeliveryTag(), multiple);
-            } catch (RuntimeException e) {
-                notAckedOrNacked.set(true);
-                throw e;
-            } catch (IOException e) {
-                notAckedOrNacked.set(true);
-                throw new RabbitFluxException(e);
+                basicAck(multiple);
+            } catch (Exception e) {
+                retry(e, (delivery) -> delivery.basicAck(multiple));
             }
         }
     }
@@ -68,6 +70,7 @@ public void ack(boolean multiple) {
      * Acknowledges this message if it has not been previously acked or nacked.
      * Subsequent calls to the method for previously acknowledged message will not produce errors and will simply
      * return instantly.
+     * In case of connection failure, {@link AcknowledgableDelivery#exceptionHandler} is executed.
      */
     public void ack() {
         ack(false);
@@ -77,20 +80,17 @@ public void ack() {
      * Rejects this message if it has not been previously acked or nacked.
      * Subsequent calls to the method for previously acknowledged or rejected message will not produce errors and
      * will simply return instantly.
+     * In case of connection failure, {@link AcknowledgableDelivery#exceptionHandler} is executed.
      * @param multiple Defines whether all messages up to and including the supplied delivery tag should be
      * rejected or not.
      * @param requeue Defines if the message should be added to the queue again instead of being discarded.
      */
     public void nack(boolean multiple, boolean requeue) {
         if (notAckedOrNacked.getAndSet(false)) {
             try {
-                channel.basicNack(getEnvelope().getDeliveryTag(), multiple, requeue);
-            } catch (RuntimeException e) {
-                notAckedOrNacked.set(true);
-                throw e;
-            } catch (IOException e) {
-                notAckedOrNacked.set(true);
-                throw new RabbitFluxException(e);
+                basicNack(multiple, requeue);
+            } catch (Exception e) {
+                retry(e, (delivery) -> delivery.basicNack(multiple, requeue));
             }
         }
     }
@@ -99,8 +99,38 @@ public void nack(boolean multiple, boolean requeue) {
      * Rejects this message if it has not been previously acked or nacked.
      * Subsequent calls to the method for previously acknowledged or rejected message will not produce errors and
      * will simply return instantly.
+     * In case of connection failure, {@link AcknowledgableDelivery#exceptionHandler} is executed.
      */
     public void nack(boolean requeue) {
         nack(false, requeue);
     }
+
+    private void basicAck(boolean multiple) {
+        try {
+            channel.basicAck(getEnvelope().getDeliveryTag(), multiple);
+        } catch (RuntimeException e) {
+            throw e;
+        } catch (IOException e) {
+            throw new RabbitFluxException(e);
+        }
+    }
+
+    private void basicNack(boolean multiple, boolean requeue) {
+        try {
+            channel.basicNack(getEnvelope().getDeliveryTag(), multiple, requeue);
+        } catch (RuntimeException e) {
+            throw e;
+        } catch (IOException e) {
+            throw new RabbitFluxException(e);
+        }
+    }
+
+    private void retry(Exception e, Consumer<AcknowledgableDelivery> consumer) {
+        try {
+            exceptionHandler.accept(new Receiver.AcknowledgmentContext(this, consumer), e);
+        } catch (Exception e2) {
+            notAckedOrNacked.set(true);
+            throw e2;
+        }
+    }
 }

src/main/java/reactor/rabbitmq/ExceptionHandlers.java

@@ -128,7 +128,7 @@ public RetryAcknowledgmentExceptionHandler(Duration timeout, Duration waitingTim
         @Override
         public void accept(Receiver.AcknowledgmentContext acknowledgmentContext, Exception e) {
             retryTemplate.retry(() -> {
-                acknowledgmentContext.getDelivery().ack();
+                acknowledgmentContext.ackOrNack();
                 return null;
             }, e);
         }

src/main/java/reactor/rabbitmq/Receiver.java

@@ -34,6 +34,7 @@
 import java.io.IOException;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /**
@@ -134,13 +135,9 @@ public Flux<Delivery> consumeAutoAck(final String queue) {
 
     public Flux<Delivery> consumeAutoAck(final String queue, ConsumeOptions options) {
         // TODO why acking here and not just after emitter.next()?
-        return consumeManualAck(queue, options).doOnNext(msg -> {
-            try {
-                msg.ack();
-            } catch (Exception e) {
-                options.getExceptionHandler().accept(new AcknowledgmentContext(msg), e);
-            }
-        }).map(ackableMsg -> ackableMsg);
+        return consumeManualAck(queue, options)
+                .doOnNext(AcknowledgableDelivery::ack)
+                .map(ackableMsg -> ackableMsg);
     }
 
     public Flux<AcknowledgableDelivery> consumeManualAck(final String queue) {
@@ -157,7 +154,7 @@ public Flux<AcknowledgableDelivery> consumeManualAck(final String queue, Consume
                 }
 
                 DeliverCallback deliverCallback = (consumerTag, message) -> {
-                    AcknowledgableDelivery delivery = new AcknowledgableDelivery(message, channel);
+                    AcknowledgableDelivery delivery = new AcknowledgableDelivery(message, channel, options.getExceptionHandler());
                     if (options.getHookBeforeEmitBiFunction().apply(emitter.requestedFromDownstream(), delivery)) {
                         emitter.next(delivery);
                     }
@@ -215,13 +212,15 @@ public void close() {
     public static class AcknowledgmentContext {
 
         private final AcknowledgableDelivery delivery;
+        private final Consumer<AcknowledgableDelivery> consumer;
 
-        public AcknowledgmentContext(AcknowledgableDelivery delivery) {
+        public AcknowledgmentContext(AcknowledgableDelivery delivery, Consumer<AcknowledgableDelivery> consumer) {
             this.delivery = delivery;
+            this.consumer = consumer;
         }
 
-        public AcknowledgableDelivery getDelivery() {
-            return delivery;
+        public void ackOrNack() {
+            consumer.accept(delivery);
         }
     }
 

src/test/java/reactor/rabbitmq/ConnectionRecoveryTests.java

@@ -260,17 +260,12 @@ void consumeManualAckRetryOnAck() throws Exception {
         BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler = new ExceptionHandlers.RetryAcknowledgmentExceptionHandler(
             ofSeconds(5), ofMillis(100), ExceptionHandlers.CONNECTION_RECOVERY_PREDICATE
         );
-        receiver.consumeManualAck("whatever")
+        receiver.consumeManualAck("whatever", new ConsumeOptions().exceptionHandler(exceptionHandler))
             .subscribe(msg -> {
                 // do business stuff
                 // ...
-                try {
-                    // trying to ack
-                    msg.ack();
-                } catch (Exception e) {
-                    // when ack-ing fail, retry-ing
-                    exceptionHandler.accept(new Receiver.AcknowledgmentContext(msg), e);
-                }
+                // trying to ack
+                msg.ack();
                 ackedMessages.incrementAndGet();
             });
 

src/test/java/reactor/rabbitmq/RabbitFluxTests.java

@@ -47,11 +47,13 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.*;
 import static org.junit.jupiter.params.provider.Arguments.of;
 import static org.mockito.ArgumentMatchers.anyBoolean;
@@ -148,10 +150,11 @@ public void receiverCloseIsIdempotent() throws Exception {
     @Test
     public void acknowledgableDeliveryAckNackIsIdempotent() throws Exception {
         Channel channel = mock(Channel.class);
+        BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler = mock(ExceptionHandlers.RetryAcknowledgmentExceptionHandler.class);
         doNothing().doThrow(new IOException()).when(channel).basicAck(anyLong(), anyBoolean());
         doThrow(new IOException()).when(channel).basicNack(anyLong(), anyBoolean(), anyBoolean());
         AcknowledgableDelivery msg = new AcknowledgableDelivery(
-                new Delivery(new Envelope(0, true, null, null), null, null), channel
+                new Delivery(new Envelope(0, true, null, null), null, null), channel, exceptionHandler
         );
 
         msg.ack();
@@ -161,8 +164,44 @@ public void acknowledgableDeliveryAckNackIsIdempotent() throws Exception {
 
         verify(channel, times(1)).basicAck(anyLong(), anyBoolean());
         verify(channel, never()).basicNack(anyLong(), anyBoolean(), anyBoolean());
+        verify(exceptionHandler, never()).accept(any(Receiver.AcknowledgmentContext.class), any(Exception.class));
     }
 
+    @Test
+    public void acknowledgableDeliveryWithSuccessfulRetry() throws Exception {
+        Channel channel = mock(Channel.class);
+        BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler =
+                (acknowledgmentContext, e) -> acknowledgmentContext.ackOrNack();
+
+        doThrow(new IOException()).doNothing().when(channel).basicAck(anyLong(), anyBoolean());
+
+        AcknowledgableDelivery msg = new AcknowledgableDelivery(
+                new Delivery(new Envelope(0, true, null, null), null, null), channel, exceptionHandler
+        );
+
+        msg.ack();
+
+        verify(channel, times(2)).basicAck(anyLong(), anyBoolean());
+    }
+
+    @Test
+    public void acknowledgableDeliveryWithUnsuccessfulRetry() throws Exception {
+        Channel channel = mock(Channel.class);
+        BiConsumer<Receiver.AcknowledgmentContext, Exception> exceptionHandler =
+                (acknowledgmentContext, e) -> acknowledgmentContext.ackOrNack();
+
+        doThrow(new RuntimeException("exc")).when(channel).basicAck(anyLong(), anyBoolean());
+
+        AcknowledgableDelivery msg = new AcknowledgableDelivery(
+                new Delivery(new Envelope(0, true, null, null), null, null), channel, exceptionHandler
+        );
+
+        assertThatThrownBy(msg::ack).hasMessage("exc");
+
+        verify(channel, times(2)).basicAck(anyLong(), anyBoolean());
+    }
+
+
     @Test
     public void receiverConsumeNoAck() throws Exception {
         Channel channel = connection.createChannel();

src/test/java/reactor/rabbitmq/docs/ApiGuideReceiver.java

@@ -91,15 +91,12 @@ void manualAckRetry() {
                 Duration.ofSeconds(20), Duration.ofMillis(500),
                 ExceptionHandlers.CONNECTION_RECOVERY_PREDICATE
         );
-        receiver.consumeManualAck("queue")
-            .subscribe(msg -> {
-                // ...                                                                   // <2>
-                try {
-                    msg.ack();                                                           // <3>
-                } catch (Exception e) {
-                    exceptionHandler.accept(new Receiver.AcknowledgmentContext(msg), e); // <4>
-                }
-            });
+        receiver.consumeManualAck("queue",
+                    new ConsumeOptions().exceptionHandler(exceptionHandler))
+                .subscribe(msg -> {
+                    // ...                                                              // <2>
+                    msg.ack();                                                          // <3>
+                });
         // end::manual-ack-retry[]
     }