API ServiceBusErrorSource to represent source of error

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

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.messaging.servicebus;
+
+/**
+ * Represent the operation user was performing when the error happened.
+ */
+public enum ServiceBusErrorSource {
+    /** Error while sending the message(s).*/
+    SEND,
+
+    /** Error while receiving the message(s).*/
+    RECEIVE,
+
+    /** Error while abandoning the message.*/
+    ABANDONED,
+
+    /** Error while completing the message.*/
+    COMPLETE,
+
+    /** Error while deferring the message.*/
+    DEFER,
+
+    /** Error while dead-lettering the message.*/
+    DEAD_LETTER;
+}

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

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.messaging.servicebus;
+
+import com.azure.core.amqp.exception.AmqpException;
+
+/**
+ * Defines {@link ServiceBusException} which has addition properties. You can {@link ServiceBusErrorSource} to
+ * determine source of error.
+ *
+ * @see ServiceBusErrorSource
+ */
+public class ServiceBusException extends AmqpException {
+    private final ServiceBusErrorSource errorSource;
+
+    /**
+     * @param amqpException for the error hapened.
+     * @param errorSource indicating which api caused the error.
+     */
+    ServiceBusException(AmqpException amqpException, ServiceBusErrorSource errorSource) {
+        super(amqpException.isTransient(), amqpException.getErrorCondition(), amqpException.getMessage(),
+            amqpException.getCause(), amqpException.getContext());
+        this.errorSource = errorSource;
+    }
+
+    /**
+     *  Gets the {@link ServiceBusErrorSource} in case of any errors.
+     *
+     * @return the {@link ServiceBusErrorSource}
+     */
+    public ServiceBusErrorSource getErrorSource() {
+        return errorSource;
+    }
+}

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

@@ -611,7 +611,14 @@ public Flux<ServiceBusReceivedMessageContext> receiveMessages() {
             withAutoComplete = withAutoLockRenewal;
         }
 
