[Service Bus] Allow 0 prefetch and dynamically use batch size to request link credits

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/AmqpReceiveLink.java

@@ -25,7 +25,7 @@ public interface AmqpReceiveLink extends AmqpLink {
     Flux<Message> receive();
 
     /**
-     * Adds the specified number of credits to the link.
+     * Schedule to adds the specified number of credits to the link.
      *
      * The number of link credits initialises to zero. It is the application's responsibility to call this method to
      * allow the receiver to receive {@code credits} more deliveries.
@@ -34,6 +34,21 @@ public interface AmqpReceiveLink extends AmqpLink {
      */
     void addCredits(int credits);
 
+    /**
+     * Adds the specified number of credits to the link.
+     *
+     * The number of link credits initialises to zero. It is the application's responsibility to call this method to
+     * allow the receiver to receive {@code credits} more deliveries.
+     *
+     * It will update the credits in local memory instantly so {@link #getCredits()} will get
+     * the updated credits immediately. But the service side may get the credits added with a latency.
+     * As a contrast, {@link #getCredits()} may return an unchanged value for a short while after
+     * {@link #addCredits(int)} is called to schedule the credit addition and before the job dispatcher executes it.
+     *
+     * @param credits Number of credits to add to the receive link.
+     */
+    void addCreditsInstantly(int credits);
+
     /**
      * Gets the current number of credits this link has.
      *

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorReceiver.java

@@ -107,6 +107,11 @@ public void addCredits(int credits) {
         }
     }
 
+    @Override
+    public void addCreditsInstantly(int credits) {
+        receiver.flow(credits);
+    }
+
     @Override
     public int getCredits() {
         return receiver.getRemoteCredit();

sdk/servicebus/azure-messaging-servicebus/pom.xml

@@ -52,7 +52,7 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core-amqp</artifactId>
-      <version>1.7.0-beta.2</version> <!-- {x-version-update;beta_com.azure:azure-core-amqp;dependency} -->
+      <version>1.7.0-beta.3</version> <!-- {x-version-update;unreleased_com.azure:azure-core-amqp;dependency} -->
     </dependency>
     <dependency>
       <groupId>com.azure</groupId>

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusAsyncConsumer.java

@@ -34,9 +34,7 @@ class ServiceBusAsyncConsumer implements AutoCloseable {
         this.linkProcessor = linkProcessor;
         this.messageSerializer = messageSerializer;
         this.processor = linkProcessor
-            .map(message -> this.messageSerializer.deserialize(message, ServiceBusReceivedMessage.class))
-            .publish(receiverOptions.getPrefetchCount())
-            .autoConnect(1);
+            .map(message -> this.messageSerializer.deserialize(message, ServiceBusReceivedMessage.class));
     }
 
     /**

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java

@@ -94,7 +94,7 @@ public final class ServiceBusClientBuilder {
 
     // Using 0 pre-fetch count for both receive modes, to avoid message lock lost exceptions in application
     // receiving messages at a slow rate. Applications can set it to a higher value if they need better performance.
-    private static final int DEFAULT_PREFETCH_COUNT = 1;
+    private static final int DEFAULT_PREFETCH_COUNT = 0;
     private static final String NAME_KEY = "name";
     private static final String VERSION_KEY = "version";
     private static final String UNKNOWN = "UNKNOWN";
@@ -670,11 +670,13 @@ public ServiceBusSessionProcessorClientBuilder maxConcurrentSessions(int maxConc
 
         /**
          * Sets the prefetch count of the processor. For both {@link ReceiveMode#PEEK_LOCK PEEK_LOCK} and {@link
-         * ReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} modes the default value is 1.
+         * ReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} modes the default value is 0.
          *
          * Prefetch speeds up the message flow by aiming to have a message readily available for local retrieval when
          * and before the application starts the processor.
          * Setting a non-zero value will prefetch that number of messages. Setting the value to zero turns prefetch off.
+         * Using a non-zero prefetch risks of losing messages even though it has better performance.
+         * @see <a href="https://docs.microsoft.com/en-us/azure/service-bus-messaging/service-bus-prefetch">Service Bus Prefetch</a>
          *
          * @param prefetchCount The prefetch count.
          *
@@ -1441,9 +1443,9 @@ ServiceBusReceiverAsyncClient buildAsyncClient(boolean isAutoCompleteAllowed) {
     }
 
     private void validateAndThrow(int prefetchCount) {
-        if (prefetchCount < 1) {
+        if (prefetchCount < 0) {
             throw logger.logExceptionAsError(new IllegalArgumentException(String.format(
-                "prefetchCount (%s) cannot be less than 1.", prefetchCount)));
+                "prefetchCount (%s) cannot be less than 0.", prefetchCount)));
         }
     }
 

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java

@@ -590,13 +590,14 @@ public void close() {
     private void queueWork(int maximumMessageCount, Duration maxWaitTime,
                            FluxSink<ServiceBusReceivedMessage> emitter) {
         final long id = idGenerator.getAndIncrement();
-        final SynchronousReceiveWork work = new SynchronousReceiveWork(id, maximumMessageCount, maxWaitTime, emitter);
-
+        int prefetch = asyncClient.getReceiverOptions().getPrefetchCount();
+        int toRequest = prefetch != 0 ? Math.min(maximumMessageCount, prefetch) : maximumMessageCount;
+        final SynchronousReceiveWork work = new SynchronousReceiveWork(id,
+            toRequest,
+            maxWaitTime, emitter);
         SynchronousMessageSubscriber messageSubscriber = synchronousMessageSubscriber.get();
         if (messageSubscriber == null) {
-            long prefetch = asyncClient.getReceiverOptions().getPrefetchCount();
-            SynchronousMessageSubscriber newSubscriber = new SynchronousMessageSubscriber(prefetch, work);
-
+            SynchronousMessageSubscriber newSubscriber = new SynchronousMessageSubscriber(toRequest, work);
             if (!synchronousMessageSubscriber.compareAndSet(null, newSubscriber)) {
                 newSubscriber.dispose();
                 SynchronousMessageSubscriber existing = synchronousMessageSubscriber.get();

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SynchronousMessageSubscriber.java

@@ -57,8 +57,6 @@ protected void hookOnSubscribe(Subscription subscription) {
 
         if (Operators.setOnce(UPSTREAM, this, subscription)) {
             this.subscription = subscription;
-            remaining.addAndGet(requested);
-            subscription.request(requested);
             subscriberInitialized = true;
             drain();
         } else {
@@ -140,7 +138,6 @@ private void drainQueue() {
 
             while ((currentWork = workQueue.peek()) != null
                 && (!currentWork.isProcessingStarted() || bufferMessages.size() > 0)) {
-
                 // Additional check for safety, but normally this work should never be terminal
                 if (currentWork.isTerminal()) {
                     // This work already finished by either timeout or no more messages to send, process next work.
@@ -155,6 +152,9 @@ private void drainQueue() {
                     // timer to complete the currentWork in case of timeout trigger
                     currentTimeoutOperation = getTimeoutOperation(currentWork);
                     currentWork.startedProcessing();
+                    final long calculatedRequest = currentWork.getNumberOfEvents() - remaining.get();
+                    remaining.addAndGet(calculatedRequest);
+                    subscription.request(calculatedRequest);
                 }
 
                 // Send messages to currentWork from buffer
@@ -174,15 +174,6 @@ private void drainQueue() {
                         currentTimeoutOperation.dispose();
                     }
                     logger.verbose("The work [{}] is complete.", currentWork.getId());
-                } else {
-                    // Since this work is not complete, find out how much we should request from upstream
-                    long creditToAdd = currentWork.getRemaining() - (remaining.get() + bufferMessages.size());
-                    if (creditToAdd > 0) {
-                        remaining.addAndGet(creditToAdd);
-                        subscription.request(creditToAdd);
-                        logger.verbose("Requesting [{}] from upstream for work [{}].", creditToAdd,
-                            currentWork.getId());
-                    }
                 }
             }
         }

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiveLinkProcessor.java

@@ -45,7 +45,6 @@ public class ServiceBusReceiveLinkProcessor extends FluxProcessor<ServiceBusRece
     private final Object queueLock = new Object();
     private final AtomicBoolean isTerminated = new AtomicBoolean();
     private final AtomicInteger retryAttempts = new AtomicInteger();
-    private final AtomicBoolean linkCreditsAdded = new AtomicBoolean();
     private final AtomicReference<String> linkName = new AtomicReference<>();
 
     // Queue containing all the prefetched messages.
@@ -199,12 +198,7 @@ public void onNext(ServiceBusReceiveLink next) {
             oldSubscription = currentLinkSubscriptions;
 
             currentLink = next;
-            next.setEmptyCreditListener(() -> {
-                final int creditsToAdd = getCreditsToAdd(0);
-                linkCreditsAdded.set(creditsToAdd > 0);
-
-                return creditsToAdd;
-            });
+            next.setEmptyCreditListener(() -> 0);
 
             currentLinkSubscriptions = Disposables.composite(
                 next.receive().publishOn(Schedulers.boundedElastic()).subscribe(message -> {
@@ -498,6 +492,9 @@ private void drainQueue() {
                     if (receiveMode != ReceiveMode.PEEK_LOCK) {
                         pendingMessages.decrementAndGet();
                     }
+                    if (prefetch > 0) { // re-fill messageQueue if there is prefetch configured.
+                        checkAndAddCredits(currentLink);
+                    }
                 } catch (Exception e) {
                     logger.error("Exception occurred while handling downstream onNext operation.", e);
                     throw logger.logExceptionAsError(Exceptions.propagate(
@@ -544,18 +541,12 @@ private void checkAndAddCredits(AmqpReceiveLink link) {
             return;
         }
 
-        // Credits have already been added to the link. We won't try again.
-        if (linkCreditsAdded.getAndSet(true)) {
-            return;
-        }
-
         final int credits = getCreditsToAdd(link.getCredits());
-        linkCreditsAdded.set(credits > 0);
 
         logger.info("Link credits to add. Credits: '{}'", credits);
 
         if (credits > 0) {
-            link.addCredits(credits);
+            link.addCreditsInstantly(credits);
         }
     }
 
@@ -570,21 +561,28 @@ private int getCreditsToAdd(int linkCredits) {
         }
 
         final int creditsToAdd;
-        if (messageQueue.isEmpty() && !hasBackpressure) {
-            creditsToAdd = prefetch;
+        final int expectedTotalCredit;
+        if (prefetch == 0) {
+            if (r <= Integer.MAX_VALUE) {
+                expectedTotalCredit = (int) r;
+            } else {
+                expectedTotalCredit = Integer.MAX_VALUE;
+            }
         } else {
-            synchronized (queueLock) {
-                final int queuedMessages = pendingMessages.get();
-                final int pending = queuedMessages + linkCredits;
-
-                if (hasBackpressure) {
-                    creditsToAdd = Math.max(Long.valueOf(r).intValue() - pending, 0);
-                } else {
-                    // If the queue has less than 1/3 of the prefetch, then add the difference to keep the queue full.
-                    creditsToAdd = minimumNumberOfMessages >= queuedMessages
-                        ? Math.max(prefetch - pending, 1)
-                        : 0;
-                }
+            expectedTotalCredit = prefetch;
+        }
+
+        synchronized (queueLock) {
+            final int queuedMessages = pendingMessages.get();
+            final int pending = queuedMessages + linkCredits;
+
+            if (hasBackpressure) {
+                creditsToAdd = Math.max(expectedTotalCredit - pending, 0);
+            } else {
+                // If the queue has less than 1/3 of the prefetch, then add the difference to keep the queue full.
+                creditsToAdd = minimumNumberOfMessages >= queuedMessages
+                    ? Math.max(expectedTotalCredit - pending, 0)
+                    : 0;
             }
         }
 

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusClientBuilderTest.java

@@ -203,7 +203,7 @@ void invalidPrefetch() {
             .receiveMode(ReceiveMode.PEEK_LOCK);
 
         // Act & Assert
-        assertThrows(IllegalArgumentException.class, () -> receiverBuilder.prefetchCount(0));
+        assertThrows(IllegalArgumentException.class, () -> receiverBuilder.prefetchCount(-1));
     }
 
     @MethodSource("getProxyConfigurations")