API ServiceBusErrorSource to represent source of error

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

@@ -0,0 +1,34 @@
+// 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 ServiceBusAmqpException} which has additional information about the operation that caused the error.
+ *
+ * @see ServiceBusErrorSource
+ */
+public final class ServiceBusAmqpException extends AmqpException {
+    private final ServiceBusErrorSource errorSource;
+
+    /**
+     * @param amqpException for the error hapened.
+     * @param errorSource indicating which api caused the error.
+     */
+    ServiceBusAmqpException(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/ServiceBusErrorSource.java

@@ -0,0 +1,52 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.messaging.servicebus;
+
+import com.azure.core.util.ExpandableStringEnum;
+
+import java.io.Serializable;
+
+/**
+ * Represent the operation this sdk was performing when the error happened.
+ */
+public final class ServiceBusErrorSource extends ExpandableStringEnum<ServiceBusErrorSource> implements Serializable {
+
+    private static final long serialVersionUID = -2819764417333954922L;
+
+    /** Error while abandoning the message.*/
+    public static final ServiceBusErrorSource ABANDONED = fromString("ABANDONED", ServiceBusErrorSource.class);
+
+    /** Error while completing the message.*/
+    public static final ServiceBusErrorSource COMPLETE = fromString("COMPLETE", ServiceBusErrorSource.class);
+
+    /** Error while deferring the message.*/
+    public static final ServiceBusErrorSource DEFER = fromString("DEFER", ServiceBusErrorSource.class);
+
+    /** Error while dead-lettering the message.*/
+    public static final ServiceBusErrorSource DEAD_LETTER = fromString("DEAD_LETTER",
+        ServiceBusErrorSource.class);
+
+    /** Error while receiving the message(s).*/
+    public static final ServiceBusErrorSource RECEIVE = fromString("RECEIVE", ServiceBusErrorSource.class);
+
+    /** Error while renewing lock.*/
+    public static final ServiceBusErrorSource RENEW_LOCK = fromString("RENEW_LOCK", ServiceBusErrorSource.class);
+
+    /** Error when we could not determine the source.*/
+    public static final ServiceBusErrorSource UNKNOWN = fromString("UNKNOWN", ServiceBusErrorSource.class);
+
+    /** Error while user's code is running for a message.*/
+    public static final ServiceBusErrorSource USER_CALLBACK = fromString("USER_CALLBACK",
+        ServiceBusErrorSource.class);
+
+    /** Error while session is accepted.*/
+    public static final ServiceBusErrorSource ACCEPT_MESSAGE_SESSION = fromString("ACCEPT_MESSAGE_SESSION",
+        ServiceBusErrorSource.class);
+
+    /** Error while session is closed.*/
+    public static final ServiceBusErrorSource CLOSE_MESSAGE_SESSION = fromString("CLOSE_MESSAGE_SESSION",
+        ServiceBusErrorSource.class);
+
+
+}

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

@@ -611,7 +611,8 @@ public Flux<ServiceBusReceivedMessageContext> receiveMessages() {
             withAutoComplete = withAutoLockRenewal;
         }
 
-        return withAutoComplete;
+        return withAutoComplete
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE));
     }
 
     /**
@@ -733,7 +734,8 @@ public Mono<OffsetDateTime> renewMessageLock(ServiceBusReceivedMessage message)
                 String.format("Cannot renew message lock [%s] for a session receiver.", message.getLockToken())));
         }
 
-        return renewMessageLock(message.getLockToken());
+        return renewMessageLock(message.getLockToken())
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK));
     }
 
     /**
@@ -751,7 +753,8 @@ Mono<OffsetDateTime> renewMessageLock(String lockToken) {
             .flatMap(serviceBusManagementNode ->
                 serviceBusManagementNode.renewMessageLock(lockToken, getLinkName(null)))
             .map(offsetDateTime -> managementNodeLocks.addOrUpdate(lockToken, offsetDateTime,
-                offsetDateTime));
+                offsetDateTime))
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK));
     }
 
     /**
@@ -790,7 +793,8 @@ public Mono<Void> renewMessageLock(ServiceBusReceivedMessage message, Duration m
         renewalContainer.addOrUpdate(message.getLockToken(), OffsetDateTime.now().plus(maxLockRenewalDuration),
             operation);
 
-        return operation.getCompletionOperation();
+        return operation.getCompletionOperation()
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK));
     }
 
     /**
@@ -815,7 +819,8 @@ public Mono<OffsetDateTime> renewSessionLock(String sessionId) {
 
         return connectionProcessor
             .flatMap(connection -> connection.getManagementNode(entityPath, entityType))
-            .flatMap(channel -> channel.renewSessionLock(sessionId, linkName));
+            .flatMap(channel -> channel.renewSessionLock(sessionId, linkName))
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK));
     }
 
     /**
@@ -851,7 +856,9 @@ public Mono<Void> renewSessionLock(String sessionId, Duration maxLockRenewalDura
             this::renewSessionLock);
 
         renewalContainer.addOrUpdate(sessionId, OffsetDateTime.now().plus(maxLockRenewalDuration), operation);
-        return operation.getCompletionOperation();
+        return operation
+            .getCompletionOperation()
+            .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK));
     }
 
     /**
@@ -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,48 @@ 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:
+                            errorSource = ServiceBusErrorSource.UNKNOWN;
+                    }
+
+                    return new ServiceBusAmqpException((AmqpException) throwable, errorSource);
+
+                } else {
+                    return throwable;
+                }
+
+            });
     }
 
     private ServiceBusAsyncConsumer getOrCreateConsumer() {
@@ -1136,4 +1172,15 @@ private String getLinkName(String sessionId) {
             return existing != null ? existing.getLinkName() : null;
         }
     }
+
+    /**
+     * Map the error to {@link ServiceBusAmqpException}
+     */
+    private Throwable mapError(Throwable throwable, ServiceBusErrorSource errorSource) {
+        if (throwable instanceof AmqpException) {
+            return new ServiceBusAmqpException((AmqpException) throwable, errorSource);
+        } else {
+            return throwable;
+        }
+    }
 }