-        return withAutoComplete;
+        return withAutoComplete
+            .onErrorMap(throwable -> {
+                if (throwable instanceof AmqpException) {
+                    return new ServiceBusException((AmqpException) throwable, ServiceBusErrorSource.RECEIVE);
+                } else {
+                    return throwable;
+                }
+            });
     }
 
     /**
@@ -1048,9 +1055,10 @@ deadLetterErrorDescription, propertiesToModify, sessionId, getLinkName(sessionId
                 renewalContainer.remove(lockToken);
             }));
 
+        Mono<Void> updateDispositionOperation;
         if (sessionManager != null) {
-            return sessionManager.updateDisposition(lockToken, sessionId, dispositionStatus, propertiesToModify,
-                deadLetterReason, deadLetterErrorDescription, transactionContext)
+            updateDispositionOperation =  sessionManager.updateDisposition(lockToken, sessionId, dispositionStatus,
+                propertiesToModify, deadLetterReason, deadLetterErrorDescription, transactionContext)
                 .flatMap(isSuccess -> {
                     if (isSuccess) {
                         renewalContainer.remove(lockToken);
@@ -1060,20 +1068,49 @@ deadLetterErrorDescription, propertiesToModify, sessionId, getLinkName(sessionId
                     logger.info("Could not perform on session manger. Performing on management node.");
                     return performOnManagement;
                 });
-        }
-
-        final ServiceBusAsyncConsumer existingConsumer = consumer.get();
-        if (isManagementToken(lockToken) || existingConsumer == null) {
-            return performOnManagement;
         } else {
-            return existingConsumer.updateDisposition(lockToken, dispositionStatus, deadLetterReason,
-                deadLetterErrorDescription, propertiesToModify, transactionContext)
-                .then(Mono.fromRunnable(() -> {
-                    logger.info("{}: Update completed. Disposition: {}. Lock: {}.",
-                        entityPath, dispositionStatus, lockToken);
-                    renewalContainer.remove(lockToken);
-                }));
+            final ServiceBusAsyncConsumer existingConsumer = consumer.get();
+            if (isManagementToken(lockToken) || existingConsumer == null) {
+                updateDispositionOperation = performOnManagement;
+            } else {
+                updateDispositionOperation = existingConsumer.updateDisposition(lockToken, dispositionStatus,
+                    deadLetterReason, deadLetterErrorDescription, propertiesToModify, transactionContext)
+                    .then(Mono.fromRunnable(() -> {
+                        logger.info("{}: Update completed. Disposition: {}. Lock: {}.",
+                            entityPath, dispositionStatus, lockToken);
+                        renewalContainer.remove(lockToken);
+                    }));
+            }
         }
+        return updateDispositionOperation
+            .onErrorMap(throwable -> {
+                ServiceBusErrorSource errorSource;
+                if (throwable instanceof AmqpException) {
+                    switch (dispositionStatus) {
+                        case COMPLETED:
+                            errorSource = ServiceBusErrorSource.COMPLETE;
+                            break;
+                        case DEFERRED:
+                            errorSource = ServiceBusErrorSource.DEFER;
+                            break;
+                        case SUSPENDED:
+                            errorSource = ServiceBusErrorSource.DEAD_LETTER;
+                            break;
+                        case ABANDONED:
+                            errorSource = ServiceBusErrorSource.ABANDONED;
+                            break;
+                        default:
+                            throw logger.logExceptionAsError(new UnsupportedOperationException(String.format(
+                                "'%s' is not supported.", dispositionStatus)));
+                    }
+
+                    return new ServiceBusException((AmqpException) throwable, errorSource);
+
+                } else {
+                    return throwable;
+                }
+
+            });
     }
 
     private ServiceBusAsyncConsumer getOrCreateConsumer() {

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

@@ -503,7 +503,8 @@ private Mono<Void> sendIterable(Iterable<ServiceBusMessage> messages, ServiceBus
         return createMessageBatch().flatMap(messageBatch -> {
             messages.forEach(message -> messageBatch.tryAddMessage(message));
             return sendInternal(messageBatch, transaction);
-        });
+        })
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.SEND));
     }
 
 
@@ -607,7 +608,8 @@ private Mono<Void> sendInternal(ServiceBusMessageBatch batch, ServiceBusTransact
                 if (isTracingEnabled) {
                     tracerProvider.endSpan(parentContext.get(), signal);
                 }
-            });
+            })
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.SEND));
 
     }
 
@@ -621,8 +623,8 @@ private Mono<Void> sendInternal(Flux<ServiceBusMessage> messages, ServiceBusTran
                     return messages.collect(new AmqpMessageCollector(batchOptions, 1,
                         link::getErrorContext, tracerProvider, messageSerializer, entityName,
                         link.getHostname()));
-                })
-                .flatMap(list -> sendInternalBatch(Flux.fromIterable(list), transactionContext)));
+                }).flatMap(list -> sendInternalBatch(Flux.fromIterable(list), transactionContext)))
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.SEND));
     }
 
     private Mono<Void> sendInternalBatch(Flux<ServiceBusMessageBatch> eventBatches,
@@ -729,4 +731,16 @@ public Set<Characteristics> characteristics() {
             return Collections.emptySet();
         }
     }
+
+    /**
+     * Maps the exception into ServiceBusException
+     */
+    private Throwable mapError(Throwable throwable, ServiceBusErrorSource errorSource) {
+        if (throwable instanceof AmqpException) {
+            return new ServiceBusException((AmqpException) throwable, errorSource);
+        } else {
+            return throwable;
+        }
+    }
+
 }

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

@@ -7,6 +7,7 @@
 import com.azure.core.amqp.AmqpRetryOptions;
 import com.azure.core.amqp.AmqpTransportType;
 import com.azure.core.amqp.ProxyOptions;
+import com.azure.core.amqp.exception.AmqpException;
 import com.azure.core.amqp.implementation.CbsAuthorizationType;
 import com.azure.core.amqp.implementation.ConnectionOptions;
 import com.azure.core.amqp.implementation.MessageSerializer;
@@ -29,6 +30,7 @@
 import com.azure.messaging.servicebus.models.ReceiveMode;
 import org.apache.qpid.proton.amqp.messaging.Accepted;
 import org.apache.qpid.proton.amqp.messaging.Rejected;
+import org.apache.qpid.proton.amqp.transport.DeliveryState;
 import org.apache.qpid.proton.amqp.transport.DeliveryState.DeliveryStateType;
 import org.apache.qpid.proton.engine.SslDomain;
 import org.apache.qpid.proton.message.Message;
@@ -40,7 +42,9 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInfo;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
@@ -65,6 +69,7 @@
 import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import static com.azure.messaging.servicebus.TestUtils.getMessage;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -443,6 +448,92 @@ void deadLetterWithDescription() {
         verify(amqpReceiveLink).updateDisposition(eq(lockToken1), isA(Rejected.class));
     }
 
