Fixed the issue of ServiceBusReceiverClient recreating the link after…

… closing the client (#34665) * Fixed the issue of ServiceBusReceiverClient recreating the link after closing the client * Changelog update * unit tests to assert expected error from Iterable iteration.
Azure · Apr 28, 2023 · 6db1030 · 6db1030
1 parent 4dc93b8
commit 6db1030
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 2 deletions.
diff --git a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md
@@ -10,6 +10,7 @@
 
 - Fixed issue where receiving messages from `ServiceBusSessionReceiverAsyncClient` would never complete. ([#34597](https://github.com/Azure/azure-sdk-for-java/issues/34597))
 - Fixed issue causing some messages to not be returned when calling peek on receiver client. 
+- Fixed the issue of `ServiceBusReceiverClient` recreating the link after closing the client. ([#34664](https://github.com/Azure/azure-sdk-for-java/issues/34664))
 
 ### Other Changes
 

diff --git a/...ing-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java b/...ing-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java
@@ -796,7 +796,7 @@ public void rollbackTransaction(ServiceBusTransactionContext transactionContext)
      */
     @Override
     public void close() {
-        SynchronousMessageSubscriber messageSubscriber = synchronousMessageSubscriber.getAndSet(null);
+        SynchronousMessageSubscriber messageSubscriber = synchronousMessageSubscriber.get();
         if (messageSubscriber != null && !messageSubscriber.isDisposed()) {
             messageSubscriber.dispose();
         }

diff --git a/...servicebus/src/main/java/com/azure/messaging/servicebus/SynchronousMessageSubscriber.java b/...servicebus/src/main/java/com/azure/messaging/servicebus/SynchronousMessageSubscriber.java
@@ -269,6 +269,11 @@ protected void hookOnError(Throwable throwable) {
         dispose("Errors occurred upstream", throwable);
     }
 
+    @Override
+    protected void hookOnComplete() {
+        dispose("Upstream signaled completion", null);
+    }
+
     @Override
     protected void hookOnCancel() {
         this.dispose();
@@ -397,4 +402,3 @@ int getWorkQueueSize() {
         return this.workQueue.size();
     }
 }
-
diff --git a/...servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverClientTest.java b/...servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverClientTest.java
@@ -3,6 +3,8 @@
 
 package com.azure.messaging.servicebus;
 
+import com.azure.core.amqp.exception.AmqpErrorContext;
+import com.azure.core.amqp.exception.AmqpException;
 import com.azure.core.util.IterableStream;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.messaging.servicebus.implementation.instrumentation.ServiceBusReceiverInstrumentation;
@@ -17,6 +19,7 @@
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
+import org.opentest4j.AssertionFailedError;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.publisher.TestPublisher;
@@ -43,6 +46,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
@@ -859,4 +863,98 @@ void setSessionState() {
         // Assert
         verify(asyncClient).setSessionState(SESSION_ID, contents);
     }
+
+    @Test
+    void iterationReplaysUpstreamTerminalError() {
+        // Arrange
+        final int messageCount = 10;
+        final AmqpException terminalError = new AmqpException(false, "non-retriable terminal error.", new AmqpErrorContext("contoso.com"));
+        final AtomicInteger state = new AtomicInteger(0);
+
+        Flux<ServiceBusReceivedMessage> messageSink = Flux.create(sink -> {
+            sink.onRequest(r -> {
+                final int s = state.getAndIncrement();
+                if (s == 0) {
+                    for (int m = 0; m < r; m++) {
+                        sink.next(mock(ServiceBusReceivedMessage.class));
+                        if (m >= messageCount) {
+                            sink.error(new AssertionFailedError(String.format("Received request for %d when expected to emit %d messages.", r, messageCount)));
+                        }
+                    }
+                } else if (s == 1) {
+                    sink.error(terminalError);
+                } else {
+                    sink.error((new AssertionFailedError("Unexpected request after termination.")));
+                }
+            });
+
+            sink.onCancel(() -> {
+                LOGGER.info("Cancelled. Completing sink.");
+                sink.complete();
+            });
+        });
+        when(asyncClient.receiveMessagesNoBackPressure()).thenReturn(messageSink);
+
+        // Assert the first receive get messages.
+        final IterableStream<ServiceBusReceivedMessage> messages0 = client.receiveMessages(messageCount);
+        assertNotNull(messages0);
+        final long collected = messages0.stream().count();
+        assertEquals(messageCount, collected);
+
+        // Assert the second receive iteration get terminal error.
+        final IterableStream<ServiceBusReceivedMessage> messages1 = client.receiveMessages(messageCount);
+        assertNotNull(messages1);
+        final AmqpException e1 = assertThrows(AmqpException.class,  () -> messages1.stream().count());
+        assertEquals(terminalError, e1);
+
+        // Assert the same terminal error 'replayed' to further receive iterations.
+        final IterableStream<ServiceBusReceivedMessage> messages2 = client.receiveMessages(messageCount);
+        assertNotNull(messages2);
+        final AmqpException e2 = assertThrows(AmqpException.class,  () -> messages2.stream().count());
+        assertEquals(terminalError, e2);
+    }
+
+    @Test
+    void iterationAfterCloseEmitsError() {
+        // Arrange
+        final int messageCount = 10;
+        final AtomicInteger state = new AtomicInteger(0);
+
+        Flux<ServiceBusReceivedMessage> messageSink = Flux.create(sink -> {
+            sink.onRequest(r -> {
+                final int s = state.getAndIncrement();
+                if (s == 0) {
+                    for (int m = 0; m < r; m++) {
+                        sink.next(mock(ServiceBusReceivedMessage.class));
+                        if (m >= messageCount) {
+                            sink.error(new AssertionFailedError(String.format("Received request for %d when expected to emit %d messages.", r, messageCount)));
+                        }
+                    }
+                } else {
+                    sink.error((new AssertionFailedError("Unexpected request after termination.")));
+                }
+            });
+
+            sink.onCancel(() -> {
+                LOGGER.info("Cancelled. Completing sink.");
+                sink.complete();
+            });
+        });
+        when(asyncClient.receiveMessagesNoBackPressure()).thenReturn(messageSink);
+        doNothing().when(asyncClient).close();
+
+        // Assert the first receive get messages.
+        final IterableStream<ServiceBusReceivedMessage> messages0 = client.receiveMessages(messageCount);
+        assertNotNull(messages0);
+        final long collected = messages0.stream().count();
+        assertEquals(messageCount, collected);
+
+        client.close();
+
+        // Assert any iteration on iterable obtained after client close get error.
+        final IterableStream<ServiceBusReceivedMessage> messages1 = client.receiveMessages(messageCount);
+        assertNotNull(messages1);
+        final RuntimeException e = assertThrows(RuntimeException.class,  () -> messages1.stream().count());
+        assertEquals("The receiver client is terminated. Re-create the client to continue receive attempt.", e.getMessage());
+    }
 }