ServiceBus: small optimizations and logging improvements (Azure#36536)

Liudmila Molkova · web-flow · commit 0f1712f4ebaa · 2023-08-25T23:40:30.000Z
* optimizations and logging improvements
diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/exception/AmqpResponseCode.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/exception/AmqpResponseCode.java
@@ -3,9 +3,6 @@
 
 package com.azure.core.amqp.exception;
 
-import java.util.HashMap;
-import java.util.Map;
-
 /**
  * Error response codes returned from AMQP.
  */
@@ -200,11 +197,11 @@ public enum AmqpResponseCode {
      */
     HTTP_VERSION_NOT_SUPPORTED(505);
 
-    private static final Map<Integer, AmqpResponseCode> VALUE_MAP = new HashMap<>();
+    private static final AmqpResponseCode[] VALUE_MAP = new AmqpResponseCode[1024];
 
     static {
         for (AmqpResponseCode code : AmqpResponseCode.values()) {
-            VALUE_MAP.put(code.value, code);
+            VALUE_MAP[code.value] = code;
         }
     }
 
@@ -222,7 +219,11 @@ public enum AmqpResponseCode {
      * is found.
      */
     public static AmqpResponseCode fromValue(final int value) {
-        return VALUE_MAP.get(value);
+        if (value >= 0 && value < VALUE_MAP.length) {
+            return VALUE_MAP[value];
+        }
+
+        return null;
     }
 
     /**
diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ManagementChannel.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ManagementChannel.java
@@ -23,6 +23,7 @@
 import java.util.Objects;
 
 import static com.azure.core.amqp.implementation.AmqpLoggingUtils.addKeyValueIfNotNull;
+import static com.azure.core.amqp.implementation.ClientConstants.DELIVERY_STATE_KEY;
 import static com.azure.core.amqp.implementation.ClientConstants.ENTITY_PATH_KEY;
 import static com.azure.core.amqp.implementation.ClientConstants.ERROR_CONDITION_KEY;
 import static com.azure.core.amqp.implementation.ClientConstants.ERROR_DESCRIPTION_KEY;
@@ -59,9 +60,7 @@ public Mono<AmqpAnnotatedMessage> send(AmqpAnnotatedMessage message) {
             return channel.sendWithAck(protonJMessage)
                 .handle((Message responseMessage, SynchronousSink<AmqpAnnotatedMessage> sink) ->
                     handleResponse(responseMessage, sink, channel.getErrorContext()))
-                .switchIfEmpty(Mono.error(() -> new AmqpException(true, String.format(
-                    "entityPath[%s] No response received from management channel.", entityPath),
-                    channel.getErrorContext())));
+                .switchIfEmpty(errorIfEmpty(channel, null));
         }));
     }
 
@@ -74,9 +73,7 @@ public Mono<AmqpAnnotatedMessage> send(AmqpAnnotatedMessage message, DeliveryOut
             return channel.sendWithAck(protonJMessage, protonJDeliveryState)
                 .handle((Message responseMessage, SynchronousSink<AmqpAnnotatedMessage> sink) ->
                     handleResponse(responseMessage, sink, channel.getErrorContext()))
-                .switchIfEmpty(Mono.error(() -> new AmqpException(true, String.format(
-                    "entityPath[%s] outcome[%s] No response received from management channel.", entityPath,
-                    deliveryOutcome.getDeliveryState()), channel.getErrorContext())));
+                .switchIfEmpty(errorIfEmpty(channel, deliveryOutcome.getDeliveryState()));
         }));
     }
 
@@ -126,6 +123,17 @@ private void handleResponse(Message response, SynchronousSink<AmqpAnnotatedMessa
         sink.error(throwable);
     }
 
+    private <T> Mono<T> errorIfEmpty(RequestResponseChannel channel, com.azure.core.amqp.models.DeliveryState deliveryState) {
+        return Mono.error(() -> {
+            String error = String.format(
+                "entityPath[%s] deliveryState[%s] No response received from management channel.", entityPath, deliveryState);
+            AmqpException exception = new AmqpException(true, error, channel.getErrorContext());
+            return logger.atError()
+                .addKeyValue(DELIVERY_STATE_KEY, deliveryState)
+                .log(exception);
+        });
+    }
+
     private Mono<Void> isAuthorized() {
         return tokenManager.getAuthorizationResults()
             .next()
diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/RequestResponseChannel.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/RequestResponseChannel.java
@@ -396,13 +396,18 @@ protected Message decodeDelivery(Delivery delivery) {
     }
 
     private void settleMessage(Message message) {
-        final String id = String.valueOf(message.getCorrelationId());
-        final UnsignedLong correlationId = UnsignedLong.valueOf(id);
+        UnsignedLong correlationId;
+        if (message.getCorrelationId() instanceof UnsignedLong) {
+            correlationId = (UnsignedLong) message.getCorrelationId();
+        } else {
+            String id = String.valueOf(message.getCorrelationId());
+            correlationId = UnsignedLong.valueOf(id);
+        }
         final MonoSink<Message> sink = unconfirmedSends.remove(correlationId);
 
         if (sink == null) {
             logger.atWarning()
-                .addKeyValue("messageId", id)
+                .addKeyValue("messageId", message.getCorrelationId())
                 .log("Received delivery without pending message.");
             return;
         }
diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/handler/ReactorHandler.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/handler/ReactorHandler.java
@@ -39,6 +39,8 @@ public void onReactorInit(Event e) {
 
     @Override
     public void onReactorFinal(Event e) {
-        logger.info("reactor.onReactorFinal. event: {}", e);
+        logger.atInfo()
+            .addKeyValue("event", e)
+            .log("reactor.onReactorFinal.");
     }
 }
diff --git a/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/exception/AmqpResponseCodeTest.java b/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/exception/AmqpResponseCodeTest.java
@@ -5,6 +5,10 @@
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import static org.junit.jupiter.api.Assertions.assertNull;
 
 public class AmqpResponseCodeTest {
     /**
@@ -20,6 +24,12 @@ public void createFromInteger() {
         Assertions.assertEquals(expected, actual);
     }
 
+    @ParameterizedTest
+    @ValueSource(ints = {-1, 42, 1024})
+    public void createFromInvalidInteger(int notDefined) {
+        assertNull(AmqpResponseCode.fromValue(notDefined));
+    }
+
     /**
      * Verifies that we can parse the AmqpResponseCode an integer.
      */
diff --git a/sdk/core/azure-core-metrics-opentelemetry/src/main/java/com/azure/core/metrics/opentelemetry/OpenTelemetryUtils.java b/sdk/core/azure-core-metrics-opentelemetry/src/main/java/com/azure/core/metrics/opentelemetry/OpenTelemetryUtils.java
@@ -9,8 +9,6 @@
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 
-import java.util.Optional;
-
 import static com.azure.core.util.tracing.Tracer.PARENT_TRACE_CONTEXT_KEY;
 
 class OpenTelemetryUtils {
@@ -51,17 +49,20 @@ static void addAttribute(io.opentelemetry.api.common.AttributesBuilder attribute
      * If not context is found, returns {@link io.opentelemetry.context.Context#current()}.
      */
     static io.opentelemetry.context.Context getTraceContextOrCurrent(Context azContext) {
-        Optional<Object> traceContextOpt = azContext.getData(PARENT_TRACE_CONTEXT_KEY);
-        if (traceContextOpt.isPresent()) {
-            Object traceContextObj = traceContextOpt.get();
-            if (traceContextObj instanceof io.opentelemetry.context.Context) {
+        Object traceContextObj = azContext.getData(PARENT_TRACE_CONTEXT_KEY).orElse(null);
+        if (traceContextObj != null) {
+            if (io.opentelemetry.context.Context.class.isAssignableFrom(traceContextObj.getClass())) {
                 return (io.opentelemetry.context.Context) traceContextObj;
-            } else if (traceContextObj != null) {
-                // TODO (limolkova) somehow we can get shaded otel agent context here
-                if (!warnedOnContextType) {
-                    LOGGER.warning("Expected instance of `io.opentelemetry.context.Context` under `PARENT_TRACE_CONTEXT_KEY`, but got {}, ignoring it.", traceContextObj.getClass().getName());
-                    warnedOnContextType = true;
-                }
+            } else if (!warnedOnContextType) {
+                // TODO (limolkova) https://github.com/Azure/azure-sdk-for-java/issues/36537
+                // The context we have here is created by azure-core-tracing-opentelemetry.
+                // but if otel or applicationInsights agents are used, the azure-core-tracing-opentelemetry is shaded and records shaded context.
+                // if azure-core-metrics-opentelemetry is NOT shaded, we won't be able to reuse this context since it's not compatible.
+                // I.e. it works fine if both are shaded or not shaded and warns (once) if just one of them is shaded.
+                // We should fix this by adding azure-core-metrics-opentelemetry to otel agent
+                // In the meantime, it's ok - we return current context and exemplars should work reasonably well.
+                LOGGER.warning("Expected instance of `io.opentelemetry.context.Context` under `PARENT_TRACE_CONTEXT_KEY`, but got {}, ignoring it.", traceContextObj.getClass().getName());
+                warnedOnContextType = true;
             }
         }
 
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/FluxAutoComplete.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/FluxAutoComplete.java
@@ -72,7 +72,9 @@ static final class AutoCompleteSubscriber extends BaseSubscriber<ServiceBusMessa
 
         @Override
         protected void hookOnSubscribe(Subscription subscription) {
-            logger.info("Subscription received. Subscribing downstream. {}", subscription);
+            logger.atVerbose()
+                .addKeyValue("subscription", subscription)
+                .log("Subscription received. Subscribing downstream.");
             downstream.onSubscribe(this);
         }
 
@@ -149,7 +151,9 @@ private void applyWithCatch(Function<ServiceBusMessageContext, Mono<Void>> funct
             try {
                 function.apply(context).block();
             } catch (Exception e) {
-                logger.warning("Unable to '{}' message.", operation, e);
+                logger.atWarning()
+                    .addKeyValue("operation", operation)
+                    .log("Operation on message failed.", e);
 
                 // On an error, we'll stop requesting from upstream and pass the error downstream.
                 upstream().cancel();
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusMessageSerializer.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusMessageSerializer.java
@@ -341,7 +341,9 @@ private List<ServiceBusReceivedMessage> deserializeListOfMessages(Message amqpMe
         final AmqpResponseCode statusCode = RequestResponseUtils.getStatusCode(amqpMessage);
 
         if (statusCode != AmqpResponseCode.OK) {
-            LOGGER.warning("AMQP response did not contain OK status code. Actual: {}", statusCode);
+            LOGGER.atWarning()
+                .addKeyValue("statusCode", statusCode)
+                .log("AMQP response did not contain OK status code.");
             return Collections.emptyList();
         }
 
@@ -351,25 +353,32 @@ private List<ServiceBusReceivedMessage> deserializeListOfMessages(Message amqpMe
             LOGGER.warning("AMQP response did not contain a body.");
             return Collections.emptyList();
         } else if (!(responseBodyMap instanceof Map)) {
-            LOGGER.warning("AMQP response body is not correct instance. Expected: {}. Actual: {}",
-                Map.class, responseBodyMap.getClass());
+            LOGGER.atWarning()
+                .addKeyValue("expectedType", Map.class)
+                .addKeyValue("actualType", responseBodyMap.getClass())
+                .log("AMQP response body is not correct instance.");
             return Collections.emptyList();
         }
 
         final Object messages = ((Map) responseBodyMap).get(ManagementConstants.MESSAGES);
         if (messages == null) {
-            LOGGER.warning("Response body did not contain key: {}", ManagementConstants.MESSAGES);
+            LOGGER.atWarning().addKeyValue("expectedKey", ManagementConstants.MESSAGES)
+                .log("AMQP response body did not contain key.");
             return Collections.emptyList();
         } else if (!(messages instanceof Iterable)) {
-            LOGGER.warning("Response body contents is not the correct type. Expected: {}. Actual: {}",
-                Iterable.class, messages.getClass());
+            LOGGER.atWarning()
+                .addKeyValue("expectedType", Iterable.class)
+                .addKeyValue("actualType", messages.getClass())
+                .log("Response body contents is not the correct type.");
             return Collections.emptyList();
         }
 
         for (Object message : (Iterable) messages) {
             if (!(message instanceof Map)) {
-                LOGGER.warning("Message inside iterable of message is not correct type. Expected: {}. Actual: {}",
-                    Map.class, message.getClass());
+                LOGGER.atWarning()
+                    .addKeyValue("expectedType", Map.class)
+                    .addKeyValue("actualType", message.getClass())
+                    .log("Message inside iterable of message is not correct type.");
                 continue;
             }
 
@@ -407,11 +416,17 @@ private ServiceBusReceivedMessage deserializeMessage(Message amqpMessage) {
                 amqpMessageBody = AmqpMessageBody.fromSequence(messageData);
 
             } else {
-                LOGGER.warning(String.format(Messages.MESSAGE_NOT_OF_TYPE, body.getType()));
+                LOGGER.atWarning()
+                    .addKeyValue("actualType", body.getType())
+                    .log("Message body is not correct. Not setting body contents.");
+
                 amqpMessageBody = AmqpMessageBody.fromData(EMPTY_BYTE_ARRAY);
             }
         } else {
-            LOGGER.warning(String.format(Messages.MESSAGE_NOT_OF_TYPE, "null"));
+            LOGGER.atWarning()
+                .addKeyValue("actualType", "null")
+                .log("Message body is not correct. Not setting body contents.");
+
             amqpMessageBody = AmqpMessageBody.fromData(EMPTY_BYTE_ARRAY);
         }
 
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SynchronousReceiveWork.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SynchronousReceiveWork.java
@@ -152,13 +152,17 @@ synchronized boolean emitNext(ServiceBusReceivedMessage message) {
         final int numberLeft = remaining.decrementAndGet();
 
         if (numberLeft < 0) {
-            logger.info("Number left {} < 0. Not emitting downstream.", numberLeft);
+            logger.atInfo()
+                .addKeyValue("numberLeft", numberLeft)
+                .log("Not emitting downstream.");
             return false;
         }
 
         final Sinks.EmitResult result = downstreamEmitter.tryEmitNext(message);
         if (result != Sinks.EmitResult.OK) {
-            logger.info("Could not emit downstream. EmitResult: {}", result);
+            logger.atInfo()
+                .addKeyValue("emitResult", result)
+                .log("Could not emit downstream.");
             return false;
         }
 
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/administration/implementation/EntityHelper.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/administration/implementation/EntityHelper.java
@@ -495,7 +495,9 @@ private static List<AuthorizationRuleImpl> toImplementation(List<AuthorizationRu
                 implementation.setType("SharedAccessAuthorizationRule");
             } else {
                 final String className = rule.getClass().getName();
-                LOGGER.warning("AuthorizationRule type '{}' is unknown.", className);
+                LOGGER.atWarning()
+                    .addKeyValue("type", className)
+                    .log("AuthorizationRule type is unknown.");
                 implementation.setType(className);
             }
 
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ManagementChannel.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ManagementChannel.java
@@ -592,11 +592,17 @@ private Mono<Message> sendWithVerify(RequestResponseChannel channel, Message mes
 
                 sink.error(throwable);
             })
-            .switchIfEmpty(Mono.error(new AmqpException(true, "No response received from management channel.",
-                channel.getErrorContext())))
+            .switchIfEmpty(errorIfEmpty(channel))
             .onErrorMap(this::mapError);
     }
 
+    private <T> Mono<T> errorIfEmpty(RequestResponseChannel channel) {
+        return Mono.error(() -> {
+            String error = String.format("entityPath[%s] No response received from management channel.", entityPath);
+            return logger.logExceptionAsError(new AmqpException(true, error, channel.getErrorContext()));
+        });
+    }
+
     private Mono<Void> isAuthorized(String operation) {
         return tokenManager.getAuthorizationResults()
             .onErrorMap(this::mapError)
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/Messages.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/Messages.java
@@ -25,7 +25,6 @@ public class Messages {
     public static final String INVALID_OPERATION_DISPOSED_RECEIVER = getMessage("INVALID_OPERATION_DISPOSED_RECEIVER");
     public static final String INVALID_OPERATION_DISPOSED_SENDER = getMessage("INVALID_OPERATION_DISPOSED_SENDER");
     public static final String INVALID_OPERATION_DISPOSED_RULE_MANAGER = getMessage("INVALID_OPERATION_DISPOSED_RULE_MANAGER");
-    public static final String MESSAGE_NOT_OF_TYPE = getMessage("MESSAGE_NOT_OF_TYPE");
 
     private static synchronized Properties getProperties() {
         if (properties != null) {
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusSharedKeyCredential.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusSharedKeyCredential.java
@@ -195,8 +195,9 @@ private OffsetDateTime getExpirationTime(String sharedAccessSignature) {
                     long epochSeconds = Long.parseLong(expirationTimeStr);
                     return Instant.ofEpochSecond(epochSeconds).atOffset(ZoneOffset.UTC);
                 } catch (NumberFormatException exception) {
-                    LOGGER.verbose("Invalid expiration time format in the SAS token: {}. Falling back to max "
-                        + "expiration time.", expirationTimeStr);
+                    LOGGER.atVerbose()
+                        .addKeyValue("expirationTime", expirationTimeStr)
+                        .log("Invalid expiration time format in the SAS token. Falling back to max expiration time.");
                     return OffsetDateTime.MAX;
                 }
             })
diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/resources/azure-messaging-servicebus.properties b/sdk/servicebus/azure-messaging-servicebus/src/main/resources/azure-messaging-servicebus.properties
@@ -1,6 +1,5 @@
 name=${project.artifactId}
 version=${project.version}
 CLASS_NOT_A_SUPPORTED_TYPE=Class '%s' is not a supported deserializable type.
-MESSAGE_NOT_OF_TYPE=Message body type is not of type Data, but type: %s. Not setting body contents.
 INVALID_OPERATION_DISPOSED_RECEIVER=Cannot perform operation '%s' on a disposed receiver.
 INVALID_OPERATION_DISPOSED_SENDER=Cannot perform operation '%s' on a disposed sender.

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,8 @@ public void onReactorInit(Event e) {`
`39`	`39`
`40`	`40`	`@Override`
`41`	`41`	`public void onReactorFinal(Event e) {`
`42`		`- logger.info("reactor.onReactorFinal. event: {}", e);`
	`42`	`+ logger.atInfo()`
	`43`	`+ .addKeyValue("event", e)`
	`44`	`+ .log("reactor.onReactorFinal.");`
`43`	`45`	`}`
`44`	`46`	`}`