+    /**
+     * Verifies that error source is populated when there is any error during message settlement.
+     */
+    @ParameterizedTest
+    @MethodSource
+    void errorSourceOnSettlement(DispositionStatus dispositionStatus, ServiceBusErrorSource expectedErrorSource,
+        DeliveryStateType expectedDeliveryState) {
+        final String lockToken1 = UUID.randomUUID().toString();
+
+        final OffsetDateTime expiration = OffsetDateTime.now().plus(Duration.ofMinutes(5));
+
+        final MessageWithLockToken message = mock(MessageWithLockToken.class);
+
+        when(messageSerializer.deserialize(message, ServiceBusReceivedMessage.class)).thenReturn(receivedMessage);
+
+        when(receivedMessage.getLockToken()).thenReturn(lockToken1);
+        when(receivedMessage.getLockedUntil()).thenReturn(expiration);
+
+        when(amqpReceiveLink.updateDisposition(eq(lockToken1), argThat(e -> e.getType() == expectedDeliveryState)))
+            .thenReturn(Mono.error(new AmqpException(false, "some error occured.", null)));
+
+        // Act & Assert
+        StepVerifier.create(receiver.receiveMessages().take(1)
+            .flatMap(context -> {
+                final Mono<Void> operation;
+                switch (dispositionStatus) {
+                    case DEFERRED:
+                        operation = receiver.defer(receivedMessage);
+                        break;
+                    case ABANDONED:
+                        operation = receiver.abandon(receivedMessage);
+                        break;
+                    case COMPLETED:
+                        operation = receiver.complete(receivedMessage);
+                        break;
+                    case SUSPENDED:
+                        operation = receiver.deadLetter(receivedMessage);
+                        break;
+                    default:
+                        throw new IllegalArgumentException("Unrecognized operation: " + dispositionStatus);
+                }
+                return operation;
+            }))
+            .then(() -> messageSink.next(message))
+            .expectNext()
+            .verifyErrorMatches(throwable -> {
+                Assertions.assertTrue(throwable instanceof  ServiceBusException);
+                final ServiceBusErrorSource actual = ((ServiceBusException) throwable).getErrorSource();
+                Assertions.assertEquals(expectedErrorSource, actual);
+                return true;
+            });
+
+        verify(amqpReceiveLink).updateDisposition(eq(lockToken1), any(DeliveryState.class));
+    }
+
+    /**
+     * Verifies that error source is populated when there is any error during receiving of message.
+     */
+    @Test
+    void errorSourceOnReceiveMessage() {
+        final String lockToken1 = UUID.randomUUID().toString();
+
+        final OffsetDateTime expiration = OffsetDateTime.now().plus(Duration.ofMinutes(5));
+
+        final MessageWithLockToken message = mock(MessageWithLockToken.class);
+
+        when(messageSerializer.deserialize(message, ServiceBusReceivedMessage.class)).thenReturn(receivedMessage);
+
+        when(receivedMessage.getLockToken()).thenReturn(lockToken1);
+        when(receivedMessage.getLockedUntil()).thenReturn(expiration);
+
+        when(connection.createReceiveLink(anyString(), anyString(), any(ReceiveMode.class), any(),
+            any(MessagingEntityType.class))).thenReturn(Mono.error(new AmqpException(false, "some receive link Error.", null)));
+
+        // Act & Assert
+        StepVerifier.create(receiver.receiveMessages().take(1))
+            .verifyErrorMatches(throwable -> {
+                Assertions.assertTrue(throwable instanceof  ServiceBusException);
+                final ServiceBusErrorSource actual = ((ServiceBusException) throwable).getErrorSource();
+                Assertions.assertEquals(ServiceBusErrorSource.RECEIVE, actual);
+                return true;
+            });
+
+        verify(amqpReceiveLink, never()).updateDisposition(eq(lockToken1), any(DeliveryState.class));
+    }
+
     /**
      * Verifies that the user can complete settlement methods on received message.
      */
@@ -970,4 +1061,12 @@ private List<Message> getMessages() {
             .mapToObj(index -> getMessage(PAYLOAD_BYTES, messageTrackingUUID, map))
             .collect(Collectors.toList());
     }
+
+    private static Stream<Arguments> errorSourceOnSettlement() {
+        return Stream.of(
+            Arguments.of(DispositionStatus.DEFERRED, ServiceBusErrorSource.DEFER, DeliveryStateType.Modified),
+            Arguments.of(DispositionStatus.COMPLETED, ServiceBusErrorSource.COMPLETE, DeliveryStateType.Accepted),
+            Arguments.of(DispositionStatus.SUSPENDED, ServiceBusErrorSource.DEAD_LETTER, DeliveryStateType.Rejected),
+            Arguments.of(DispositionStatus.ABANDONED, ServiceBusErrorSource.ABANDONED, DeliveryStateType.Modified));
+    }
 }