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

@@ -201,7 +201,7 @@ public String getEntityPath() {
      */
     public Mono<Void> abandon(ServiceBusReceivedMessage message) {
         return updateDisposition(message, DispositionStatus.ABANDONED, null, null,
-            null, null);
+            null, null, ServiceBusErrorSource.ABANDONED);
     }
 
     /**
@@ -233,7 +233,7 @@ public Mono<Void> abandon(ServiceBusReceivedMessage message, AbandonOptions opti
         }
 
         return updateDisposition(message, DispositionStatus.ABANDONED, null, null,
-            options.getPropertiesToModify(), options.getTransactionContext());
+            options.getPropertiesToModify(), options.getTransactionContext(), ServiceBusErrorSource.ABANDONED);
     }
 
     /**
@@ -248,7 +248,7 @@ public Mono<Void> abandon(ServiceBusReceivedMessage message, AbandonOptions opti
      */
     public Mono<Void> complete(ServiceBusReceivedMessage message) {
         return updateDisposition(message, DispositionStatus.COMPLETED, null, null,
-            null, null);
+            null, null, ServiceBusErrorSource.COMPLETE);
     }
 
     /**
@@ -277,7 +277,7 @@ public Mono<Void> complete(ServiceBusReceivedMessage message, CompleteOptions op
         }
 
         return updateDisposition(message, DispositionStatus.COMPLETED, null, null,
-            null, options.getTransactionContext());
+            null, options.getTransactionContext(), ServiceBusErrorSource.COMPLETE);
     }
 
     /**
@@ -292,7 +292,7 @@ public Mono<Void> complete(ServiceBusReceivedMessage message, CompleteOptions op
      */
     public Mono<Void> defer(ServiceBusReceivedMessage message) {
         return updateDisposition(message, DispositionStatus.DEFERRED, null, null,
-            null, null);
+            null, null, ServiceBusErrorSource.DEFER);
     }
 
     /**
@@ -323,7 +323,7 @@ public Mono<Void> defer(ServiceBusReceivedMessage message, DeferOptions options)
         }
 
         return updateDisposition(message, DispositionStatus.DEFERRED, null, null,
-            options.getPropertiesToModify(), options.getTransactionContext());
+            options.getPropertiesToModify(), options.getTransactionContext(), ServiceBusErrorSource.DEFER);
     }
 
     /**
@@ -370,7 +370,7 @@ public Mono<Void> deadLetter(ServiceBusReceivedMessage message, DeadLetterOption
         }
         return  updateDisposition(message, DispositionStatus.SUSPENDED, options.getDeadLetterReason(),
             options.getDeadLetterErrorDescription(), options.getPropertiesToModify(),
-            options.getTransactionContext());
+            options.getTransactionContext(), ServiceBusErrorSource.DEAD_LETTER);
     }
 
     /**
@@ -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;
+                }
+            });
     }
 
     /**
@@ -1008,8 +1015,21 @@ private boolean isManagementToken(String lockToken) {
 
     private Mono<Void> updateDisposition(ServiceBusReceivedMessage message, DispositionStatus dispositionStatus,
         String deadLetterReason, String deadLetterErrorDescription, Map<String, Object> propertiesToModify,
-        ServiceBusTransactionContext transactionContext) {
+        ServiceBusTransactionContext transactionContext, ServiceBusErrorSource errorSource) {
+        return updateDispositionInternal(message, dispositionStatus, deadLetterReason, deadLetterErrorDescription,
+            propertiesToModify, transactionContext)
+            .onErrorMap(throwable -> {
+                if (throwable instanceof AmqpException) {
+                    return new ServiceBusException((AmqpException) throwable, errorSource);
+                }
+                return throwable;
 
+            });
+    }
+
+    private Mono<Void> updateDispositionInternal(ServiceBusReceivedMessage message, DispositionStatus dispositionStatus,
+        String deadLetterReason, String deadLetterErrorDescription, Map<String, Object> propertiesToModify,
+        ServiceBusTransactionContext transactionContext) {
         if (isDisposed.get()) {
             return monoError(logger, new IllegalStateException(
                 String.format(INVALID_OPERATION_DISPOSED_RECEIVER, dispositionStatus.getValue())));

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

@@ -503,6 +503,12 @@ private Mono<Void> sendIterable(Iterable<ServiceBusMessage> messages, ServiceBus
         return createMessageBatch().flatMap(messageBatch -> {
             messages.forEach(message -> messageBatch.tryAddMessage(message));
             return sendInternal(messageBatch, transaction);
+        }).onErrorMap(throwable -> {
+            if (throwable instanceof AmqpException) {
+                return new ServiceBusException((AmqpException) throwable, ServiceBusErrorSource.SEND);
+            } else {
+                return throwable;
+            }
         });
     }
 
@@ -607,6 +613,12 @@ private Mono<Void> sendInternal(ServiceBusMessageBatch batch, ServiceBusTransact
                 if (isTracingEnabled) {
                     tracerProvider.endSpan(parentContext.get(), signal);
                 }
+            }).onErrorMap(throwable -> {
+                if (throwable instanceof AmqpException) {
+                    return new ServiceBusException((AmqpException) throwable, ServiceBusErrorSource.SEND);
+                } else {
+                    return throwable;
+                }
             });
 
     }
@@ -621,8 +633,14 @@ 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 -> {
+                if (throwable instanceof AmqpException) {
+                    return new ServiceBusException((AmqpException) throwable, ServiceBusErrorSource.SEND);
+                } else {
+                    return throwable;
+                }
+            });
     }
 
     private Mono<Void> sendInternalBatch(Flux<ServiceBusMessageBatch> eventBatches,

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));
+    }
 }