Sender publish confirm backpressure

Author: flooq

src/jmh/java/reactor/rabbitmq/SenderMaxInFlightBenchmark.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright (c) 2019 Pivotal Software Inc, All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package reactor.rabbitmq;
+
+import com.rabbitmq.client.Connection;
+import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
+import reactor.core.publisher.Flux;
+
+import java.util.concurrent.TimeUnit;
+
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@State(Scope.Benchmark)
+@Warmup(iterations = 1, time = 5, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 1, time = 5, timeUnit = TimeUnit.SECONDS)
+@Fork(1)
+@Threads(2)
+public class SenderMaxInFlightBenchmark {
+
+    Connection connection;
+    Sender sender;
+    String queue;
+    Flux<OutboundMessage> msgFlux;
+
+    @Param({"1", "10", "100", "1000"})
+    public int nbMessages;
+
+    @Param({"1", "256", Integer.MAX_VALUE + ""})
+    public int maxInFlight;
+
+    @Setup
+    public void setupConnection() throws Exception {
+        connection = SenderBenchmarkUtils.newConnection();
+    }
+
+    @TearDown
+    public void closeConnection() throws Exception {
+        connection.close();
+    }
+
+    @Setup(Level.Iteration)
+    public void setupSender() throws Exception {
+        queue = SenderBenchmarkUtils.declareQueue(connection);
+        sender = RabbitFlux.createSender();
+        msgFlux = SenderBenchmarkUtils.outboundMessageFlux(queue, nbMessages);
+    }
+
+    @TearDown(Level.Iteration)
+    public void tearDownSender() throws Exception {
+        SenderBenchmarkUtils.deleteQueue(connection, queue);
+        if (sender != null) {
+            sender.close();
+        }
+    }
+
+    @Benchmark
+    public void sendWithPublishConfirmsMaxInFlight(Blackhole blackhole) {
+        blackhole.consume(sender.sendWithPublishConfirms(msgFlux, new SendOptions().maxInFlight(maxInFlight)).blockLast());
+    }
+
+}

src/main/java/reactor/rabbitmq/SendOptions.java

@@ -21,7 +21,6 @@
 import reactor.core.publisher.SignalType;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
-import reactor.util.concurrent.Queues;
 
 import java.time.Duration;
 import java.util.function.BiConsumer;
