Update the synchronous work from queue when current work is completed (Azure#27573)

* Queue work when current work is not terminated

* Update work when request is 0

* Add syncSubscribed to sync subscribe

* Update CHANGELOG

* Add UT to test updating work when current work completed

Author: liukun-msft

sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- Fixed a bug that when the current `SynchronousReceiveWork` is completed, the queued `SynchronousReceiveWork` is not updated. ([#27578](https://github.com/Azure/azure-sdk-for-java/issues/27578))
 
 ### Other Changes
 

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

@@ -18,6 +18,7 @@
 import java.time.Duration;
 import java.time.OffsetDateTime;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
@@ -60,6 +61,8 @@ public final class ServiceBusReceiverClient implements AutoCloseable {
 
     /* To hold each receive work item to be processed.*/
     private final AtomicReference<SynchronousMessageSubscriber> synchronousMessageSubscriber = new AtomicReference<>();
+    /* To ensure synchronousMessageSubscriber is subscribed only once. */
+    private final AtomicBoolean syncSubscribed = new AtomicBoolean(false);
 
     /**
      * Creates a synchronous receiver given its asynchronous counterpart.
@@ -719,28 +722,36 @@ private void queueWork(int maximumMessageCount, Duration maxWaitTime,
 
         final long id = idGenerator.getAndIncrement();
         final SynchronousReceiveWork work = new SynchronousReceiveWork(id, maximumMessageCount, maxWaitTime, emitter);
-        final SynchronousMessageSubscriber messageSubscriber = synchronousMessageSubscriber.get();
+        SynchronousMessageSubscriber messageSubscriber = synchronousMessageSubscriber.get();
 
         if (messageSubscriber != null) {
             messageSubscriber.queueWork(work);
             return;
         }
 
-        final SynchronousMessageSubscriber newSubscriber = new SynchronousMessageSubscriber(asyncClient,
-            work,
-            isPrefetchDisabled,
-            operationTimeout);
+        messageSubscriber = synchronousMessageSubscriber.updateAndGet(subscriber -> {
+            // Ensuring we create SynchronousMessageSubscriber only once.
+            if (subscriber == null) {
+                return new SynchronousMessageSubscriber(asyncClient,
+                    work,
+                    isPrefetchDisabled,
+                    operationTimeout);
+            } else {
+                return subscriber;
+            }
+        });
 
         // NOTE: We asynchronously send the credit to the service as soon as receiveMessage() API is called (for first
         // time).
         // This means that there may be messages internally buffered before users start iterating the IterableStream.
         // If users do not iterate through the stream and their lock duration expires, it is possible that the
         // Service Bus message's delivery count will be incremented.
-        if (synchronousMessageSubscriber.compareAndSet(null, newSubscriber)) {
-            asyncClient.receiveMessagesNoBackPressure().subscribeWith(newSubscriber);
+        if (!syncSubscribed.getAndSet(true)) {
+            // The 'subscribeWith' has side effects hence must not be called from
+            // the above updateFunction of AtomicReference::updateAndGet.
+            asyncClient.receiveMessagesNoBackPressure().subscribeWith(messageSubscriber);
         } else {
-            newSubscriber.dispose();
-            synchronousMessageSubscriber.get().queueWork(work);
+            messageSubscriber.queueWork(work);
         }
 
         logger.atVerbose()

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

@@ -136,9 +136,9 @@ void queueWork(SynchronousReceiveWork work) {
             .addKeyValue("numberOfEvents", work.getNumberOfEvents())
             .addKeyValue("timeout", work.getTimeout());
 
-        // If previous work items were completed, the message queue is empty and currentWork == null. Update the
-        // current work and request items upstream if we need to.
-        if (workQueue.peek() == work) {
+        // If previous work items were completed, the message queue is empty, and currentWork is null or terminal,
+        // Update the current work and request items upstream if we need to.
+        if (workQueue.peek() == work && (currentWork == null || currentWork.isTerminal())) {
             logBuilder.log("First work in queue. Requesting upstream if needed.");
             getOrUpdateCurrentWork();
         } else {
@@ -236,6 +236,11 @@ private void drainQueue() {
                 numberRequested = REQUESTED.addAndGet(this, -numberConsumed);
             }
         }
+        if (numberRequested == 0L) {
+            logger.atVerbose()
+                .log("Current work is completed. Schedule next work.");
+            getOrUpdateCurrentWork();
+        }
     }
 
     /**
@@ -272,35 +277,17 @@ private SynchronousReceiveWork getOrUpdateCurrentWork() {
             }
 
             currentWork = workQueue.poll();
-            while (currentWork != null) {
-                // For the terminal work, subtract the remaining number of messages from our current request
-                // count. This is so we don't keep adding credits for work that was expired, but we never
-                // received messages for.
-                if (currentWork.isTerminal()) {
-                    REQUESTED.updateAndGet(this, currentRequest -> {
-                        final int remainingEvents = currentWork.getRemainingEvents();
-
-                        // The work had probably emitted all its messages and then terminated.
-                        // The currentRequest is fine.
-                        if (remainingEvents < 1) {
-                            return currentRequest;
-                        }
-
-                        final long difference = currentRequest - remainingEvents;
-
-                        logger.atVerbose()
-                            .addKeyValue(NUMBER_OF_REQUESTED_MESSAGES_KEY, currentRequest)
-                            .addKeyValue("remainingEvents", remainingEvents)
-                            .addKeyValue("difference", difference)
-                            .log("Updating REQUESTED because current work item is terminal.");
-
-                        return difference < 0 ? 0 : difference;
-                    });
-
-                    currentWork = workQueue.poll();
-                    continue;
-                }
+            //The work in queue will not be terminal, here is double check
+            while (currentWork != null && currentWork.isTerminal()) {
+                logger.atVerbose()
+                    .addKeyValue(WORK_ID_KEY, currentWork.getId())
+                    .addKeyValue("numberOfEvents", currentWork.getNumberOfEvents())
+                    .log("This work from queue is terminal. Skip it.");
 
+                currentWork = workQueue.poll();
+            }
+
+            if (currentWork != null) {
                 final SynchronousReceiveWork work = currentWork;
                 logger.atVerbose()
                     .addKeyValue(WORK_ID_KEY, work.getId())
@@ -312,8 +299,6 @@ private SynchronousReceiveWork getOrUpdateCurrentWork() {
                 // Now that we updated REQUESTED to account for credits already on the line, we're good to
                 // place any credits for this new work item.
                 requestUpstream(work.getNumberOfEvents());
-
-                return work;
             }
 
             return currentWork;

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

@@ -17,15 +17,12 @@
 import reactor.core.publisher.Mono;
 
 import java.time.Duration;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
@@ -124,15 +121,11 @@ public void queueWorkTest() {
     }
 
     /**
-     * Verifies that this processes multiple work items.
+     * Verifies that this processes multiple work items and current work encounter timeout
      */
     @Test
-    public void processesMultipleWorkItems() {
+    public void processesMultipleWorkItemsAndCurrentWorkTimeout() {
         // Arrange
-        final SynchronousReceiveWork work3 = mock(SynchronousReceiveWork.class);
-        when(work3.getId()).thenReturn(3L);
-        when(work3.getNumberOfEvents()).thenReturn(1);
-
         final ServiceBusReceivedMessage message1 = mock(ServiceBusReceivedMessage.class);
         final ServiceBusReceivedMessage message2 = mock(ServiceBusReceivedMessage.class);
         final ServiceBusReceivedMessage message3 = mock(ServiceBusReceivedMessage.class);
@@ -154,9 +147,15 @@ public void processesMultipleWorkItems() {
         doAnswer(invocation -> isTerminal.get()).when(work1).isTerminal();
         doAnswer(invocation -> remaining.get()).when(work1).getRemainingEvents();
 
+        // WORK 2 is update to current work after the work1 is terminal and successfully emits message3
         when(work2.emitNext(message3)).thenReturn(true);
         when(work2.isTerminal()).thenReturn(false);
 
+        // WORK 3 is placed in queue
+        final SynchronousReceiveWork work3 = mock(SynchronousReceiveWork.class);
+        when(work3.getId()).thenReturn(3L);
+        when(work3.getNumberOfEvents()).thenReturn(1);
+
         syncSubscriber = new SynchronousMessageSubscriber(asyncClient, work1, false, operationTimeout);
         syncSubscriber.queueWork(work2);
         syncSubscriber.queueWork(work3);
@@ -188,15 +187,79 @@ public void processesMultipleWorkItems() {
         //     - work2.getNumberOfEvents() - (REQUESTED - work1.getRemainingItems());
         verify(subscription, times(2)).request(subscriptionArgumentCaptor.capture());
         final List<Long> allRequests = subscriptionArgumentCaptor.getAllValues();
-        final Set<Long> expected = new HashSet<>();
-        expected.add((long) work1.getNumberOfEvents());
+        assertEquals(NUMBER_OF_WORK_ITEMS, allRequests.get(0));
 
         final long requestedAfterWork1 = NUMBER_OF_WORK_ITEMS - remaining.get();
         final long expectedDifference = work2.getNumberOfEvents() - requestedAfterWork1;
-        expected.add(expectedDifference);
+        assertEquals(expectedDifference, allRequests.get(1));
+    }
+
+    /**
+     * Verifies that this processes multiple work items and current work can emit all messages successfully
+     */
+    @Test
+    public void processesMultipleWorkItemsAndCurrentWorkEmitAllMessages() {
+        // Arrange
+        final ServiceBusReceivedMessage message1 = mock(ServiceBusReceivedMessage.class);
+        final ServiceBusReceivedMessage message2 = mock(ServiceBusReceivedMessage.class);
+        final ServiceBusReceivedMessage message3 = mock(ServiceBusReceivedMessage.class);
+        final ServiceBusReceivedMessage message4 = mock(ServiceBusReceivedMessage.class);
+        final ServiceBusReceivedMessage message5 = mock(ServiceBusReceivedMessage.class);
+
+        // WORK 1 successfully emits all messages and is terminal after emit message4.
+        final AtomicBoolean isTerminal = new AtomicBoolean(false);
+        final AtomicInteger remaining = new AtomicInteger(NUMBER_OF_WORK_ITEMS);
+        doAnswer(invocation -> {
+            ServiceBusReceivedMessage arg = invocation.getArgument(0);
+            remaining.decrementAndGet();
+            if (arg == message4) {
+                isTerminal.set(true);
+            }
+            return true;
+        }).when(work1).emitNext(any(ServiceBusReceivedMessage.class));
+        doAnswer(invocation -> isTerminal.get()).when(work1).isTerminal();
+        doAnswer(invocation -> remaining.get()).when(work1).getRemainingEvents();
+
+        // WORK 2 is updated to current work after work1 completed and successfully emit message5
+        when(work2.isTerminal()).thenReturn(false);
+        when(work2.emitNext(message5)).thenReturn(true);
 
-        assertEquals(expected.size(), allRequests.size());
-        allRequests.forEach(r -> assertTrue(expected.contains(r)));
+        // WORK 3 is placed in queue
+        final SynchronousReceiveWork work3 = mock(SynchronousReceiveWork.class);
+        when(work3.getId()).thenReturn(3L);
+        when(work3.getNumberOfEvents()).thenReturn(1);
+
+        syncSubscriber = new SynchronousMessageSubscriber(asyncClient, work1, false, operationTimeout);
+        syncSubscriber.queueWork(work2);
+        syncSubscriber.queueWork(work3);
+
+        syncSubscriber.hookOnSubscribe(subscription);
+
+        assertEquals(2, syncSubscriber.getWorkQueueSize());
+
+        // Act
+        syncSubscriber.hookOnNext(message1);
+        syncSubscriber.hookOnNext(message2);
+        syncSubscriber.hookOnNext(message3);
+        syncSubscriber.hookOnNext(message4);
+
+        // Assert
+        verify(work2).start();
+
+        // work2 emits message5
+        syncSubscriber.hookOnNext(message5);
+
+        verify(work2).emitNext(message5);
+
+        assertEquals(1, syncSubscriber.getWorkQueueSize());
+
+        // Verify that we requested:
+        // 1st time: hookOnSubscribe(work1.getNumberOfEvents())
+        // 2nd time: requestUpstream(work2.getNumberOfEvents()) and REQUESTED = 0
+        verify(subscription, times(2)).request(subscriptionArgumentCaptor.capture());
+        final List<Long> allRequests = subscriptionArgumentCaptor.getAllValues();
+        assertEquals(NUMBER_OF_WORK_ITEMS, allRequests.get(0));
+        assertEquals(NUMBER_OF_WORK_ITEMS_2, allRequests.get(1));
     }
 
     /**