Servicebus track2 sync queue up multiple receive calls

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

@@ -0,0 +1,168 @@
+package com.azure.messaging.servicebus;
+
+import com.azure.core.util.logging.ClientLogger;
+import org.reactivestreams.Subscription;
+import reactor.core.Disposable;
+import reactor.core.publisher.BaseSubscriber;
+import reactor.core.publisher.Mono;
+
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+ class LongLivedMessageSubscriber extends BaseSubscriber<ServiceBusReceivedMessageContext> {
+
+    private final ClientLogger logger = new ClientLogger(SynchronousMessageSubscriber.class);
+    private final AtomicBoolean isDisposed = new AtomicBoolean(false);
+    private final long prefetch;
+    private final AtomicInteger wip = new AtomicInteger();
+
+    private volatile Subscription subscription;
+
+    private Queue<SynchronousReceiveWork> workQueue = new ConcurrentLinkedQueue<>();
+    private SynchronousReceiveWork currentWork = null;
+    private Disposable timeoutOperation;
+     private Disposable drainQueueDisposable;
+
+
+    LongLivedMessageSubscriber(long prefetch) {
+        this.prefetch = prefetch;
+    }
+
+    /**
+     * On an initial subscription, will take the first work item, and request that amount of work for it.
+     *
+     * @param subscription Subscription for upstream.
+     */
+    @Override
+    protected void hookOnSubscribe(Subscription subscription) {
+        if (this.subscription == null) {
+            this.subscription = subscription;
+        }
+    }
+
+    /**
+     * Publishes the event to the current {@link SynchronousReceiveWork}. If that work item is complete, will dispose of
+     * the subscriber.
+     *
+     * @param message Event to publish.
+     */
+    @Override
+    protected void hookOnNext(ServiceBusReceivedMessageContext message) {
+
+        currentWork.next(message);
+
+        logger.verbose("[{}] received message with Sequence Number [{}].", currentWork.getId(), message.getMessage().getSequenceNumber());
+
+        if (currentWork.isTerminal()) {
+            logger.info("[{}] Completed. Closing Flux and cancelling subscription.", currentWork.getId());
+            completeCurrentWork(currentWork);
+        }
+    }
+
+    private void completeCurrentWork(SynchronousReceiveWork currentWork) {
+
+        if (isTerminated()) {
+            return;
+        }
+
+        currentWork.complete();
+        logger.verbose("[{}] work completed.", currentWork.getId());
+
+        if (timeoutOperation != null && !timeoutOperation.isDisposed()) {
+            timeoutOperation.dispose();
+        }
+        if (drainQueueDisposable != null && !drainQueueDisposable.isDisposed()) {
+            drainQueueDisposable.dispose();
+        }
+
+        if (wip.decrementAndGet() != 0) {
+            logger.warning("There is another worker in drainLoop. But there should only be 1 worker. Value:"+wip.get());
+        }
+
+        // After current work finished and there more receive requests
+        if (workQueue.size() > 0 ) {
+            drain();
+        }
+    }
+
+    void queueWork(SynchronousReceiveWork work) {
+        workQueue.add(work);
+        drain();
+    }
+
+    private void drain() {
+        // If someone is already in this loop, then we are already clearing the queue.
+        if (!wip.compareAndSet(0, 1)) {
+            return;
+        }
+        // Drain queue..
+        drainQueueDisposable = Mono.just(true)
+            .subscribe(l -> {
+                drainQueue();
+            });
+    }
+
+    private void drainQueue() {
+        if (isTerminated()) {
+            return;
+        }
+        currentWork = workQueue.poll();
+        if (currentWork == null) {
+            return;
+        }
+        subscription.request(currentWork.getNumberOfEvents());
+
+        // timer to complete the current in case of timeout trigger
+        timeoutOperation = Mono.delay(currentWork.getTimeout())
+            .subscribe(l -> {
+                if (!currentWork.isTerminal()) {
+                    completeCurrentWork(currentWork);
+                }
+            });
+    }
+
+
+    @Override
+    protected void hookOnComplete() {
+        logger.info("[{}] Completed. No events to listen to.", currentWork.getId());
+        dispose();
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    protected void hookOnError(Throwable throwable) {
+        logger.error("[{}] Errors occurred upstream", currentWork.getId(), throwable);
+        currentWork.error(throwable);
+        dispose();
+    }
+
+    @Override
+    protected void hookOnCancel() {
+        dispose();
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void dispose() {
+        if (isDisposed.getAndSet(true)) {
+            return;
+        }
+
+        currentWork.complete();
+        subscription.cancel();
+        if (timeoutOperation != null && !timeoutOperation.isDisposed() ) {
+            timeoutOperation.dispose();
+        }
+        super.dispose();
+    }
+
+    private boolean isTerminated(){
+        return isDisposed.get();
+    }
+}

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

@@ -41,9 +41,12 @@ public final class ServiceBusReceiverClient implements AutoCloseable {
         .setIsAutoCompleteEnabled(false)
         .setMaxAutoLockRenewalDuration(Duration.ZERO);
 
-    private final AtomicReference<EmitterProcessor<ServiceBusReceivedMessageContext>> messageProcessor =
+    private final AtomicReference<EmitterProcessor<ServiceBusReceivedMessageContext>> messageSource =
         new AtomicReference<>();
 
+    /* To hold each receive work item to be processed.*/
+    private final AtomicReference<LongLivedMessageSubscriber> longLivedMessageSubscriber = new AtomicReference<>();
+
     /**
      * Creates a synchronous receiver given its asynchronous counterpart.
      *
@@ -479,7 +482,7 @@ public IterableStream<ServiceBusReceivedMessageContext> receive(int maxMessages)
      *
      * @throws IllegalArgumentException if {@code maxMessages} or {@code maxWaitTime} is zero or a negative value.
      */
-    public IterableStream<ServiceBusReceivedMessageContext> receive(int maxMessages, Duration maxWaitTime) {
+    public synchronized IterableStream<ServiceBusReceivedMessageContext> receive(int maxMessages, Duration maxWaitTime) {
         if (maxMessages <= 0) {
             throw logger.logExceptionAsError(new IllegalArgumentException(
                 "'maxMessages' cannot be less than or equal to 0. maxMessages: " + maxMessages));
@@ -613,9 +616,9 @@ public void setSessionState(String sessionId, byte[] sessionState) {
     public void close() {
         asyncClient.close();
 
-        EmitterProcessor<ServiceBusReceivedMessageContext> processor = messageProcessor.getAndSet(null);
-        if (processor != null) {
-            processor.onComplete();
+        LongLivedMessageSubscriber messageSubscriber = longLivedMessageSubscriber.getAndSet(null);
+        if (messageSubscriber != null && !messageSubscriber.isDisposed()) {
+            messageSubscriber.dispose();
         }
     }
 
@@ -624,23 +627,23 @@ public void close() {
      * entity.
      */
     private void queueWork(int maximumMessageCount, Duration maxWaitTime,
-        FluxSink<ServiceBusReceivedMessageContext> emitter) {
+                                   FluxSink<ServiceBusReceivedMessageContext> emitter) {
         synchronized (lock) {
             final long id = idGenerator.getAndIncrement();
-            EmitterProcessor<ServiceBusReceivedMessageContext> emitterProcessor = messageProcessor.get();
-
             final SynchronousReceiveWork work = new SynchronousReceiveWork(id, maximumMessageCount, maxWaitTime,
                 emitter);
-            final SynchronousMessageSubscriber syncSubscriber = new SynchronousMessageSubscriber(work);
-            logger.info("[{}]: Started synchronous message subscriber.", id);
 
-            if (emitterProcessor == null) {
-                emitterProcessor = this.asyncClient.receive(DEFAULT_RECEIVE_OPTIONS)
-                    .subscribeWith(EmitterProcessor.create(asyncClient.getReceiverOptions().getPrefetchCount(), false));
-                messageProcessor.set(emitterProcessor);
+            LongLivedMessageSubscriber messageSubscriber = longLivedMessageSubscriber.get();
+            if (messageSubscriber == null) {
+                messageSubscriber = asyncClient.receive(DEFAULT_RECEIVE_OPTIONS)
+                    .subscribeWith(new LongLivedMessageSubscriber(asyncClient.getReceiverOptions().getPrefetchCount()));
+                logger.verbose("Created source for receiving messages from [{}]", asyncClient.getEntityPath());
+                longLivedMessageSubscriber.set(messageSubscriber);
             }
 
-            emitterProcessor.subscribe(syncSubscriber);
+            messageSubscriber.queueWork(work);
+            logger.verbose("[{}] Receive request queued up.", work.getId());
         }
     }
+
 }

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

@@ -98,6 +98,10 @@ public void dispose() {
         super.dispose();
     }
 
+    SynchronousReceiveWork getWork() {
+        return work;
+    }
+
     private static final class ReceiveTimeoutTask extends TimerTask {
         private final ClientLogger logger = new ClientLogger(ReceiveTimeoutTask.class);
         private final long workId;

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

@@ -47,7 +47,7 @@
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
 public abstract class IntegrationTestBase extends TestBase {
-    protected static final Duration OPERATION_TIMEOUT = Duration.ofSeconds(30);
+    protected static final Duration OPERATION_TIMEOUT = Duration.ofSeconds(10 * 30);
     protected static final Duration TIMEOUT = Duration.ofSeconds(60);
     protected static final AmqpRetryOptions RETRY_OPTIONS = new AmqpRetryOptions().setTryTimeout(TIMEOUT);
     protected final ClientLogger logger;

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

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -101,11 +102,11 @@ protected void afterTest() {
      */
     @MethodSource("messagingEntityWithSessions")
     @ParameterizedTest
-    void receiveByTwoSubscriber(MessagingEntityType entityType, boolean isSessionEnabled) {
+    void multipleReceiveByOneSubscriber(MessagingEntityType entityType, boolean isSessionEnabled) {
         // Arrange
         setSenderAndReceiver(entityType, isSessionEnabled);
-        final int maxMessages = 1;
-        final int totalReceive = 3;
+        final int maxMessages = 2;
+        final int totalReceive = 2;
         final Duration shortTimeOut = Duration.ofSeconds(8);
 
         final String messageId = UUID.randomUUID().toString();
@@ -136,6 +137,63 @@ void receiveByTwoSubscriber(MessagingEntityType entityType, boolean isSessionEna
         assertEquals(totalReceive * maxMessages, totalReceivedCount);
     }
 
+    /**
+     * Verifies that we can only call receive() multiple times.
+     */
+    @MethodSource("messagingEntityWithSessions")
+    @ParameterizedTest
+    void parallelReceiveByOneSubscriber(MessagingEntityType entityType, boolean isSessionEnabled) throws InterruptedException {
+        // Arrange
+        setSenderAndReceiver(entityType, isSessionEnabled);
+        final int maxMessagesEachReceive = 2;
+        final int totalReceiver = 10;
+        final Duration shortTimeOut = Duration.ofSeconds(8);
+
+        final String messageId = UUID.randomUUID().toString();
+        final ServiceBusMessage message = getMessage(messageId, isSessionEnabled);
+
+        for (int i = 0; i < totalReceiver * maxMessagesEachReceive; ++i) {
+            sendMessage(message);
+        }
+
+        // Act & Assert
+        AtomicInteger totalReceivedMessages = new AtomicInteger();
+        List<Thread> receiverThreads = new ArrayList<>();
+        for (int i = 0; i< totalReceiver; ++i) {
+            int finalI = i;
+            Thread thread = new Thread(() -> {
+                IterableStream<ServiceBusReceivedMessageContext> messages1 = receiver.receive(maxMessagesEachReceive, shortTimeOut);
+                int receivedMessageCount = 0;
+                long lastSequenceReceiver = 0;
+                for (ServiceBusReceivedMessageContext receivedMessage : messages1) {
+                    logger.verbose("Receiver [{}}] Received SQ: " ,(finalI + 1)  ,receivedMessage.getMessage().getSequenceNumber());
+                    assertMessageEquals(receivedMessage, messageId, isSessionEnabled);
+                    receiver.complete(receivedMessage.getMessage());
+                    assertTrue( receivedMessage.getMessage().getSequenceNumber() > lastSequenceReceiver);
+                    lastSequenceReceiver = receivedMessage.getMessage().getSequenceNumber();
+                    messagesPending.decrementAndGet();
+                    ++receivedMessageCount;
+                }
+                totalReceivedMessages.addAndGet(receivedMessageCount);
+                assertEquals(maxMessagesEachReceive, receivedMessageCount);
+                logger.verbose("!!!! Receiver [{}}] . Test Completed receivers  ", (finalI + 1) );
+            });
+            receiverThreads.add(thread);
+        }
+
+        receiverThreads.forEach(t -> t.start());
+
+        TimeUnit.SECONDS.sleep(1);
+        receiverThreads.forEach(t -> {
+            try {
+                t.join();
+            } catch (InterruptedException e) {
+                e.printStackTrace();
+            }
+        });
+        assertEquals( totalReceiver * maxMessagesEachReceive, totalReceivedMessages.get());
+    }
+
     /**
      * Verifies that we can send and receive two messages.
      */