@@ -41,14 +40,14 @@ public class SendOptions {
      *
      * @since 1.1.1
      */
-    private int maxInFlight = Queues.SMALL_BUFFER_SIZE;
+    private Integer maxInFlight;
 
     /**
      * The scheduler used for publishing send results.
      *
      * @since 1.1.1
      */
-    private Scheduler scheduler = Schedulers.newSingle("rabbitmq-sender-publishing");
+    private Scheduler scheduler = Schedulers.immediate();
 
     /**
      * Channel mono to use to send messages.
@@ -69,9 +68,9 @@ public class SendOptions {
      * from the outbound record publisher while publisher confirms are pending.
      *
      * @return maximum number of in-flight records
-     * @since 1.1.18
+     * @since 1.1.1
      */
-    public int getMaxInFlight() {
+    public Integer getMaxInFlight() {
         return maxInFlight;
     }
 
@@ -89,25 +88,29 @@ public SendOptions maxInFlight(int maxInFlight) {
     }
 
     /**
-     * The scheduler used for publishing send results.
+     * Set the maximum number of in-flight records that are fetched
+     * from the outbound record publisher while publisher confirms are pending.
+     * Results are run on a supplied {@link Scheduler}
      *
-     * @return scheduler used for publishing send results
+     * @param maxInFlight
+     * @param scheduler
+     * @return this {@link SendOptions} instance
      * @since 1.1.1
      */
-    public Scheduler getScheduler() {
-        return scheduler;
+    public SendOptions maxInFlight(int maxInFlight, Scheduler scheduler) {
+        this.maxInFlight = maxInFlight;
+        this.scheduler = scheduler;
+        return this;
     }
 
     /**
-     * Sets the scheduler used for publishing send results.
+     * The scheduler used for publishing send results.
      *
-     * @param scheduler
-     * @return this {@link SendOptions} instance
+     * @return scheduler used for publishing send results
      * @since 1.1.1
      */
-    public SendOptions scheduler(Scheduler scheduler) {
-        this.scheduler = scheduler;
-        return this;
+    public Scheduler getScheduler() {
+        return scheduler;
     }
 
     public BiConsumer<Sender.SendContext, Exception> getExceptionHandler() {

src/main/java/reactor/rabbitmq/Sender.java

@@ -143,7 +143,7 @@ public Flux<OutboundMessageResult> sendWithPublishConfirms(Publisher<OutboundMes
         final Mono<? extends Channel> currentChannelMono = getChannelMono(options);
         final BiConsumer<SignalType, Channel> channelCloseHandler = getChannelCloseHandler(options);
 
-        return currentChannelMono.map(channel -> {
+        Flux<OutboundMessageResult> result = currentChannelMono.map(channel -> {
                 try {
                     channel.confirmSelect();
                 } catch (IOException e) {
@@ -161,8 +161,12 @@ public Flux<OutboundMessageResult> sendWithPublishConfirms(Publisher<OutboundMes
                     // so we need to complete in another thread
                     channelCloseThreadPool.execute(() -> channelCloseHandler.accept(signalType, channel));
                 }
-            }))
-            .publishOn(sendOptions.getScheduler(), sendOptions.getMaxInFlight());
+            }));
+
+        if (sendOptions.getMaxInFlight() != null) {
+            result = result.publishOn(sendOptions.getScheduler(), sendOptions.getMaxInFlight());
+        }
+        return result;
     }
 
     // package-protected for testing
@@ -525,7 +529,7 @@ private static class PublishConfirmSubscriber implements
 
         private final AtomicReference<SubscriberState> state = new AtomicReference<>(SubscriberState.INIT);
 
-        private final AtomicReference<Throwable> firstException = new AtomicReference<Throwable>();
+        private final AtomicReference<Throwable> firstException = new AtomicReference<>();
 
         private final ConcurrentNavigableMap<Long, OutboundMessage> unconfirmed = new ConcurrentSkipListMap<>();
 
@@ -535,7 +539,7 @@ private static class PublishConfirmSubscriber implements
 
         private final BiConsumer<SendContext, Exception> exceptionHandler;
 
-        Subscription s;
+        private Subscription subscription;
 
         private PublishConfirmSubscriber(Channel channel, Subscriber<? super OutboundMessageResult> subscriber, SendOptions options) {
             this.channel = channel;
@@ -545,12 +549,12 @@ private PublishConfirmSubscriber(Channel channel, Subscriber<? super OutboundMes
 
         @Override
         public void request(long n) {
-            s.request(n);
+            subscription.request(n);
         }
 
         @Override
         public void cancel() {
-            s.cancel();
+            subscription.cancel();
         }
 
         @Override
@@ -594,8 +598,8 @@ private void handleAckNack(long deliveryTag, boolean multiple, boolean ack) {
                 }
             });
             state.set(SubscriberState.ACTIVE);
-            if (Operators.validate(this.s, subscription)) {
-                this.s = subscription;
+            if (Operators.validate(this.subscription, subscription)) {
+                this.subscription = subscription;
                 subscriber.onSubscribe(this);
             }
         }

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

@@ -708,6 +708,55 @@ public void publishConfirms() throws Exception {
         assertEquals(nbMessages, counter.get());
     }
 
+    @Test
+    public void publishConfirmsBackpressure() throws Exception {
+        int nbMessages = 10;
+        int subscriberRequest = 3;
+        CountDownLatch consumedLatch = new CountDownLatch(subscriberRequest);
+        CountDownLatch confirmedLatch = new CountDownLatch(subscriberRequest);
+        AtomicInteger counter = new AtomicInteger();
+        Channel channel = connection.createChannel();
+        channel.basicConsume(queue, true, (consumerTag, delivery) -> {
+            counter.incrementAndGet();
+            consumedLatch.countDown();
+        }, consumerTag -> {
+        });
+
+        Flux<OutboundMessage> msgFlux = Flux.range(0, nbMessages).map(i -> new OutboundMessage("", queue, "".getBytes()));
+
+        sender = createSender();
+        sender.sendWithPublishConfirms(msgFlux).subscribe(new BaseSubscriber<OutboundMessageResult>() {
+                @Override
+                protected void hookOnSubscribe(Subscription subscription) {
+                    subscription.request(subscriberRequest);
+                }
+
+                @Override
+                protected void hookOnNext(OutboundMessageResult outboundMessageResult) {
+                    if (outboundMessageResult.getOutboundMessage() != null) {
+                        confirmedLatch.countDown();
+                    }
+                }
+        });
+
+        assertTrue(consumedLatch.await(1, TimeUnit.SECONDS));
+        assertTrue(confirmedLatch.await(1, TimeUnit.SECONDS));
+        assertEquals(subscriberRequest, counter.get());
+    }
+
+    @Test
+    public void publishConfirmsEmptyPublisher() throws Exception {
+        CountDownLatch finallyLatch = new CountDownLatch(1);
+        Flux<OutboundMessage> msgFlux = Flux.empty();
+
+        sender = createSender();
+        sender.sendWithPublishConfirms(msgFlux)
+                .doFinally(signalType -> finallyLatch.countDown())
+                .subscribe();
+
+        assertTrue(finallyLatch.await(1, TimeUnit.SECONDS));
+    }
+
     @Test
     void publishConfirmsMaxInFlight() throws InterruptedException {
         int maxConcurrency = 4;