diff --git a/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml b/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml index 725adce69d21..f57dcb91ab03 100755 --- a/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml +++ b/eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml @@ -98,6 +98,7 @@ + diff --git a/sdk/communication/azure-communication-callautomation/src/main/java/com/azure/communication/callautomation/models/CallConnectionProperties.java b/sdk/communication/azure-communication-callautomation/src/main/java/com/azure/communication/callautomation/models/CallConnectionProperties.java index 618a42e18b68..79179db87b6b 100644 --- a/sdk/communication/azure-communication-callautomation/src/main/java/com/azure/communication/callautomation/models/CallConnectionProperties.java +++ b/sdk/communication/azure-communication-callautomation/src/main/java/com/azure/communication/callautomation/models/CallConnectionProperties.java @@ -139,8 +139,7 @@ public String getCallConnectionId() { * * @return the mediaSubscriptionId value. */ - public String getMediaSubscriptionId() { + public String getMediaSubscriptionId() { return mediaSubscriptionId; } - } diff --git a/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java b/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java index 4cf115d3b117..0587e5e32008 100644 --- a/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java +++ b/sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java @@ -362,7 +362,7 @@ private SpanBuilder createSpanBuilder(String spanName, SpanBuilder spanBuilder = tracer.spanBuilder(spanNameKey) .setSpanKind(spanKind); - io.opentelemetry.context.Context parentContext = getTraceContextOrDefault(context, io.opentelemetry.context.Context.current()); + io.opentelemetry.context.Context parentContext = getTraceContextOrDefault(context, io.opentelemetry.context.Context.current()); // if remote parent is provided, it has higher priority if (remoteParentContext != null) { spanBuilder.setParent(parentContext.with(Span.wrap(remoteParentContext))); diff --git a/sdk/servicebus/azure-messaging-servicebus/pom.xml b/sdk/servicebus/azure-messaging-servicebus/pom.xml index 8e8af51e1933..6ae86d91f3e1 100644 --- a/sdk/servicebus/azure-messaging-servicebus/pom.xml +++ b/sdk/servicebus/azure-messaging-servicebus/pom.xml @@ -118,6 +118,27 @@ 4.5.1 test + + + com.azure + azure-core-tracing-opentelemetry + 1.0.0-beta.28 + test + + + + io.opentelemetry + opentelemetry-api + 1.14.0 + test + + + + io.opentelemetry + opentelemetry-sdk + 1.14.0 + test + diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/FluxTrace.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/FluxTrace.java new file mode 100644 index 000000000000..0f7a49f0ca52 --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/FluxTrace.java @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.Context; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; +import org.reactivestreams.Subscription; +import reactor.core.CoreSubscriber; +import reactor.core.publisher.BaseSubscriber; +import reactor.core.publisher.Flux; +import reactor.core.publisher.FluxOperator; + +import java.util.Objects; + +/** + * Flux operator that traces receive and process calls + */ +final class FluxTrace extends FluxOperator { + static final String PROCESS_ERROR_KEY = "process-error"; + private final ServiceBusReceiverTracer tracer; + + FluxTrace(Flux upstream, ServiceBusReceiverTracer tracer) { + super(upstream); + this.tracer = tracer; + } + + @Override + public void subscribe(CoreSubscriber coreSubscriber) { + Objects.requireNonNull(coreSubscriber, "'coreSubscriber' cannot be null."); + + source.subscribe(new TracingSubscriber(coreSubscriber, tracer)); + } + + private static class TracingSubscriber extends BaseSubscriber { + + private final CoreSubscriber downstream; + private final ServiceBusReceiverTracer tracer; + TracingSubscriber(CoreSubscriber downstream, ServiceBusReceiverTracer tracer) { + this.downstream = downstream; + this.tracer = tracer; + } + + @Override + public reactor.util.context.Context currentContext() { + return downstream.currentContext(); + } + + @Override + protected void hookOnSubscribe(Subscription subscription) { + downstream.onSubscribe(this); + } + + @Override + protected void hookOnNext(ServiceBusMessageContext message) { + if (tracer == null || tracer.isSync()) { + downstream.onNext(message); + return; + } + + Throwable exception = null; + Context span = tracer.startProcessSpan("ServiceBus.process", message.getMessage(), Context.NONE); + AutoCloseable scope = tracer.makeSpanCurrent(span); + + try { + downstream.onNext(message); + } catch (Throwable t) { + exception = t; + } finally { + Context context = message.getMessage().getContext(); + if (context != null) { + Object processorException = context.getData(PROCESS_ERROR_KEY).orElse(null); + if (processorException instanceof Throwable) { + exception = (Throwable) processorException; + } + } + tracer.endSpan(exception, span, scope); + } + } + + @Override + protected void hookOnError(Throwable throwable) { + downstream.onError(throwable); + } + + @Override + protected void hookOnComplete() { + downstream.onComplete(); + } + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java index a8449e77cd00..d705c78be934 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java @@ -17,7 +17,6 @@ import com.azure.core.amqp.implementation.ReactorProvider; import com.azure.core.amqp.implementation.StringUtil; import com.azure.core.amqp.implementation.TokenManagerProvider; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.amqp.models.CbsAuthorizationType; import com.azure.core.annotation.ServiceClientBuilder; import com.azure.core.annotation.ServiceClientProtocol; @@ -34,14 +33,15 @@ import com.azure.core.util.Configuration; import com.azure.core.util.CoreUtils; import com.azure.core.util.logging.ClientLogger; -import com.azure.core.util.tracing.Tracer; import com.azure.messaging.servicebus.implementation.MessageUtils; import com.azure.messaging.servicebus.implementation.MessagingEntityType; import com.azure.messaging.servicebus.implementation.ServiceBusAmqpConnection; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; import com.azure.messaging.servicebus.implementation.ServiceBusConstants; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.implementation.ServiceBusReactorAmqpConnection; import com.azure.messaging.servicebus.implementation.ServiceBusSharedKeyCredential; +import com.azure.messaging.servicebus.implementation.ServiceBusTracer; import com.azure.messaging.servicebus.implementation.models.ServiceBusProcessorClientOptions; import com.azure.messaging.servicebus.models.ServiceBusReceiveMode; import com.azure.messaging.servicebus.models.SubQueue; @@ -59,7 +59,6 @@ import java.util.Locale; import java.util.Map; import java.util.Objects; -import java.util.ServiceLoader; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; @@ -212,7 +211,6 @@ public final class ServiceBusClientBuilder implements private final Object connectionLock = new Object(); private final MessageSerializer messageSerializer = new ServiceBusMessageSerializer(); - private final TracerProvider tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class)); private ClientOptions clientOptions; private Configuration configuration; @@ -963,8 +961,9 @@ public ServiceBusSenderAsyncClient buildAsyncClient() { clientIdentifier = UUID.randomUUID().toString(); } + final ServiceBusSenderTracer tracer = new ServiceBusSenderTracer(ServiceBusTracer.getDefaultTracer(), connectionProcessor.getFullyQualifiedNamespace(), entityName); return new ServiceBusSenderAsyncClient(entityName, entityType, connectionProcessor, retryOptions, - tracerProvider, messageSerializer, ServiceBusClientBuilder.this::onClientClose, null, clientIdentifier); + tracer, messageSerializer, ServiceBusClientBuilder.this::onClientClose, null, clientIdentifier); } /** @@ -1047,8 +1046,7 @@ public final class ServiceBusSessionProcessorClientBuilder { private ServiceBusSessionProcessorClientBuilder() { sessionReceiverClientBuilder = new ServiceBusSessionReceiverClientBuilder(); processorClientOptions = new ServiceBusProcessorClientOptions() - .setMaxConcurrentCalls(1) - .setTracerProvider(tracerProvider); + .setMaxConcurrentCalls(1); sessionReceiverClientBuilder.maxConcurrentSessions(1); } @@ -1444,11 +1442,11 @@ ServiceBusReceiverAsyncClient buildAsyncClientForProcessor() { } final ServiceBusSessionManager sessionManager = new ServiceBusSessionManager(entityPath, entityType, - connectionProcessor, tracerProvider, messageSerializer, receiverOptions, clientIdentifier); - + connectionProcessor, messageSerializer, receiverOptions, clientIdentifier); + final ServiceBusReceiverTracer tracer = new ServiceBusReceiverTracer(ServiceBusTracer.getDefaultTracer(), connectionProcessor.getFullyQualifiedNamespace(), entityPath, false); return new ServiceBusReceiverAsyncClient(connectionProcessor.getFullyQualifiedNamespace(), entityPath, entityType, receiverOptions, connectionProcessor, ServiceBusConstants.OPERATION_TIMEOUT, - tracerProvider, messageSerializer, ServiceBusClientBuilder.this::onClientClose, sessionManager); + tracer, messageSerializer, ServiceBusClientBuilder.this::onClientClose, sessionManager); } /** @@ -1466,7 +1464,7 @@ ServiceBusReceiverAsyncClient buildAsyncClientForProcessor() { * queueName()} or {@link #topicName(String) topicName()}, respectively. */ public ServiceBusSessionReceiverAsyncClient buildAsyncClient() { - return buildAsyncClient(true); + return buildAsyncClient(true, false); } /** @@ -1484,12 +1482,12 @@ public ServiceBusSessionReceiverAsyncClient buildAsyncClient() { */ public ServiceBusSessionReceiverClient buildClient() { final boolean isPrefetchDisabled = prefetchCount == 0; - return new ServiceBusSessionReceiverClient(buildAsyncClient(false), + return new ServiceBusSessionReceiverClient(buildAsyncClient(false, true), isPrefetchDisabled, MessageUtils.getTotalTimeout(retryOptions)); } - private ServiceBusSessionReceiverAsyncClient buildAsyncClient(boolean isAutoCompleteAllowed) { + private ServiceBusSessionReceiverAsyncClient buildAsyncClient(boolean isAutoCompleteAllowed, boolean syncConsumer) { final MessagingEntityType entityType = validateEntityPaths(connectionStringEntityName, topicName, queueName); final String entityPath = getEntityPath(entityType, queueName, topicName, subscriptionName, @@ -1520,8 +1518,10 @@ private ServiceBusSessionReceiverAsyncClient buildAsyncClient(boolean isAutoComp clientIdentifier = UUID.randomUUID().toString(); } + final ServiceBusReceiverTracer tracer = new ServiceBusReceiverTracer(ServiceBusTracer.getDefaultTracer(), + connectionProcessor.getFullyQualifiedNamespace(), entityPath, syncConsumer); return new ServiceBusSessionReceiverAsyncClient(connectionProcessor.getFullyQualifiedNamespace(), - entityPath, entityType, receiverOptions, connectionProcessor, tracerProvider, messageSerializer, + entityPath, entityType, receiverOptions, connectionProcessor, tracer, messageSerializer, ServiceBusClientBuilder.this::onClientClose, clientIdentifier); } } @@ -1582,8 +1582,7 @@ public final class ServiceBusProcessorClientBuilder { private ServiceBusProcessorClientBuilder() { serviceBusReceiverClientBuilder = new ServiceBusReceiverClientBuilder(); processorClientOptions = new ServiceBusProcessorClientOptions() - .setMaxConcurrentCalls(1) - .setTracerProvider(tracerProvider); + .setMaxConcurrentCalls(1); } /** @@ -1908,7 +1907,7 @@ public ServiceBusReceiverClientBuilder topicName(String topicName) { * queueName()} or {@link #topicName(String) topicName()}, respectively. */ public ServiceBusReceiverAsyncClient buildAsyncClient() { - return buildAsyncClient(true); + return buildAsyncClient(true, false); } /** @@ -1926,12 +1925,12 @@ public ServiceBusReceiverAsyncClient buildAsyncClient() { */ public ServiceBusReceiverClient buildClient() { final boolean isPrefetchDisabled = prefetchCount == 0; - return new ServiceBusReceiverClient(buildAsyncClient(false), + return new ServiceBusReceiverClient(buildAsyncClient(false, true), isPrefetchDisabled, MessageUtils.getTotalTimeout(retryOptions)); } - ServiceBusReceiverAsyncClient buildAsyncClient(boolean isAutoCompleteAllowed) { + ServiceBusReceiverAsyncClient buildAsyncClient(boolean isAutoCompleteAllowed, boolean syncConsumer) { final MessagingEntityType entityType = validateEntityPaths(connectionStringEntityName, topicName, queueName); final String entityPath = getEntityPath(entityType, queueName, topicName, subscriptionName, @@ -1962,9 +1961,11 @@ ServiceBusReceiverAsyncClient buildAsyncClient(boolean isAutoCompleteAllowed) { clientIdentifier = UUID.randomUUID().toString(); } + final ServiceBusReceiverTracer tracer = new ServiceBusReceiverTracer(ServiceBusTracer.getDefaultTracer(), + connectionProcessor.getFullyQualifiedNamespace(), entityPath, syncConsumer); return new ServiceBusReceiverAsyncClient(connectionProcessor.getFullyQualifiedNamespace(), entityPath, entityType, receiverOptions, connectionProcessor, ServiceBusConstants.OPERATION_TIMEOUT, - tracerProvider, messageSerializer, ServiceBusClientBuilder.this::onClientClose, clientIdentifier); + tracer, messageSerializer, ServiceBusClientBuilder.this::onClientClose, clientIdentifier); } } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusMessageBatch.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusMessageBatch.java index 3308d1aac920..6a8ea03376fe 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusMessageBatch.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusMessageBatch.java @@ -7,7 +7,6 @@ import com.azure.core.amqp.exception.AmqpException; import com.azure.core.amqp.implementation.ErrorContextProvider; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.util.logging.ClientLogger; import org.apache.qpid.proton.message.Message; @@ -17,7 +16,6 @@ import java.util.Locale; import java.util.Objects; -import static com.azure.messaging.servicebus.implementation.MessageUtils.traceMessageSpan; /** * A class for aggregating {@link ServiceBusMessage messages} into a single, size-limited, batch. It is treated as a @@ -31,11 +29,11 @@ public final class ServiceBusMessageBatch { private final List serviceBusMessageList; private final byte[] eventBytes; private int sizeInBytes; - private final TracerProvider tracerProvider; + private final ServiceBusSenderTracer tracer; private final String entityPath; private final String hostname; - ServiceBusMessageBatch(int maxMessageSize, ErrorContextProvider contextProvider, TracerProvider tracerProvider, + ServiceBusMessageBatch(int maxMessageSize, ErrorContextProvider contextProvider, ServiceBusSenderTracer tracer, MessageSerializer serializer, String entityPath, String hostname) { this.maxMessageSize = maxMessageSize; this.contextProvider = contextProvider; @@ -43,7 +41,7 @@ public final class ServiceBusMessageBatch { this.serviceBusMessageList = new ArrayList<>(); this.sizeInBytes = (maxMessageSize / 65536) * 1024; // reserve 1KB for every 64KB this.eventBytes = new byte[maxMessageSize]; - this.tracerProvider = tracerProvider; + this.tracer = tracer; this.entityPath = entityPath; this.hostname = hostname; } @@ -94,15 +92,11 @@ public boolean tryAddMessage(final ServiceBusMessage serviceBusMessage) { if (serviceBusMessage == null) { throw LOGGER.logExceptionAsWarning(new NullPointerException("'serviceBusMessage' cannot be null")); } - ServiceBusMessage serviceBusMessageUpdated = - tracerProvider.isEnabled() - ? traceMessageSpan(serviceBusMessage, serviceBusMessage.getContext(), hostname, entityPath, - tracerProvider) - : serviceBusMessage; + tracer.createMessageSpan(serviceBusMessage); final int size; try { - size = getSize(serviceBusMessageUpdated, serviceBusMessageList.isEmpty()); + size = getSize(serviceBusMessage, serviceBusMessageList.isEmpty()); } catch (BufferOverflowException exception) { final RuntimeException ex = new ServiceBusException( new AmqpException(false, AmqpErrorCondition.LINK_PAYLOAD_SIZE_EXCEEDED, @@ -117,7 +111,7 @@ public boolean tryAddMessage(final ServiceBusMessage serviceBusMessage) { } this.sizeInBytes += size; - this.serviceBusMessageList.add(serviceBusMessageUpdated); + this.serviceBusMessageList.add(serviceBusMessage); return true; } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorClient.java index 55d2d83e7f55..99212d649516 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorClient.java @@ -3,39 +3,22 @@ package com.azure.messaging.servicebus; -import com.azure.core.amqp.implementation.TracerProvider; -import com.azure.core.util.Context; import com.azure.core.util.logging.ClientLogger; -import com.azure.core.util.tracing.ProcessKind; import com.azure.messaging.servicebus.ServiceBusClientBuilder.ServiceBusProcessorClientBuilder; import com.azure.messaging.servicebus.ServiceBusClientBuilder.ServiceBusSessionProcessorClientBuilder; import com.azure.messaging.servicebus.implementation.models.ServiceBusProcessorClientOptions; import org.reactivestreams.Subscription; import reactor.core.CoreSubscriber; import reactor.core.Disposable; -import reactor.core.publisher.Signal; import reactor.core.scheduler.Schedulers; - -import java.util.Locale; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; -import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; -import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; -import static com.azure.core.util.tracing.Tracer.ENTITY_PATH_KEY; -import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; -import static com.azure.core.util.tracing.Tracer.MESSAGE_ENQUEUED_TIME; -import static com.azure.core.util.tracing.Tracer.SCOPE_KEY; -import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; -import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_NAMESPACE_VALUE; -import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_SERVICE_NAME; - /** * The processor client for processing Service Bus messages. {@link ServiceBusProcessorClient} provides a push-based * mechanism that invokes the message processing callback when a message is received or the error handler when an error @@ -133,7 +116,6 @@ public final class ServiceBusProcessorClient implements AutoCloseable { private final Map receiverSubscriptions = new ConcurrentHashMap<>(); private final AtomicReference asyncClient = new AtomicReference<>(); private final AtomicBoolean isRunning = new AtomicBoolean(); - private final TracerProvider tracerProvider; private final String queueName; private final String topicName; private final String subscriptionName; @@ -160,9 +142,9 @@ public final class ServiceBusProcessorClient implements AutoCloseable { this.processMessage = Objects.requireNonNull(processMessage, "'processMessage' cannot be null"); this.processError = Objects.requireNonNull(processError, "'processError' cannot be null"); this.processorOptions = Objects.requireNonNull(processorOptions, "'processorOptions' cannot be null"); + this.asyncClient.set(sessionReceiverBuilder.buildAsyncClientForProcessor()); this.receiverBuilder = null; - this.tracerProvider = processorOptions.getTracerProvider(); this.queueName = queueName; this.topicName = topicName; this.subscriptionName = subscriptionName; @@ -189,7 +171,6 @@ public final class ServiceBusProcessorClient implements AutoCloseable { this.processorOptions = Objects.requireNonNull(processorOptions, "'processorOptions' cannot be null"); this.asyncClient.set(receiverBuilder.buildAsyncClient()); this.sessionReceiverBuilder = null; - this.tracerProvider = processorOptions.getTracerProvider(); this.queueName = queueName; this.topicName = topicName; this.subscriptionName = subscriptionName; @@ -345,22 +326,15 @@ public void onNext(ServiceBusMessageContext serviceBusMessageContext) { if (serviceBusMessageContext.hasError()) { handleError(serviceBusMessageContext.getThrowable()); } else { - Context processSpanContext = null; try { ServiceBusReceivedMessageContext serviceBusReceivedMessageContext = new ServiceBusReceivedMessageContext(receiverClient, serviceBusMessageContext); - processSpanContext = - startProcessTracingSpan(serviceBusMessageContext.getMessage(), - receiverClient.getEntityPath(), receiverClient.getFullyQualifiedNamespace()); - if (processSpanContext.getData(SPAN_CONTEXT_KEY).isPresent()) { - serviceBusMessageContext.getMessage().addContext(SPAN_CONTEXT_KEY, processSpanContext); - } processMessage.accept(serviceBusReceivedMessageContext); - endProcessTracingSpan(processSpanContext, Signal.complete()); } catch (Exception ex) { + serviceBusMessageContext.getMessage().addContext(FluxTrace.PROCESS_ERROR_KEY, ex); handleError(new ServiceBusException(ex, ServiceBusErrorSource.USER_CALLBACK)); - endProcessTracingSpan(processSpanContext, Signal.error(ex)); + if (!processorOptions.isDisableAutoComplete()) { LOGGER.warning("Error when processing message. Abandoning message.", ex); abandonMessage(serviceBusMessageContext, receiverClient); @@ -407,54 +381,6 @@ public void onComplete() { } } - private void endProcessTracingSpan(Context processSpanContext, Signal signal) { - if (processSpanContext == null) { - return; - } - - Optional spanScope = processSpanContext.getData(SCOPE_KEY); - // Disposes of the scope when the trace span closes. - if (!spanScope.isPresent() || !tracerProvider.isEnabled()) { - return; - } - if (spanScope.get() instanceof AutoCloseable) { - AutoCloseable close = (AutoCloseable) processSpanContext.getData(SCOPE_KEY).get(); - try { - close.close(); - } catch (Exception exception) { - LOGGER.error("endTracingSpan().close() failed with an error {}", exception); - } - - } else { - LOGGER.warning(String.format(Locale.US, - "Process span scope type is not of type AutoCloseable, but type: %s. Not closing the scope" - + " and span", spanScope.get() != null ? spanScope.getClass() : "null")); - } - tracerProvider.endSpan(processSpanContext, signal); - } - - private Context startProcessTracingSpan(ServiceBusReceivedMessage receivedMessage, String entityPath, - String fullyQualifiedNamespace) { - - Object diagnosticId = receivedMessage.getApplicationProperties().get(DIAGNOSTIC_ID_KEY); - if (tracerProvider == null || !tracerProvider.isEnabled()) { - return Context.NONE; - } - - Context spanContext = Objects.isNull(diagnosticId) ? Context.NONE : tracerProvider.extractContext(diagnosticId.toString(), Context.NONE); - - spanContext = spanContext - .addData(ENTITY_PATH_KEY, entityPath) - .addData(HOST_NAME_KEY, fullyQualifiedNamespace) - .addData(AZ_TRACING_NAMESPACE_KEY, AZ_TRACING_NAMESPACE_VALUE); - spanContext = receivedMessage.getEnqueuedTime() == null - ? spanContext - : spanContext.addData(MESSAGE_ENQUEUED_TIME, - receivedMessage.getEnqueuedTime().toInstant().getEpochSecond()); - - return tracerProvider.startSpan(AZ_TRACING_SERVICE_NAME, spanContext, ProcessKind.PROCESS); - } - private void abandonMessage(ServiceBusMessageContext serviceBusMessageContext, ServiceBusReceiverAsyncClient receiverClient) { try { diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessage.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessage.java index 840df93f1302..48eaa893bec6 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessage.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessage.java @@ -533,6 +533,18 @@ ServiceBusReceivedMessage addContext(String key, Object value) { return this; } + /** + * Adds a new key value pair to the existing context on Message. + * + * @param key The key for this context object + * @param value The value for this context object. + * + * @return The updated {@link ServiceBusMessage}. + * @throws NullPointerException if {@code key} or {@code value} is null. + */ + Context getContext() { + return this.context; + } /** * Gets whether the message has been settled. * diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java index 7cc782acb59b..6ad0133db0ee 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java @@ -10,7 +10,6 @@ import com.azure.core.amqp.implementation.RequestResponseChannelClosedException; import com.azure.core.amqp.implementation.RetryUtil; import com.azure.core.amqp.implementation.StringUtil; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.annotation.ServiceClient; import com.azure.core.util.BinaryData; import com.azure.core.util.CoreUtils; @@ -21,6 +20,7 @@ import com.azure.messaging.servicebus.implementation.LockContainer; import com.azure.messaging.servicebus.implementation.MessagingEntityType; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.implementation.ServiceBusReceiveLink; import com.azure.messaging.servicebus.implementation.ServiceBusReceiveLinkProcessor; import com.azure.messaging.servicebus.models.AbandonOptions; @@ -33,6 +33,7 @@ import reactor.core.publisher.Mono; import java.time.Duration; +import java.time.Instant; import java.time.OffsetDateTime; import java.util.Collections; import java.util.Map; @@ -230,7 +231,7 @@ public final class ServiceBusReceiverAsyncClient implements AutoCloseable { private final MessagingEntityType entityType; private final ReceiverOptions receiverOptions; private final ServiceBusConnectionProcessor connectionProcessor; - private final TracerProvider tracerProvider; + private final ServiceBusReceiverTracer tracer; private final MessageSerializer messageSerializer; private final Runnable onClientClose; private final ServiceBusSessionManager sessionManager; @@ -249,20 +250,20 @@ public final class ServiceBusReceiverAsyncClient implements AutoCloseable { * @param entityType The type of the Service Bus resource. * @param receiverOptions Options when receiving messages. * @param connectionProcessor The AMQP connection to the Service Bus resource. - * @param tracerProvider Tracer for telemetry. + * @param tracer ServiceBus tracing helper * @param messageSerializer Serializes and deserializes Service Bus messages. * @param onClientClose Operation to run when the client completes. */ ServiceBusReceiverAsyncClient(String fullyQualifiedNamespace, String entityPath, MessagingEntityType entityType, ReceiverOptions receiverOptions, ServiceBusConnectionProcessor connectionProcessor, Duration cleanupInterval, - TracerProvider tracerProvider, MessageSerializer messageSerializer, Runnable onClientClose, String identifier) { + ServiceBusReceiverTracer tracer, MessageSerializer messageSerializer, Runnable onClientClose, String identifier) { this.fullyQualifiedNamespace = Objects.requireNonNull(fullyQualifiedNamespace, "'fullyQualifiedNamespace' cannot be null."); this.entityPath = Objects.requireNonNull(entityPath, "'entityPath' cannot be null."); this.entityType = Objects.requireNonNull(entityType, "'entityType' cannot be null."); this.receiverOptions = Objects.requireNonNull(receiverOptions, "'receiveOptions cannot be null.'"); this.connectionProcessor = Objects.requireNonNull(connectionProcessor, "'connectionProcessor' cannot be null."); - this.tracerProvider = Objects.requireNonNull(tracerProvider, "'tracerProvider' cannot be null."); + this.tracer = Objects.requireNonNull(tracer, "'tracer' cannot be null"); this.messageSerializer = Objects.requireNonNull(messageSerializer, "'messageSerializer' cannot be null."); this.onClientClose = Objects.requireNonNull(onClientClose, "'onClientClose' cannot be null."); @@ -280,16 +281,16 @@ public final class ServiceBusReceiverAsyncClient implements AutoCloseable { } ServiceBusReceiverAsyncClient(String fullyQualifiedNamespace, String entityPath, MessagingEntityType entityType, - ReceiverOptions receiverOptions, ServiceBusConnectionProcessor connectionProcessor, Duration cleanupInterval, - TracerProvider tracerProvider, MessageSerializer messageSerializer, Runnable onClientClose, - ServiceBusSessionManager sessionManager) { + ReceiverOptions receiverOptions, ServiceBusConnectionProcessor connectionProcessor, Duration cleanupInterval, + ServiceBusReceiverTracer tracer, MessageSerializer messageSerializer, Runnable onClientClose, + ServiceBusSessionManager sessionManager) { this.fullyQualifiedNamespace = Objects.requireNonNull(fullyQualifiedNamespace, "'fullyQualifiedNamespace' cannot be null."); this.entityPath = Objects.requireNonNull(entityPath, "'entityPath' cannot be null."); this.entityType = Objects.requireNonNull(entityType, "'entityType' cannot be null."); this.receiverOptions = Objects.requireNonNull(receiverOptions, "'receiveOptions cannot be null.'"); this.connectionProcessor = Objects.requireNonNull(connectionProcessor, "'connectionProcessor' cannot be null."); - this.tracerProvider = Objects.requireNonNull(tracerProvider, "'tracerProvider' cannot be null."); + this.tracer = Objects.requireNonNull(tracer, "'tracer' cannot be null"); this.messageSerializer = Objects.requireNonNull(messageSerializer, "'messageSerializer' cannot be null."); this.onClientClose = Objects.requireNonNull(onClientClose, "'onClientClose' cannot be null."); this.sessionManager = Objects.requireNonNull(sessionManager, "'sessionManager' cannot be null."); @@ -596,19 +597,20 @@ Mono peekMessage(String sessionId) { return monoError(LOGGER, new IllegalStateException( String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "peek"))); } - + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) - .flatMap(channel -> { - final long sequence = lastPeekedSequenceNumber.get() + 1; + .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) + .flatMap(channel -> { + final long sequence = lastPeekedSequenceNumber.get() + 1; - LOGGER.atVerbose() - .addKeyValue(SEQUENCE_NUMBER_KEY, sequence) - .log("Peek message."); + LOGGER.atVerbose() + .addKeyValue(SEQUENCE_NUMBER_KEY, sequence) + .log("Peek message."); - return channel.peek(sequence, sessionId, getLinkName(sessionId)); - }) - .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)) + return channel.peek(sequence, sessionId, getLinkName(sessionId)); + }) + .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)) + .doOnEach(signal -> tracer.reportReceiveSpan("ServiceBus.peekMessage", startTime, signal)) .handle((message, sink) -> { final long current = lastPeekedSequenceNumber .updateAndGet(value -> Math.max(value, message.getSequenceNumber())); @@ -654,10 +656,12 @@ Mono peekMessage(long sequenceNumber, String sessionI return monoError(LOGGER, new IllegalStateException( String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "peekAt"))); } + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) - .flatMap(node -> node.peek(sequenceNumber, sessionId, getLinkName(sessionId))) - .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); + .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) + .flatMap(node -> node.peek(sequenceNumber, sessionId, getLinkName(sessionId))) + .doOnEach(signal -> tracer.reportReceiveSpan("ServiceBus.peekMessage", startTime, signal)) + .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); } /** @@ -673,7 +677,9 @@ Mono peekMessage(long sequenceNumber, String sessionI * @see Message browsing */ public Flux peekMessages(int maxMessages) { - return peekMessages(maxMessages, receiverOptions.getSessionId()); + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; + return peekMessages(maxMessages, receiverOptions.getSessionId()) + .doOnEach(signal -> tracer.reportReceiveSpan("ServiceBus.peekMessages", startTime, signal)); } /** @@ -698,7 +704,8 @@ Flux peekMessages(int maxMessages, String sessionId) return fluxError(LOGGER, new IllegalArgumentException("'maxMessages' is not positive.")); } - return connectionProcessor + return + connectionProcessor .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) .flatMapMany(node -> { final long nextSequenceNumber = lastPeekedSequenceNumber.get() + 1; @@ -770,10 +777,11 @@ Flux peekMessages(int maxMessages, long sequenceNumbe if (maxMessages <= 0) { return fluxError(LOGGER, new IllegalArgumentException("'maxMessages' is not positive.")); } - + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; return connectionProcessor .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) .flatMapMany(node -> node.peek(sequenceNumber, sessionId, getLinkName(sessionId), maxMessages)) + .doOnEach(signal -> tracer.reportReceiveSpan("ServiceBus.peekMessages", startTime, signal)) .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); } @@ -842,12 +850,14 @@ Flux receiveMessagesWithContext(int highTide) { ? sessionManager.receive() : getOrCreateConsumer().receive().map(ServiceBusMessageContext::new); + final Flux messageFluxWithTracing = new FluxTrace(messageFlux, tracer); final Flux withAutoLockRenewal; + if (!receiverOptions.isSessionReceiver() && receiverOptions.isAutoLockRenewEnabled()) { - withAutoLockRenewal = new FluxAutoLockRenew(messageFlux, receiverOptions, + withAutoLockRenewal = new FluxAutoLockRenew(messageFluxWithTracing, receiverOptions, renewalContainer, this::renewMessageLock); } else { - withAutoLockRenewal = messageFlux; + withAutoLockRenewal = messageFluxWithTracing; } Flux result; @@ -899,6 +909,8 @@ Mono receiveDeferredMessage(long sequenceNumber, Stri return monoError(LOGGER, new IllegalStateException( String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "receiveDeferredMessage"))); } + + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; return connectionProcessor .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) .flatMap(node -> node.receiveDeferredMessages(receiverOptions.getReceiveMode(), @@ -915,6 +927,7 @@ sessionId, getLinkName(sessionId), Collections.singleton(sequenceNumber)).last() return receivedMessage; }) + .doOnEach(signal -> tracer.reportReceiveSpan("ServiceBus.receiveDeferredMessage", startTime, signal)) .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); } @@ -931,7 +944,9 @@ sessionId, getLinkName(sessionId), Collections.singleton(sequenceNumber)).last() * @throws ServiceBusException if deferred messages cannot be received. */ public Flux receiveDeferredMessages(Iterable sequenceNumbers) { - return receiveDeferredMessages(sequenceNumbers, receiverOptions.getSessionId()); + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; + return receiveDeferredMessages(sequenceNumbers, receiverOptions.getSessionId()) + .doOnEach(signal -> tracer.reportReceiveSpan("ServiceBus.receiveDeferredMessages", startTime, signal)); } /** @@ -1015,7 +1030,7 @@ public Mono renewMessageLock(ServiceBusReceivedMessage message) return monoError(LOGGER, new IllegalStateException(errorMessage)); } - return renewMessageLock(message.getLockToken()) + return tracer.traceMonoWithLink(renewMessageLock(message.getLockToken()), message, "ServiceBus.renewMessageLock") .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK)); } @@ -1080,7 +1095,7 @@ public Mono renewMessageLock(ServiceBusReceivedMessage message, Duration m renewalContainer.addOrUpdate(message.getLockToken(), OffsetDateTime.now().plus(maxLockRenewalDuration), operation); - return operation.getCompletionOperation() + return tracer.traceMonoWithLink(operation.getCompletionOperation(), message, "ServiceBus.renewMessageLock") .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK)); } @@ -1162,10 +1177,11 @@ public Mono createTransaction() { String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "createTransaction"))); } - return connectionProcessor - .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) - .flatMap(transactionSession -> transactionSession.createTransaction()) - .map(transaction -> new ServiceBusTransactionContext(transaction.getTransactionId())) + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) + .flatMap(transactionSession -> transactionSession.createTransaction()) + .map(transaction -> new ServiceBusTransactionContext(transaction.getTransactionId())), + null, "ServiceBus.commitTransaction") .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); } @@ -1215,10 +1231,11 @@ public Mono commitTransaction(ServiceBusTransactionContext transactionCont return monoError(LOGGER, new NullPointerException("'transactionContext.transactionId' cannot be null.")); } - return connectionProcessor - .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) - .flatMap(transactionSession -> transactionSession.commitTransaction(new AmqpTransaction( - transactionContext.getTransactionId()))) + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) + .flatMap(transactionSession -> transactionSession.commitTransaction(new AmqpTransaction( + transactionContext.getTransactionId()))), + null, "ServiceBus.commitTransaction") .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); } @@ -1267,11 +1284,13 @@ public Mono rollbackTransaction(ServiceBusTransactionContext transactionCo return monoError(LOGGER, new NullPointerException("'transactionContext.transactionId' cannot be null.")); } - return connectionProcessor - .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) - .flatMap(transactionSession -> transactionSession.rollbackTransaction(new AmqpTransaction( - transactionContext.getTransactionId()))) - .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) + .flatMap(transactionSession -> transactionSession.rollbackTransaction(new AmqpTransaction( + transactionContext.getTransactionId()))), + null, "ServiceBus.rollbackTransaction") + .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RECEIVE)); + } /** @@ -1429,7 +1448,7 @@ deadLetterErrorDescription, propertiesToModify, sessionId, getLinkName(sessionId } } - return updateDispositionOperation + return tracer.traceSettlement(updateDispositionOperation, message, dispositionStatus) .onErrorMap(throwable -> { if (throwable instanceof ServiceBusException) { return throwable; @@ -1551,9 +1570,10 @@ Mono renewSessionLock(String sessionId) { ? sessionManager.getLinkName(sessionId) : null; - return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) - .flatMap(channel -> channel.renewSessionLock(sessionId, linkName)) + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) + .flatMap(channel -> channel.renewSessionLock(sessionId, linkName)), + null, "ServiceBus.renewSessionLock") .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK)); } @@ -1578,7 +1598,7 @@ Mono renewSessionLock(String sessionId, Duration maxLockRenewalDuration) { maxLockRenewalDuration, true, this::renewSessionLock); renewalContainer.addOrUpdate(sessionId, OffsetDateTime.now().plus(maxLockRenewalDuration), operation); - return operation.getCompletionOperation() + return tracer.traceMonoWithLink(operation.getCompletionOperation(), null, "ServiceBus.renewSessionLock") .onErrorMap(throwable -> mapError(throwable, ServiceBusErrorSource.RENEW_LOCK)); } @@ -1593,9 +1613,10 @@ Mono setSessionState(String sessionId, byte[] sessionState) { ? sessionManager.getLinkName(sessionId) : null; - return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) - .flatMap(channel -> channel.setSessionState(sessionId, sessionState, linkName)) + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.getManagementNode(entityPath, entityType)) + .flatMap(channel -> channel.setSessionState(sessionId, sessionState, linkName)), + null, "ServiceBus.setSessionState") .onErrorMap((err) -> mapError(err, ServiceBusErrorSource.RECEIVE)); } @@ -1617,7 +1638,12 @@ Mono getSessionState(String sessionId) { .flatMap(channel -> channel.getSessionState(sessionId, getLinkName(sessionId))); } - return result.onErrorMap((err) -> mapError(err, ServiceBusErrorSource.RECEIVE)); + return tracer.traceMonoWithLink(result, null, "ServiceBus.setSessionState") + .onErrorMap((err) -> mapError(err, ServiceBusErrorSource.RECEIVE)); + } + + ServiceBusReceiverTracer getTracer() { + return tracer; } /** diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java index 80bf6f88450c..f1e80e8ee528 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverClient.java @@ -7,6 +7,7 @@ import com.azure.core.util.IterableStream; import com.azure.core.util.logging.ClientLogger; import com.azure.messaging.servicebus.ServiceBusClientBuilder.ServiceBusReceiverClientBuilder; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.models.AbandonOptions; import com.azure.messaging.servicebus.models.CompleteOptions; import com.azure.messaging.servicebus.models.DeadLetterOptions; @@ -16,6 +17,7 @@ import reactor.core.publisher.Sinks; import java.time.Duration; +import java.time.Instant; import java.time.OffsetDateTime; import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; @@ -64,6 +66,7 @@ public final class ServiceBusReceiverClient implements AutoCloseable { /* To ensure synchronousMessageSubscriber is subscribed only once. */ private final AtomicBoolean syncSubscribed = new AtomicBoolean(false); + private final ServiceBusReceiverTracer tracer; /** * Creates a synchronous receiver given its asynchronous counterpart. * @@ -77,6 +80,7 @@ public final class ServiceBusReceiverClient implements AutoCloseable { this.asyncClient = Objects.requireNonNull(asyncClient, "'asyncClient' cannot be null."); this.operationTimeout = Objects.requireNonNull(operationTimeout, "'operationTimeout' cannot be null."); this.isPrefetchDisabled = isPrefetchDisabled; + this.tracer = asyncClient.getTracer(); } /** @@ -392,13 +396,15 @@ IterableStream peekMessages(int maxMessages, String s "'maxMessages' cannot be less than or equal to 0. maxMessages: " + maxMessages)); } + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; final Flux messages = asyncClient.peekMessages(maxMessages, sessionId) .timeout(operationTimeout); - // Subscribe so we can kick off this operation. + final Flux tracedMessages = tracer.reportSyncReceiverSpan("ServiceBus.peekMessages", startTime, messages); + // Subscribe to message flux so we can kick off this operation, but not to tracing - caller subscriber will take care of it. messages.subscribe(); - return new IterableStream<>(messages); + return new IterableStream<>(tracedMessages); } /** @@ -442,13 +448,15 @@ IterableStream peekMessages(int maxMessages, long seq "'maxMessages' cannot be less than or equal to 0. maxMessages: " + maxMessages)); } + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; final Flux messages = asyncClient.peekMessages(maxMessages, sequenceNumber, sessionId).timeout(operationTimeout); - // Subscribe so we can kick off this operation. + final Flux tracedMessages = tracer.reportSyncReceiverSpan("ServiceBus.peekMessages", startTime, messages); + // Subscribe to message flux so we can kick off this operation, but not to tracing - caller subscriber will take care of it. messages.subscribe(); - return new IterableStream<>(messages); + return new IterableStream<>(tracedMessages); } /** @@ -498,13 +506,22 @@ public IterableStream receiveMessages(int maxMessages new IllegalArgumentException("'maxWaitTime' cannot be zero or less. maxWaitTime: " + maxWaitTime)); } + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; + // There are two subscribers to this emitter. One is the timeout between messages subscription in // SynchronousReceiverWork.start() and the other is the IterableStream(emitter.asFlux()); // Since the subscriptions may happen at different times, we want to replay results to downstream subscribers. final Sinks.Many emitter = Sinks.many().replay().all(); + queueWork(maxMessages, maxWaitTime, emitter); - return new IterableStream<>(emitter.asFlux()); + final Flux messagesFlux = emitter.asFlux(); + final Flux tracedMessages = tracer.reportSyncReceiverSpan("ServiceBus.receiveMessages", startTime, messagesFlux); + + // Subscribe to message flux so we can kick off this operation, but not to tracing - caller subscriber will take care of it. + messagesFlux.subscribe(); + + return new IterableStream<>(tracedMessages); } /** @@ -571,13 +588,16 @@ public IterableStream receiveDeferredMessageBatch(Ite */ IterableStream receiveDeferredMessageBatch(Iterable sequenceNumbers, String sessionId) { + + final Instant startTime = tracer.isEnabled() ? Instant.now() : null; final Flux messages = asyncClient.receiveDeferredMessages(sequenceNumbers, sessionId).timeout(operationTimeout); + final Flux tracedMessages = tracer.reportSyncReceiverSpan("ServiceBus.receiveDeferredMessageBatch", startTime, messages); // Subscribe so we can kick off this operation. messages.subscribe(); - return new IterableStream<>(messages); + return new IterableStream<>(tracedMessages); } /** diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClient.java index 71a896dfe0b5..597160d663eb 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClient.java @@ -11,12 +11,9 @@ import com.azure.core.amqp.implementation.AmqpSendLink; import com.azure.core.amqp.implementation.ErrorContextProvider; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.annotation.ServiceClient; -import com.azure.core.util.Context; import com.azure.core.util.CoreUtils; import com.azure.core.util.logging.ClientLogger; -import com.azure.core.util.tracing.ProcessKind; import com.azure.messaging.servicebus.implementation.MessagingEntityType; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; import com.azure.messaging.servicebus.models.CreateMessageBatchOptions; @@ -48,12 +45,7 @@ import static com.azure.core.amqp.implementation.RetryUtil.withRetry; import static com.azure.core.util.FluxUtil.fluxError; import static com.azure.core.util.FluxUtil.monoError; -import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; -import static com.azure.core.util.tracing.Tracer.ENTITY_PATH_KEY; -import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; -import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; import static com.azure.messaging.servicebus.implementation.Messages.INVALID_OPERATION_DISPOSED_SENDER; -import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_SERVICE_NAME; /** * An asynchronous client to send messages to a Service Bus resource. @@ -170,7 +162,7 @@ public final class ServiceBusSenderAsyncClient implements AutoCloseable { private static final ClientLogger LOGGER = new ClientLogger(ServiceBusSenderAsyncClient.class); private final AtomicReference linkName = new AtomicReference<>(); private final AtomicBoolean isDisposed = new AtomicBoolean(); - private final TracerProvider tracerProvider; + private final ServiceBusSenderTracer tracer; private final MessageSerializer messageSerializer; private final AmqpRetryOptions retryOptions; private final AmqpRetryPolicy retryPolicy; @@ -185,7 +177,7 @@ public final class ServiceBusSenderAsyncClient implements AutoCloseable { * Creates a new instance of this {@link ServiceBusSenderAsyncClient} that sends messages to a Service Bus entity. */ ServiceBusSenderAsyncClient(String entityName, MessagingEntityType entityType, - ServiceBusConnectionProcessor connectionProcessor, AmqpRetryOptions retryOptions, TracerProvider tracerProvider, + ServiceBusConnectionProcessor connectionProcessor, AmqpRetryOptions retryOptions, ServiceBusSenderTracer tracer, MessageSerializer messageSerializer, Runnable onClientClose, String viaEntityName, String identifier) { // Caching the created link so we don't invoke another link creation. this.messageSerializer = Objects.requireNonNull(messageSerializer, @@ -194,7 +186,7 @@ public final class ServiceBusSenderAsyncClient implements AutoCloseable { this.entityName = Objects.requireNonNull(entityName, "'entityPath' cannot be null."); this.connectionProcessor = Objects.requireNonNull(connectionProcessor, "'connectionProcessor' cannot be null."); - this.tracerProvider = tracerProvider; + this.tracer = Objects.requireNonNull(tracer, "'tracer' cannot be null."); this.retryPolicy = getRetryPolicy(retryOptions); this.entityType = entityType; this.viaEntityName = viaEntityName; @@ -408,7 +400,7 @@ public Mono createMessageBatch(CreateMessageBatchOptions : maximumLinkSize; return Mono.just( - new ServiceBusMessageBatch(batchSize, link::getErrorContext, tracerProvider, messageSerializer, + new ServiceBusMessageBatch(batchSize, link::getErrorContext, tracer, messageSerializer, entityName, getFullyQualifiedNamespace())); })).onErrorMap(this::mapError); } @@ -518,10 +510,12 @@ public Flux scheduleMessages(Iterable messages, OffsetD } return messageBatch; }) - .flatMapMany(messageBatch -> connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityName, entityType)) - .flatMapMany(managementNode -> managementNode.schedule(messageBatch.getMessages(), scheduledEnqueueTime, - messageBatch.getMaxSizeInBytes(), linkName.get(), transactionContext)) + .flatMapMany(messageBatch -> + tracer.traceFluxWithLinks(connectionProcessor + .flatMap(connection -> connection.getManagementNode(entityName, entityType)) + .flatMapMany(managementNode -> managementNode.schedule(messageBatch.getMessages(), scheduledEnqueueTime, + messageBatch.getMaxSizeInBytes(), linkName.get(), transactionContext)), + messageBatch, "ServiceBus.scheduleMessages") ).onErrorMap(this::mapError); } @@ -545,10 +539,12 @@ public Mono cancelScheduledMessage(long sequenceNumber) { return monoError(LOGGER, new IllegalArgumentException("'sequenceNumber' cannot be negative.")); } - return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityName, entityType)) - .flatMap(managementNode -> managementNode.cancelScheduledMessages( - Collections.singletonList(sequenceNumber), linkName.get())).onErrorMap(this::mapError); + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.getManagementNode(entityName, entityType)) + .flatMap(managementNode -> managementNode.cancelScheduledMessages( + Collections.singletonList(sequenceNumber), linkName.get())), + null, "ServiceBus.cancelScheduledMessage") + .onErrorMap(this::mapError); } /** @@ -572,9 +568,10 @@ public Mono cancelScheduledMessages(Iterable sequenceNumbers) { return monoError(LOGGER, new NullPointerException("'messages' cannot be null.")); } - return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityName, entityType)) - .flatMap(managementNode -> managementNode.cancelScheduledMessages(sequenceNumbers, linkName.get())) + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.getManagementNode(entityName, entityType)) + .flatMap(managementNode -> managementNode.cancelScheduledMessages(sequenceNumbers, linkName.get())), + null, "ServiceBus.cancelScheduledMessages") .onErrorMap(this::mapError); } @@ -595,10 +592,11 @@ public Mono createTransaction() { String.format(INVALID_OPERATION_DISPOSED_SENDER, "createTransaction"))); } - return connectionProcessor - .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) - .flatMap(transactionSession -> transactionSession.createTransaction()) - .map(transaction -> new ServiceBusTransactionContext(transaction.getTransactionId())) + return tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) + .flatMap(transactionSession -> transactionSession.createTransaction()) + .map(transaction -> new ServiceBusTransactionContext(transaction.getTransactionId())), + null, "ServiceBus.createTransaction") .onErrorMap(this::mapError); } @@ -627,10 +625,12 @@ public Mono commitTransaction(ServiceBusTransactionContext transactionCont return monoError(LOGGER, new NullPointerException("'transactionContext.transactionId' cannot be null.")); } - return connectionProcessor - .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) - .flatMap(transactionSession -> transactionSession.commitTransaction(new AmqpTransaction( - transactionContext.getTransactionId()))) + return + tracer.traceMonoWithLink(connectionProcessor + .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) + .flatMap(transactionSession -> transactionSession.commitTransaction(new AmqpTransaction( + transactionContext.getTransactionId()))), + null, "ServiceBus.commitTransaction") .onErrorMap(this::mapError); } @@ -659,10 +659,12 @@ public Mono rollbackTransaction(ServiceBusTransactionContext transactionCo return monoError(LOGGER, new NullPointerException("'transactionContext.transactionId' cannot be null.")); } - return connectionProcessor - .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) - .flatMap(transactionSession -> transactionSession.rollbackTransaction(new AmqpTransaction( - transactionContext.getTransactionId()))) + return tracer.traceMonoWithLink( + connectionProcessor + .flatMap(connection -> connection.createSession(TRANSACTION_LINK_NAME)) + .flatMap(transactionSession -> transactionSession.rollbackTransaction(new AmqpTransaction( + transactionContext.getTransactionId()))), + null, "ServiceBus.rollbackTransaction") .onErrorMap(this::mapError); } @@ -705,18 +707,20 @@ private Mono scheduleMessageInternal(ServiceBusMessage message, OffsetDate return monoError(LOGGER, new NullPointerException("'scheduledEnqueueTime' cannot be null.")); } - return getSendLink() - .flatMap(link -> link.getLinkSize().flatMap(size -> { - int maxSize = size > 0 - ? size - : MAX_MESSAGE_LENGTH_BYTES; + return tracer.traceMonoWithLink(getSendLink() + .flatMap(link -> link.getLinkSize().flatMap(size -> { + int maxSize = size > 0 + ? size + : MAX_MESSAGE_LENGTH_BYTES; - return connectionProcessor - .flatMap(connection -> connection.getManagementNode(entityName, entityType)) - .flatMap(managementNode -> managementNode.schedule(Arrays.asList(message), scheduledEnqueueTime, - maxSize, link.getLinkName(), transactionContext) - .next()); - })).onErrorMap(this::mapError); + return connectionProcessor + .flatMap(connection -> connection.getManagementNode(entityName, entityType)) + .flatMap(managementNode -> managementNode.schedule(Arrays.asList(message), scheduledEnqueueTime, + maxSize, link.getLinkName(), transactionContext) + .next()); + })), + message, "ServiceBus.scheduleMessage") + .onErrorMap(this::mapError); } /** @@ -735,11 +739,6 @@ private Mono sendInternal(ServiceBusMessageBatch batch, ServiceBusTransact return monoError(LOGGER, new NullPointerException("'batch' cannot be null.")); } - final boolean isTracingEnabled = tracerProvider.isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; - if (batch.getMessages().isEmpty()) { LOGGER.info("Cannot send an EventBatch that is empty."); return Mono.empty(); @@ -749,16 +748,10 @@ private Mono sendInternal(ServiceBusMessageBatch batch, ServiceBusTransact .addKeyValue("batchSize", batch.getCount()) .log("Sending batch."); - AtomicReference sharedContext = new AtomicReference<>(Context.NONE); final List messages = Collections.synchronizedList(new ArrayList<>()); + + batch.getMessages().forEach(serviceBusMessage -> { - if (isTracingEnabled) { - parentContext.set(serviceBusMessage.getContext()); - if (sharedContext.get().equals(Context.NONE)) { - sharedContext.set(tracerProvider.getSharedSpanBuilder(SERVICE_BASE_NAME, parentContext.get())); - } - tracerProvider.addSpanLinks(sharedContext.get().addData(SPAN_CONTEXT_KEY, serviceBusMessage.getContext())); - } final org.apache.qpid.proton.message.Message message = messageSerializer.serialize(serviceBusMessage); final MessageAnnotations messageAnnotations = message.getMessageAnnotations() == null ? new MessageAnnotations(new HashMap<>()) @@ -768,17 +761,6 @@ private Mono sendInternal(ServiceBusMessageBatch batch, ServiceBusTransact messages.add(message); }); - if (isTracingEnabled) { - final Context finalSharedContext = sharedContext.get().equals(Context.NONE) - ? Context.NONE - : sharedContext.get() - .addData(ENTITY_PATH_KEY, entityName) - .addData(HOST_NAME_KEY, connectionProcessor.getFullyQualifiedNamespace()) - .addData(AZ_TRACING_NAMESPACE_KEY, AZ_TRACING_NAMESPACE_VALUE); - // Start send span and store updated context - parentContext.set(tracerProvider.startSpan(AZ_TRACING_SERVICE_NAME, finalSharedContext, ProcessKind.SEND)); - } - final Mono sendMessage = getSendLink().flatMap(link -> { if (transactionContext != null && transactionContext.getTransactionId() != null) { final TransactionalState deliveryState = new TransactionalState(); @@ -793,14 +775,11 @@ private Mono sendInternal(ServiceBusMessageBatch batch, ServiceBusTransact } }); - return withRetry(sendMessage, retryOptions, - String.format("entityPath[%s], partitionId[%s]: Sending messages timed out.", entityName, - batch.getCount())) - .doOnEach(signal -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), signal); - } - }).onErrorMap(this::mapError); + return tracer.traceMonoWithLinks( + withRetry(sendMessage, retryOptions, + String.format("entityPath[%s], partitionId[%s]: Sending messages timed out.", entityName, batch.getCount())), + batch, "ServiceBus.send") + .onErrorMap(this::mapError); } private Mono sendInternal(Flux messages, ServiceBusTransactionContext transactionContext) { @@ -815,7 +794,7 @@ private Mono sendInternal(Flux messages, ServiceBusTran final CreateMessageBatchOptions batchOptions = new CreateMessageBatchOptions() .setMaximumSizeInBytes(batchSize); return messages.collect(new AmqpMessageCollector(batchOptions, 1, - link::getErrorContext, tracerProvider, messageSerializer, entityName, + link::getErrorContext, tracer, messageSerializer, entityName, link.getHostname())); }) .flatMap(list -> sendInternalBatch(Flux.fromIterable(list), transactionContext))) @@ -855,7 +834,7 @@ private static class AmqpMessageCollector implements Collector 0 ? options.getMaximumSizeInBytes() : MAX_MESSAGE_LENGTH_BYTES; this.contextProvider = contextProvider; - this.tracerProvider = tracerProvider; + this.tracer = tracer; this.serializer = serializer; this.entityPath = entityPath; this.hostname = hostname; - currentBatch = new ServiceBusMessageBatch(maxMessageSize, contextProvider, tracerProvider, serializer, + currentBatch = new ServiceBusMessageBatch(maxMessageSize, contextProvider, tracer, serializer, entityPath, hostname); } @@ -900,7 +879,7 @@ public BiConsumer, ServiceBusMessage> accumulator() contextProvider.getErrorContext()); } - currentBatch = new ServiceBusMessageBatch(maxMessageSize, contextProvider, tracerProvider, serializer, + currentBatch = new ServiceBusMessageBatch(maxMessageSize, contextProvider, tracer, serializer, entityPath, hostname); currentBatch.tryAddMessage(event); list.add(batch); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderTracer.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderTracer.java new file mode 100644 index 000000000000..da20924707f7 --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSenderTracer.java @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.Context; +import com.azure.core.util.tracing.ProcessKind; +import com.azure.core.util.tracing.Tracer; +import com.azure.messaging.servicebus.implementation.ServiceBusTracer; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.util.Optional; + +import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; +import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; + +class ServiceBusSenderTracer extends ServiceBusTracer { + ServiceBusSenderTracer(Tracer tracer, String fullyQualifiedName, String entityPath) { + super(tracer, fullyQualifiedName, entityPath); + } + + Mono traceMonoWithLink(Mono publisher, ServiceBusMessage message, String spanName) { + if (tracer == null) { + return publisher; + } + + return publisher + .doOnEach(this::endSpan) + .contextWrite(reactor.util.context.Context.of(REACTOR_PARENT_TRACE_CONTEXT_KEY, startSpanWithLink(spanName, message, Context.NONE))); + } + + Mono traceMonoWithLinks(Mono publisher, ServiceBusMessageBatch batch, String spanName) { + if (tracer == null) { + return publisher; + } + + return publisher + .doOnEach(this::endSpan) + .contextWrite(reactor.util.context.Context.of(REACTOR_PARENT_TRACE_CONTEXT_KEY, startSpanWithLinks(spanName, batch, Context.NONE))); + } + + Flux traceFluxWithLinks(Flux publisher, ServiceBusMessageBatch batch, String spanName) { + if (tracer == null) { + return publisher; + } + + return publisher + .doOnEach(this::endSpan) + .contextWrite(reactor.util.context.Context.of(REACTOR_PARENT_TRACE_CONTEXT_KEY, startSpanWithLinks(spanName, batch, Context.NONE))); + } + + /** + * Used in ServiceBusMessageBatch.tryAddMessage() to start tracing for to-be-sent out messages. + */ + void createMessageSpan(ServiceBusMessage serviceBusMessage) { + Context messageContext = serviceBusMessage.getContext(); + if (tracer == null || messageContext == null || messageContext.getData(SPAN_CONTEXT_KEY).isPresent()) { + // if message has context (in case of retries), don't start a message span or add a new context + return; + } + + String traceparent = getTraceparent(serviceBusMessage.getApplicationProperties()); + if (traceparent != null) { + // if message has context (in case of retries) or if user supplied it, don't start a message span or add a new context + return; + } + + // Starting the span makes the sampling decision (nothing is logged at this time) + Context newMessageContext = setAttributes(messageContext); + + Context eventSpanContext = tracer.start("ServiceBus.message", newMessageContext, ProcessKind.MESSAGE); + Optional traceparentOpt = eventSpanContext.getData(DIAGNOSTIC_ID_KEY); + + if (traceparentOpt.isPresent()) { + serviceBusMessage.getApplicationProperties().put(DIAGNOSTIC_ID_KEY, traceparentOpt.get().toString()); + serviceBusMessage.getApplicationProperties().put(TRACEPARENT_KEY, traceparentOpt.get().toString()); + + endSpan(null, eventSpanContext, null); + + Optional spanContext = eventSpanContext.getData(SPAN_CONTEXT_KEY); + if (spanContext.isPresent()) { + serviceBusMessage.addContext(SPAN_CONTEXT_KEY, spanContext.get()); + } + } + } + + private Context startSpanWithLink(String name, ServiceBusMessage message, Context context) { + Context spanBuilder = getBuilder(name, context); + addLink(message, spanBuilder); + + return tracer.start(name, spanBuilder, ProcessKind.SEND); + } + + private void addLink(ServiceBusMessage message, Context spanBuilder) { + if (message == null) { + return; + } + + String traceparent = getTraceparent(message.getApplicationProperties()); + + if (traceparent == null) { + createMessageSpan(message); + } + + addLink(message.getApplicationProperties(), null, spanBuilder, Context.NONE); + } + + private Context startSpanWithLinks(String name, ServiceBusMessageBatch batch, Context context) { + Context spanBuilder = getBuilder(name, context); + if (batch != null) { + for (ServiceBusMessage message : batch.getMessages()) { + addLink(message, spanBuilder); + } + } + + return tracer.start(name, spanBuilder, ProcessKind.SEND); + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionManager.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionManager.java index 8ff91d60a070..a4176b51be31 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionManager.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionManager.java @@ -9,7 +9,6 @@ import com.azure.core.amqp.exception.SessionErrorContext; import com.azure.core.amqp.implementation.MessageSerializer; import com.azure.core.amqp.implementation.StringUtil; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.util.logging.ClientLogger; import com.azure.messaging.servicebus.implementation.DispositionStatus; import com.azure.messaging.servicebus.implementation.MessageUtils; @@ -61,7 +60,6 @@ class ServiceBusSessionManager implements AutoCloseable { private final ServiceBusReceiveLink receiveLink; private final ServiceBusConnectionProcessor connectionProcessor; private final Duration operationTimeout; - private final TracerProvider tracerProvider; private final MessageSerializer messageSerializer; private final String identifier; @@ -81,14 +79,13 @@ class ServiceBusSessionManager implements AutoCloseable { private volatile Flux receiveFlux; ServiceBusSessionManager(String entityPath, MessagingEntityType entityType, - ServiceBusConnectionProcessor connectionProcessor, TracerProvider tracerProvider, + ServiceBusConnectionProcessor connectionProcessor, MessageSerializer messageSerializer, ReceiverOptions receiverOptions, ServiceBusReceiveLink receiveLink, String identifier) { this.entityPath = entityPath; this.entityType = entityType; this.receiverOptions = receiverOptions; this.connectionProcessor = connectionProcessor; this.operationTimeout = connectionProcessor.getRetryOptions().getTryTimeout(); - this.tracerProvider = tracerProvider; this.messageSerializer = messageSerializer; this.maxSessionLockRenewDuration = receiverOptions.getMaxLockRenewDuration(); this.identifier = identifier; @@ -113,9 +110,9 @@ class ServiceBusSessionManager implements AutoCloseable { } ServiceBusSessionManager(String entityPath, MessagingEntityType entityType, - ServiceBusConnectionProcessor connectionProcessor, TracerProvider tracerProvider, + ServiceBusConnectionProcessor connectionProcessor, MessageSerializer messageSerializer, ReceiverOptions receiverOptions, String identifier) { - this(entityPath, entityType, connectionProcessor, tracerProvider, + this(entityPath, entityType, connectionProcessor, messageSerializer, receiverOptions, null, identifier); } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClient.java index 9be0ab93b465..aa53f3815166 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClient.java @@ -6,7 +6,6 @@ import com.azure.core.amqp.AmqpRetryOptions; import com.azure.core.amqp.exception.AmqpException; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.annotation.ReturnType; import com.azure.core.annotation.ServiceClient; import com.azure.core.annotation.ServiceMethod; @@ -15,6 +14,7 @@ import com.azure.messaging.servicebus.implementation.MessagingEntityType; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; import com.azure.messaging.servicebus.implementation.ServiceBusConstants; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import reactor.core.publisher.Mono; import java.util.Objects; @@ -94,7 +94,7 @@ public final class ServiceBusSessionReceiverAsyncClient implements AutoCloseable private final MessagingEntityType entityType; private final ReceiverOptions receiverOptions; private final ServiceBusConnectionProcessor connectionProcessor; - private final TracerProvider tracerProvider; + private final ServiceBusReceiverTracer tracer; private final MessageSerializer messageSerializer; private final Runnable onClientClose; private final ServiceBusSessionManager unNamedSessionManager; // for acceptNextSession() @@ -102,7 +102,7 @@ public final class ServiceBusSessionReceiverAsyncClient implements AutoCloseable ServiceBusSessionReceiverAsyncClient(String fullyQualifiedNamespace, String entityPath, MessagingEntityType entityType, ReceiverOptions receiverOptions, - ServiceBusConnectionProcessor connectionProcessor, TracerProvider tracerProvider, + ServiceBusConnectionProcessor connectionProcessor, ServiceBusReceiverTracer tracer, MessageSerializer messageSerializer, Runnable onClientClose, String identifier) { this.fullyQualifiedNamespace = Objects.requireNonNull(fullyQualifiedNamespace, "'fullyQualifiedNamespace' cannot be null."); @@ -110,11 +110,11 @@ public final class ServiceBusSessionReceiverAsyncClient implements AutoCloseable this.entityType = Objects.requireNonNull(entityType, "'entityType' cannot be null."); this.receiverOptions = Objects.requireNonNull(receiverOptions, "'receiveOptions cannot be null.'"); this.connectionProcessor = Objects.requireNonNull(connectionProcessor, "'connectionProcessor' cannot be null."); - this.tracerProvider = Objects.requireNonNull(tracerProvider, "'tracerProvider' cannot be null."); + this.tracer = Objects.requireNonNull(tracer, "'tracer' cannot be null."); this.messageSerializer = Objects.requireNonNull(messageSerializer, "'messageSerializer' cannot be null."); this.onClientClose = Objects.requireNonNull(onClientClose, "'onClientClose' cannot be null."); this.unNamedSessionManager = new ServiceBusSessionManager(entityPath, entityType, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, identifier); + messageSerializer, receiverOptions, identifier); this.identifier = identifier; } @@ -131,18 +131,20 @@ public final class ServiceBusSessionReceiverAsyncClient implements AutoCloseable */ @ServiceMethod(returns = ReturnType.SINGLE) public Mono acceptNextSession() { - return unNamedSessionManager.getActiveLink().flatMap(receiveLink -> receiveLink.getSessionId() + return tracer.traceMonoWithLink(unNamedSessionManager.getActiveLink().flatMap(receiveLink -> receiveLink.getSessionId() .map(sessionId -> { final ReceiverOptions newReceiverOptions = new ReceiverOptions(receiverOptions.getReceiveMode(), receiverOptions.getPrefetchCount(), receiverOptions.getMaxLockRenewDuration(), receiverOptions.isEnableAutoComplete(), sessionId, null); final ServiceBusSessionManager sessionSpecificManager = new ServiceBusSessionManager(entityPath, - entityType, connectionProcessor, tracerProvider, messageSerializer, newReceiverOptions, + entityType, connectionProcessor, messageSerializer, newReceiverOptions, receiveLink, identifier); return new ServiceBusReceiverAsyncClient(fullyQualifiedNamespace, entityPath, entityType, newReceiverOptions, connectionProcessor, ServiceBusConstants.OPERATION_TIMEOUT, - tracerProvider, messageSerializer, () -> { }, sessionSpecificManager); - })); + tracer, messageSerializer, () -> { }, sessionSpecificManager); + })), + null, "ServiceBus.acceptNextSession"); + } /** @@ -172,12 +174,15 @@ public Mono acceptSession(String sessionId) { receiverOptions.getPrefetchCount(), receiverOptions.getMaxLockRenewDuration(), receiverOptions.isEnableAutoComplete(), sessionId, null); final ServiceBusSessionManager sessionSpecificManager = new ServiceBusSessionManager(entityPath, entityType, - connectionProcessor, tracerProvider, messageSerializer, newReceiverOptions, identifier); + connectionProcessor, messageSerializer, newReceiverOptions, identifier); - return sessionSpecificManager.getActiveLink().map(receiveLink -> new ServiceBusReceiverAsyncClient( - fullyQualifiedNamespace, entityPath, entityType, newReceiverOptions, connectionProcessor, - ServiceBusConstants.OPERATION_TIMEOUT, tracerProvider, messageSerializer, () -> { }, - sessionSpecificManager)); + return tracer.traceMonoWithLink(sessionSpecificManager + .getActiveLink() + .map(receiveLink -> new ServiceBusReceiverAsyncClient( + fullyQualifiedNamespace, entityPath, entityType, newReceiverOptions, connectionProcessor, + ServiceBusConstants.OPERATION_TIMEOUT, tracer, messageSerializer, () -> { }, + sessionSpecificManager)), + null, "ServiceBus.acceptSession"); } @Override diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/MessageUtils.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/MessageUtils.java index 87e11b07d7cf..d7a53bcbd9c6 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/MessageUtils.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/MessageUtils.java @@ -6,11 +6,7 @@ import com.azure.core.amqp.AmqpRetryMode; import com.azure.core.amqp.AmqpRetryOptions; import com.azure.core.amqp.implementation.AmqpConstants; -import com.azure.core.amqp.implementation.TracerProvider; -import com.azure.core.util.Context; import com.azure.core.util.CoreUtils; -import com.azure.core.util.tracing.ProcessKind; -import com.azure.messaging.servicebus.ServiceBusMessage; import com.azure.messaging.servicebus.ServiceBusTransactionContext; import org.apache.qpid.proton.amqp.Binary; import org.apache.qpid.proton.amqp.DescribedType; @@ -23,7 +19,6 @@ import org.apache.qpid.proton.amqp.transaction.TransactionalState; import org.apache.qpid.proton.amqp.transport.DeliveryState; import org.apache.qpid.proton.amqp.transport.ErrorCondition; -import reactor.core.publisher.Signal; import java.net.URI; import java.nio.ByteBuffer; @@ -35,21 +30,11 @@ import java.util.HashMap; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.UUID; - -import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; -import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; -import static com.azure.core.util.tracing.Tracer.ENTITY_PATH_KEY; -import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; -import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; -import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_SERVICE_NAME; -import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_NAMESPACE_VALUE; import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.EPOCH_TICKS; import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.TICK_PER_SECOND; import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.TIME_LENGTH_DELTA; - /** * Contains helper methods for message conversions, reading status codes, and getting delivery state. */ @@ -309,35 +294,6 @@ private static TransactionalState getTransactionState(ByteBuffer transactionId, return transactionalState; } - /** - * Used in ServiceBusMessageBatch.tryAddMessage() to start tracing for to-be-sent out messages. - */ - public static ServiceBusMessage traceMessageSpan(ServiceBusMessage serviceBusMessage, - Context messageContext, String hostname, String entityPath, TracerProvider tracerProvider) { - Optional eventContextData = messageContext.getData(SPAN_CONTEXT_KEY); - if (eventContextData.isPresent()) { - // if message has context (in case of retries), don't start a message span or add a new context - return serviceBusMessage; - } else { - // Starting the span makes the sampling decision (nothing is logged at this time) - Context newMessageContext = messageContext - .addData(AZ_TRACING_NAMESPACE_KEY, AZ_TRACING_NAMESPACE_VALUE) - .addData(ENTITY_PATH_KEY, entityPath) - .addData(HOST_NAME_KEY, hostname); - Context eventSpanContext = tracerProvider.startSpan(AZ_TRACING_SERVICE_NAME, newMessageContext, - ProcessKind.MESSAGE); - Optional eventDiagnosticIdOptional = eventSpanContext.getData(DIAGNOSTIC_ID_KEY); - if (eventDiagnosticIdOptional.isPresent()) { - serviceBusMessage.getApplicationProperties().put(DIAGNOSTIC_ID_KEY, eventDiagnosticIdOptional.get() - .toString()); - tracerProvider.endSpan(eventSpanContext, Signal.complete()); - serviceBusMessage.addContext(SPAN_CONTEXT_KEY, eventSpanContext); - } - } - return serviceBusMessage; - } - - /** * Convert DescribedType to origin type based on the descriptor. * @param describedType Service bus defined DescribedType. diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusConstants.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusConstants.java index 7bba379ac6d1..f8daebef6131 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusConstants.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusConstants.java @@ -19,7 +19,6 @@ public class ServiceBusConstants { // Please see here // for more information on Azure resource provider namespaces. public static final String AZ_TRACING_NAMESPACE_VALUE = "Microsoft.ServiceBus"; - public static final String AZ_TRACING_SERVICE_NAME = "ServiceBus."; public static final String SERVICE_BUS_SUPPLEMENTARY_AUTHORIZATION_HEADER_NAME = "ServiceBusSupplementaryAuthorization"; public static final String SERVICE_BUS_DLQ_SUPPLEMENTARY_AUTHORIZATION_HEADER_NAME = "ServiceBusDlqSupplementaryAuthorization"; diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiveLinkProcessor.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiveLinkProcessor.java index 8b5f5fbd9ca8..b7d52bd79744 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiveLinkProcessor.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiveLinkProcessor.java @@ -48,6 +48,8 @@ public class ServiceBusReceiveLinkProcessor extends FluxProcessor implements Subscription { private static final ClientLogger LOGGER = new ClientLogger(ServiceBusReceiveLinkProcessor.class); + private static final String REACTOR_PARENT_TRACE_CONTEXT_KEY = "otel-context-key"; + protected static final String TRACEPARENT_KEY = "traceparent"; private final Object lock = new Object(); private final Object queueLock = new Object(); private final AtomicBoolean isTerminated = new AtomicBoolean(); @@ -527,7 +529,6 @@ private void drainQueue() { throw LOGGER.logExceptionAsError(Exceptions.propagate( Operators.onOperatorError(upstream, e, message, subscriber.currentContext()))); } - numberEmitted++; isEmpty = messageQueue.isEmpty(); } @@ -538,6 +539,7 @@ private void drainQueue() { } } + private boolean checkAndSetTerminated() { if (!isTerminated()) { return false; diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiverTracer.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiverTracer.java new file mode 100644 index 000000000000..2957f7092279 --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusReceiverTracer.java @@ -0,0 +1,163 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus.implementation; + +import com.azure.core.util.Context; +import com.azure.core.util.tracing.ProcessKind; +import com.azure.core.util.tracing.Tracer; +import com.azure.messaging.servicebus.ServiceBusReceivedMessage; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; + +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; +import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; +import static com.azure.core.util.tracing.Tracer.MESSAGE_ENQUEUED_TIME; +import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_NAMESPACE_VALUE; + +public class ServiceBusReceiverTracer extends ServiceBusTracer { + private static final AutoCloseable NOOP_AUTOCLOSEABLE = () -> { + }; + + private final boolean isSync; + + public ServiceBusReceiverTracer(Tracer tracer, String fullyQualifiedName, String entityPath, boolean isSync) { + super(tracer, fullyQualifiedName, entityPath); + this.isSync = isSync; + } + + public Context startProcessSpan(String name, ServiceBusReceivedMessage message, Context parent) { + if (message == null || tracer == null) { + return parent; + } + return tracer.start(name, setAttributes(message, parent), ProcessKind.PROCESS); + } + + public void reportReceiveSpan(String spanName, Instant startTime, Signal signal) { + if (tracer != null && (signal.isOnNext())) { + ServiceBusReceivedMessage message = null; + if (signal.hasValue()) { + message = signal.get(); + } + + Context span = startSpanWithLink(spanName, message, Context.NONE.addData("span-start-time", startTime)); + endSpan(signal.getThrowable(), span, null); + } + } + + public Flux reportSyncReceiverSpan(String name, Instant startTime, Flux messages) { + if (messages == null || tracer == null) { + return messages; + } + + List messageList = new ArrayList<>(); + return messages + .doOnEach(signal -> { + if (signal.isOnNext() && signal.hasValue()) { + messageList.add(signal.get()); + } else if (signal.isOnComplete() || signal.isOnError()) { + Context span = startSpanWithLinks(name, messageList, Context.NONE.addData("span-start-time", startTime)); + endSpan(signal.getThrowable(), span, null); + } + }); + } + + public Mono traceMonoWithLink(Mono publisher, ServiceBusReceivedMessage message, String spanName) { + if (tracer == null) { + return publisher; + } + + return publisher + .doOnEach(this::endSpan) + .contextWrite(ctx -> ctx.put(REACTOR_PARENT_TRACE_CONTEXT_KEY, startSpanWithLink(spanName, message, Context.NONE))); + } + + public Mono traceSettlement(Mono publisher, ServiceBusReceivedMessage message, DispositionStatus status) { + if (tracer == null) { + return publisher; + } + + return publisher + .doOnEach(this::endSpan) + .contextWrite(ctx -> ctx.put(REACTOR_PARENT_TRACE_CONTEXT_KEY, startSpanWithLink(getSettlementSpanName(status), message, Context.NONE))); + } + + public AutoCloseable makeSpanCurrent(Context span) { + return tracer == null ? NOOP_AUTOCLOSEABLE : tracer.makeSpanCurrent(span); + } + + private static String getSettlementSpanName(DispositionStatus status) { + switch (status) { + case COMPLETED: + return "ServiceBus.complete"; + case ABANDONED: + return "ServiceBus.abandon"; + case DEFERRED: + return "ServiceBus.defer"; + case SUSPENDED: + return "ServiceBus.deadLetter"; + case RELEASED: + return "ServiceBus.release"; + default: + return "ServiceBus.unknown"; + } + } + + public boolean isSync() { + return isSync; + } + + private Context startSpanWithLinks(String name, List batch, Context context) { + Context spanBuilder = getBuilder(name, context); + if (batch != null) { + for (ServiceBusReceivedMessage message : batch) { + addLink(message.getApplicationProperties(), message.getEnqueuedTime(), spanBuilder, Context.NONE); + } + } + + // TODO: need to refactor tracing in core. Currently we use ProcessKind.SEND as + // SpanKind.CLIENT + return tracer.start(name, spanBuilder, ProcessKind.SEND); + } + + private Context startSpanWithLink(String name, ServiceBusReceivedMessage message, Context parent) { + if (message == null) { + return parent; + } + Context spanBuilder = getBuilder(name, parent); + addLink(message.getApplicationProperties(), message.getEnqueuedTime(), spanBuilder, Context.NONE); + + // TODO: need to refactor tracing in core. Currently we use ProcessKind.SEND as + // SpanKind.CLIENT + return tracer.start(name, spanBuilder, ProcessKind.SEND); + } + + private Context getParent(Map properties, Context context) { + if (properties == null) { + return context; + } + + String traceparent = getTraceparent(properties); + return traceparent == null ? context : tracer.extractContext(traceparent, context); + } + + private Context setAttributes(ServiceBusReceivedMessage message, Context context) { + if (message.getEnqueuedTime() != null) { + context = context.addData(MESSAGE_ENQUEUED_TIME, message.getEnqueuedTime().toInstant().atOffset(ZoneOffset.UTC).toEpochSecond()); + } + + context = getParent(message.getApplicationProperties(), context); + + return context + .addData(Tracer.ENTITY_PATH_KEY, entityPath) + .addData(HOST_NAME_KEY, fullyQualifiedName) + .addData(AZ_TRACING_NAMESPACE_KEY, AZ_TRACING_NAMESPACE_VALUE); + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusTracer.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusTracer.java new file mode 100644 index 000000000000..a0a04a5cf435 --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusTracer.java @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus.implementation; + +import com.azure.core.amqp.exception.AmqpException; +import com.azure.core.util.Configuration; +import com.azure.core.util.Context; +import com.azure.core.util.logging.ClientLogger; +import com.azure.core.util.tracing.Tracer; +import reactor.core.publisher.Signal; + +import java.time.OffsetDateTime; +import java.util.Iterator; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.ServiceLoader; + +import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; +import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; +import static com.azure.core.util.tracing.Tracer.ENTITY_PATH_KEY; +import static com.azure.core.util.tracing.Tracer.HOST_NAME_KEY; +import static com.azure.core.util.tracing.Tracer.MESSAGE_ENQUEUED_TIME; +import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; +import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_NAMESPACE_VALUE; + +public class ServiceBusTracer { + private static final ClientLogger LOGGER = new ClientLogger(ServiceBusTracer.class); + protected static final String TRACEPARENT_KEY = "traceparent"; + protected static final String REACTOR_PARENT_TRACE_CONTEXT_KEY = "otel-context-key"; + + protected static final boolean IS_TRACING_DISABLED = Configuration.getGlobalConfiguration().get(Configuration.PROPERTY_AZURE_TRACING_DISABLED, false); + protected final Tracer tracer; + protected final String fullyQualifiedName; + protected final String entityPath; + + protected ServiceBusTracer(Tracer tracer, String fullyQualifiedName, String entityPath) { + this.tracer = IS_TRACING_DISABLED ? null : tracer; + this.fullyQualifiedName = Objects.requireNonNull(fullyQualifiedName, "'fullyQualifiedName' cannot be null"); + this.entityPath = Objects.requireNonNull(entityPath, "'entityPath' cannot be null"); + } + + public static Tracer getDefaultTracer() { + Iterable tracers = ServiceLoader.load(Tracer.class); + Iterator it = tracers.iterator(); + return it.hasNext() ? it.next() : null; + } + + public boolean isEnabled() { + return tracer != null; + } + + public void endSpan(Throwable throwable, Context span, AutoCloseable scope) { + if (tracer == null) { + return; + } + + String errorCondition = "success"; + if (throwable instanceof AmqpException) { + AmqpException exception = (AmqpException) throwable; + errorCondition = exception.getErrorCondition().getErrorCondition(); + } + + try { + if (scope != null) { + scope.close(); + } + } catch (Exception e) { + LOGGER.warning("Can't close scope", e); + } finally { + tracer.end(errorCondition, throwable, span); + } + } + + protected Context getBuilder(String spanName, Context context) { + return setAttributes(tracer.getSharedSpanBuilder(spanName, context)); + } + + protected Context setAttributes(Context context) { + return context + .addData(ENTITY_PATH_KEY, entityPath) + .addData(HOST_NAME_KEY, fullyQualifiedName) + .addData(AZ_TRACING_NAMESPACE_KEY, AZ_TRACING_NAMESPACE_VALUE); + } + + protected static String getTraceparent(Map applicationProperties) { + Object diagnosticId = applicationProperties.get(DIAGNOSTIC_ID_KEY); + if (diagnosticId == null) { + diagnosticId = applicationProperties.get(TRACEPARENT_KEY); + } + + return diagnosticId == null ? null : diagnosticId.toString(); + } + + protected void endSpan(Signal signal) { + Context span = signal.getContextView().getOrDefault(REACTOR_PARENT_TRACE_CONTEXT_KEY, Context.NONE); + endSpan(signal.getThrowable(), span, null); + } + + protected void addLink(Map applicationProperties, OffsetDateTime enqueuedTime, Context spanBuilder, Context eventContext) { + if (applicationProperties == null) { + return; + } + + Optional linkContext = eventContext.getData(SPAN_CONTEXT_KEY); + if (!linkContext.isPresent()) { + String traceparent = getTraceparent(applicationProperties); + Context link = traceparent == null ? Context.NONE : tracer.extractContext(traceparent, Context.NONE); + linkContext = link.getData(SPAN_CONTEXT_KEY); + } + + if (enqueuedTime != null) { + spanBuilder = spanBuilder.addData(MESSAGE_ENQUEUED_TIME, enqueuedTime.toInstant()); + } + + if (linkContext.isPresent()) { + tracer.addLink(spanBuilder.addData(SPAN_CONTEXT_KEY, linkContext.get())); + } + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/models/ServiceBusProcessorClientOptions.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/models/ServiceBusProcessorClientOptions.java index ce1b6094d02b..6c2c0ad35baf 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/models/ServiceBusProcessorClientOptions.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/models/ServiceBusProcessorClientOptions.java @@ -3,7 +3,6 @@ package com.azure.messaging.servicebus.implementation.models; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.annotation.Fluent; import com.azure.messaging.servicebus.ServiceBusProcessorClient; @@ -16,8 +15,6 @@ public final class ServiceBusProcessorClientOptions { private int maxConcurrentCalls = 1; private boolean disableAutoComplete; - private TracerProvider tracerProvider; - /** * Returns true if the auto-complete and auto-abandon feature is disabled. * @return true if the auto-complete and auto-abandon feature is disabled. @@ -53,24 +50,4 @@ public ServiceBusProcessorClientOptions setMaxConcurrentCalls(int maxConcurrentC this.maxConcurrentCalls = maxConcurrentCalls; return this; } - - /** - * Returns the {@link TracerProvider} instance that is used in {@link ServiceBusProcessorClient}. - * - * @return The {@link TracerProvider} instance that is used in {@link ServiceBusProcessorClient}. - */ - public TracerProvider getTracerProvider() { - return tracerProvider; - } - - /** - * Sets the {@link TracerProvider} instance to use in {@link ServiceBusProcessorClient}. - * - * @param tracerProvider The {@link TracerProvider} instance to use in {@link ServiceBusProcessorClient}. - * @return The updated instance of {@link ServiceBusProcessorClientOptions}. - */ - public ServiceBusProcessorClientOptions setTracerProvider(TracerProvider tracerProvider) { - this.tracerProvider = tracerProvider; - return this; - } } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusMessageBatchTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusMessageBatchTest.java index 19019439c777..40bab79c613f 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusMessageBatchTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusMessageBatchTest.java @@ -6,7 +6,6 @@ import com.azure.core.amqp.exception.AmqpErrorContext; import com.azure.core.amqp.implementation.ErrorContextProvider; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.util.BinaryData; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; @@ -14,8 +13,6 @@ import org.mockito.Mock; import org.mockito.MockitoAnnotations; -import java.util.Collections; - import static com.azure.messaging.servicebus.ServiceBusSenderAsyncClient.MAX_MESSAGE_LENGTH_BYTES; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.when; @@ -25,7 +22,7 @@ public class ServiceBusMessageBatchTest { private ErrorContextProvider errorContextProvider; private MessageSerializer serializer = new ServiceBusMessageSerializer(); - private TracerProvider tracerProvider = new TracerProvider(Collections.emptyList()); + private ServiceBusSenderTracer tracer = new ServiceBusSenderTracer(null, "namespace", "entity"); @BeforeEach public void setup() { @@ -34,7 +31,7 @@ public void setup() { @Test public void nullMessage() { - final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(1024, errorContextProvider, tracerProvider, serializer, null, null); + final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(1024, errorContextProvider, tracer, serializer, null, null); assertThrows(NullPointerException.class, () -> batch.tryAddMessage(null)); } @@ -46,7 +43,7 @@ public void payloadExceededException() { // Arrange when(errorContextProvider.getErrorContext()).thenReturn(new AmqpErrorContext("test-namespace")); - final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(1024, errorContextProvider, tracerProvider, serializer, null, null); + final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(1024, errorContextProvider, tracer, serializer, null, null); final ServiceBusMessage tooBig = new ServiceBusMessage(BinaryData.fromBytes(new byte[1024 * 1024 * 2])); // Act @@ -64,7 +61,7 @@ public void payloadExceededException() { @Test public void withinPayloadSize() { final int maxSize = MAX_MESSAGE_LENGTH_BYTES; - final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(maxSize, errorContextProvider, tracerProvider, + final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(maxSize, errorContextProvider, tracer, serializer, null, null); final ServiceBusMessage within = new ServiceBusMessage(BinaryData.fromBytes(new byte[1024])); @@ -81,7 +78,7 @@ public void withinPayloadSize() { public void setsPartitionId() { // Act final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(MAX_MESSAGE_LENGTH_BYTES, errorContextProvider, - tracerProvider, serializer, null, null); + tracer, serializer, null, null); // Assert Assertions.assertTrue(batch.getMessages().isEmpty()); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorTest.java index e95ef0e7f8a3..f1a75406cfc8 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorTest.java @@ -3,11 +3,11 @@ package com.azure.messaging.servicebus; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.util.BinaryData; import com.azure.core.util.Context; import com.azure.core.util.tracing.ProcessKind; import com.azure.core.util.tracing.Tracer; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.implementation.models.ServiceBusProcessorClientOptions; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -18,7 +18,6 @@ import java.time.OffsetDateTime; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -28,7 +27,7 @@ import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; import static com.azure.core.util.tracing.Tracer.MESSAGE_ENQUEUED_TIME; -import static com.azure.core.util.tracing.Tracer.PARENT_SPAN_KEY; +import static com.azure.core.util.tracing.Tracer.PARENT_TRACE_CONTEXT_KEY; import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; @@ -49,6 +48,9 @@ */ public class ServiceBusProcessorTest { + private static final String NAMESPACE = "namespace"; + private static final String ENTITY_NAME = "entity"; + /** * Tests receiving messages using a {@link ServiceBusProcessorClient}. * @@ -68,11 +70,11 @@ public void testReceivingMessagesWithProcessor() throws InterruptedException { } }); - ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux); + ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux, null); AtomicInteger messageId = new AtomicInteger(); CountDownLatch countDownLatch = new CountDownLatch(5); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { assertEquals(String.valueOf(messageId.getAndIncrement()), messageContext.getMessage().getMessageId()); @@ -112,7 +114,7 @@ public void testReceivingMultiSessionMessagesWithProcessor() throws InterruptedE AtomicInteger messageId = new AtomicInteger(); CountDownLatch countDownLatch = new CountDownLatch(numberOfMessages); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { int expectedMessageId = messageId.getAndIncrement(); @@ -139,14 +141,14 @@ public void testReceivingMultiSessionMessagesWithProcessor() throws InterruptedE public void testStartStopResume() throws InterruptedException { AtomicReference> sink = new AtomicReference<>(); Flux messageFlux = Flux.create(sink::set); - ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux); + ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux, null); AtomicInteger messageId = new AtomicInteger(); AtomicReference countDownLatch = new AtomicReference<>(); countDownLatch.set(new CountDownLatch(2)); AtomicBoolean assertionFailed = new AtomicBoolean(); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { try { @@ -223,13 +225,13 @@ public void testErrorRecovery() throws InterruptedException { return state + 1; }); - ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux); + ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux, null); AtomicInteger messageId = new AtomicInteger(); AtomicReference countDownLatch = new AtomicReference<>(); countDownLatch.set(new CountDownLatch(4)); AtomicBoolean assertionFailed = new AtomicBoolean(); StringBuffer messageIdNotMatched = new StringBuffer(); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { try { @@ -281,11 +283,13 @@ public void testUserMessageHandlerError() throws InterruptedException { when(asyncClient.receiveMessagesWithContext()).thenReturn(messageFlux); when(asyncClient.isConnectionClosed()).thenReturn(false); when(asyncClient.abandon(any(ServiceBusReceivedMessage.class))).thenReturn(Mono.empty()); + when(asyncClient.getFullyQualifiedNamespace()).thenReturn(NAMESPACE); + when(asyncClient.getEntityPath()).thenReturn(ENTITY_NAME); doNothing().when(asyncClient).close(); final AtomicInteger messageId = new AtomicInteger(); final CountDownLatch countDownLatch = new CountDownLatch(numberOfEvents); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { assertEquals(String.valueOf(messageId.getAndIncrement()), messageContext.getMessage().getMessageId()); @@ -333,11 +337,14 @@ public void testUserMessageHandlerErrorWithAutoCompleteDisabled() throws Interru when(receiverBuilder.buildAsyncClient()).thenReturn(asyncClient); when(asyncClient.receiveMessagesWithContext()).thenReturn(messageFlux); when(asyncClient.isConnectionClosed()).thenReturn(false); + when(asyncClient.getFullyQualifiedNamespace()).thenReturn(NAMESPACE); + when(asyncClient.getEntityPath()).thenReturn(ENTITY_NAME); + doNothing().when(asyncClient).close(); AtomicInteger messageId = new AtomicInteger(); CountDownLatch countDownLatch = new CountDownLatch(5); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { assertEquals(String.valueOf(messageId.getAndIncrement()), messageContext.getMessage().getMessageId()); @@ -361,9 +368,7 @@ public void testUserMessageHandlerErrorWithAutoCompleteDisabled() throws Interru @Test public void testProcessorWithTracingEnabled() throws InterruptedException { final Tracer tracer = mock(Tracer.class); - final List tracers = Collections.singletonList(tracer); final int numberOfTimes = 5; - final TracerProvider tracerProvider = new TracerProvider(tracers); String diagnosticId = "00-08ee063508037b1719dddcbf248e30e2-1365c684eb25daed-01"; @@ -377,9 +382,8 @@ public void testProcessorWithTracingEnabled() throws InterruptedException { invocation -> { Context passed = invocation.getArgument(1, Context.class); assertTrue(passed.getData(MESSAGE_ENQUEUED_TIME).isPresent()); - return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (AutoCloseable) () -> { - return; - }).addData(PARENT_SPAN_KEY, "value2"); + return passed.addData(SPAN_CONTEXT_KEY, "value1") + .addData(PARENT_TRACE_CONTEXT_KEY, "value2"); } ); Flux messageFlux = @@ -396,18 +400,18 @@ public void testProcessorWithTracingEnabled() throws InterruptedException { } }); - ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux); + ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux, tracer); AtomicInteger messageId = new AtomicInteger(); CountDownLatch countDownLatch = new CountDownLatch(numberOfTimes); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { assertEquals(String.valueOf(messageId.getAndIncrement()), messageContext.getMessage().getMessageId()); countDownLatch.countDown(); }, error -> Assertions.fail("Error occurred when receiving messages from the processor"), - new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1).setTracerProvider(tracerProvider)); + new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1)); serviceBusProcessorClient.start(); boolean success = countDownLatch.await(numberOfTimes, TimeUnit.SECONDS); @@ -426,17 +430,16 @@ public void testProcessorWithTracingEnabled() throws InterruptedException { @Test public void testProcessorWithTracingEnabledWithoutDiagnosticId() throws InterruptedException { final Tracer tracer = mock(Tracer.class); - final List tracers = Collections.singletonList(tracer); final int numberOfTimes = 5; - final TracerProvider tracerProvider = new TracerProvider(tracers); when(tracer.start(eq("ServiceBus.process"), any(), eq(ProcessKind.PROCESS))).thenAnswer( invocation -> { Context passed = invocation.getArgument(1, Context.class); assertTrue(passed.getData(MESSAGE_ENQUEUED_TIME).isPresent()); - return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (AutoCloseable) () -> { - return; - }).addData(PARENT_SPAN_KEY, "value2"); + return passed + .addData(SPAN_CONTEXT_KEY, "value1") + .addData("scope", (AutoCloseable) () -> { }) + .addData(PARENT_TRACE_CONTEXT_KEY, "value2"); } ); Flux messageFlux = @@ -452,18 +455,18 @@ public void testProcessorWithTracingEnabledWithoutDiagnosticId() throws Interrup } }); - ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux); + ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux, tracer); AtomicInteger messageId = new AtomicInteger(); CountDownLatch countDownLatch = new CountDownLatch(numberOfTimes); - ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, "queue", + ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME, null, null, messageContext -> { assertEquals(String.valueOf(messageId.getAndIncrement()), messageContext.getMessage().getMessageId()); countDownLatch.countDown(); }, error -> Assertions.fail("Error occurred when receiving messages from the processor"), - new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1).setTracerProvider(tracerProvider)); + new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1)); serviceBusProcessorClient.start(); boolean success = countDownLatch.await(numberOfTimes, TimeUnit.SECONDS); @@ -479,14 +482,19 @@ public void testProcessorWithTracingEnabledWithoutDiagnosticId() throws Interrup } private ServiceBusClientBuilder.ServiceBusReceiverClientBuilder getBuilder( - Flux messageFlux) { + Flux messageFlux, Tracer tracer) { ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = mock(ServiceBusClientBuilder.ServiceBusReceiverClientBuilder.class); ServiceBusReceiverAsyncClient asyncClient = mock(ServiceBusReceiverAsyncClient.class); when(receiverBuilder.buildAsyncClient()).thenReturn(asyncClient); - when(asyncClient.receiveMessagesWithContext()).thenReturn(messageFlux.publishOn(Schedulers.boundedElastic())); + when(asyncClient.getFullyQualifiedNamespace()).thenReturn(NAMESPACE); + when(asyncClient.getEntityPath()).thenReturn(ENTITY_NAME); + + ServiceBusReceiverTracer sbTracer = new ServiceBusReceiverTracer(tracer, NAMESPACE, ENTITY_NAME, false); + when(asyncClient.receiveMessagesWithContext()).thenReturn( + new FluxTrace(messageFlux, sbTracer).publishOn(Schedulers.boundedElastic())); when(asyncClient.isConnectionClosed()).thenReturn(false); doNothing().when(asyncClient).close(); return receiverBuilder; @@ -499,6 +507,8 @@ private ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder getSessio mock(ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder.class); ServiceBusReceiverAsyncClient asyncClient = mock(ServiceBusReceiverAsyncClient.class); + when(asyncClient.getFullyQualifiedNamespace()).thenReturn(NAMESPACE); + when(asyncClient.getEntityPath()).thenReturn(ENTITY_NAME); when(receiverBuilder.buildAsyncClientForProcessor()).thenReturn(asyncClient); when(asyncClient.receiveMessagesWithContext()).thenReturn(messageFlux); when(asyncClient.isConnectionClosed()).thenReturn(false); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClientTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClientTest.java index eab4c030b00f..4df3dfe043ae 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClientTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClientTest.java @@ -11,7 +11,6 @@ import com.azure.core.amqp.exception.AmqpException; import com.azure.core.amqp.implementation.ConnectionOptions; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.amqp.models.CbsAuthorizationType; import com.azure.core.credential.TokenCredential; import com.azure.core.exception.AzureException; @@ -26,6 +25,7 @@ import com.azure.messaging.servicebus.implementation.ServiceBusAmqpConnection; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; import com.azure.messaging.servicebus.implementation.ServiceBusConstants; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.implementation.ServiceBusManagementNode; import com.azure.messaging.servicebus.implementation.ServiceBusReactorReceiver; import com.azure.messaging.servicebus.models.AbandonOptions; @@ -132,8 +132,8 @@ class ServiceBusReceiverAsyncClientTest { private TokenCredential tokenCredential; @Mock private MessageSerializer messageSerializer; - @Mock - private TracerProvider tracerProvider; + + private ServiceBusReceiverTracer tracer = new ServiceBusReceiverTracer(null, NAMESPACE, ENTITY_PATH, false); @Mock private ServiceBusManagementNode managementNode; @Mock @@ -192,12 +192,12 @@ AmqpTransportType.AMQP, new AmqpRetryOptions(), ProxyOptions.SYSTEM_DEFAULTS, Sc receiver = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, MessagingEntityType.QUEUE, new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, false), - connectionProcessor, CLEANUP_INTERVAL, tracerProvider, messageSerializer, onClientClose, CLIENT_IDENTIFIER); + connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); sessionReceiver = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, MessagingEntityType.QUEUE, new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, false, SESSION_ID, null), - connectionProcessor, CLEANUP_INTERVAL, tracerProvider, messageSerializer, onClientClose, CLIENT_IDENTIFIER); + connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); } @AfterEach @@ -327,7 +327,7 @@ void receivesMessageLockRenewSessionOnly() { ServiceBusReceiverAsyncClient mySessionReceiver = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, MessagingEntityType.QUEUE, new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, maxLockRenewDuration, false, SESSION_ID, null), connectionProcessor, - CLEANUP_INTERVAL, tracerProvider, messageSerializer, onClientClose, CLIENT_IDENTIFIER); + CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); // This needs to be used with "try with resource" : https://javadoc.io/static/org.mockito/mockito-core/3.9.0/org/mockito/Mockito.html#static_mocks try ( @@ -419,7 +419,7 @@ void completeNullMessage() { void completeInReceiveAndDeleteMode() { final ReceiverOptions options = new ReceiverOptions(ServiceBusReceiveMode.RECEIVE_AND_DELETE, PREFETCH, null, false); ServiceBusReceiverAsyncClient client = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, options, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, options, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); final String lockToken1 = UUID.randomUUID().toString(); @@ -439,7 +439,7 @@ void completeInReceiveAndDeleteMode() { void throwsExceptionAboutSettlingPeekedMessagesWithNullLockToken() { final ReceiverOptions options = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, false); ServiceBusReceiverAsyncClient client = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, options, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, options, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); when(receivedMessage.getLockToken()).thenReturn(null); @@ -557,7 +557,7 @@ void errorSourceOnRenewMessageLock() { final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, true); final ServiceBusReceiverAsyncClient receiver2 = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); // Act & Assert @@ -582,7 +582,7 @@ void errorSourceOnSessionLock() { final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, true, SESSION_ID, null); final ServiceBusReceiverAsyncClient sessionReceiver2 = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, MessagingEntityType.QUEUE, - receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, messageSerializer, onClientClose, CLIENT_IDENTIFIER); + receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); // Act & Assert StepVerifier.create(sessionReceiver2.renewSessionLock(SESSION_ID)) @@ -654,7 +654,7 @@ void errorSourceAutoCompleteMessage() { final String lockToken = UUID.randomUUID().toString(); final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, true); final ServiceBusReceiverAsyncClient receiver2 = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); when(receivedMessage.getLockToken()).thenReturn(lockToken); @@ -702,7 +702,7 @@ void errorSourceOnReceiveMessage() { final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, true); final ServiceBusReceiverAsyncClient receiver2 = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); when(connection.createReceiveLink(anyString(), anyString(), any(ServiceBusReceiveMode.class), any(), @@ -1238,7 +1238,7 @@ void autoCompleteMessage() { final String lockToken = UUID.randomUUID().toString(); final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, true); final ServiceBusReceiverAsyncClient receiver2 = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); when(receivedMessage.getLockToken()).thenReturn(lockToken); @@ -1272,7 +1272,7 @@ void autoCompleteMessageSessionReceiver() { final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, PREFETCH, null, true, SESSION_ID, null); final ServiceBusReceiverAsyncClient sessionReceiver2 = new ServiceBusReceiverAsyncClient(NAMESPACE, ENTITY_PATH, - MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracerProvider, + MessagingEntityType.QUEUE, receiverOptions, connectionProcessor, CLEANUP_INTERVAL, tracer, messageSerializer, onClientClose, CLIENT_IDENTIFIER); final ServiceBusReceivedMessage receivedMessage3 = mock(ServiceBusReceivedMessage.class); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverClientTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverClientTest.java index bf5a1eb02b7c..c5826e614a6a 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverClientTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverClientTest.java @@ -5,6 +5,7 @@ import com.azure.core.util.IterableStream; import com.azure.core.util.logging.ClientLogger; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.models.AbandonOptions; import com.azure.messaging.servicebus.models.CompleteOptions; import com.azure.messaging.servicebus.models.DeadLetterOptions; @@ -84,6 +85,7 @@ void setup() { when(asyncClient.getReceiverOptions()).thenReturn(new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 0, null, false)); when(asyncClient.getIdentifier()).thenReturn(CLIENT_IDENTIFIER); when(sessionReceiverOptions.getSessionId()).thenReturn(SESSION_ID); + when(asyncClient.getTracer()).thenReturn(new ServiceBusReceiverTracer(null, NAMESPACE, ENTITY_PATH, false)); client = new ServiceBusReceiverClient(asyncClient, false, OPERATION_TIMEOUT); } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClientTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClientTest.java index 07b257809355..86faeaab5758 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClientTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSenderAsyncClientTest.java @@ -13,7 +13,6 @@ import com.azure.core.amqp.implementation.ConnectionOptions; import com.azure.core.amqp.implementation.ErrorContextProvider; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.amqp.models.CbsAuthorizationType; import com.azure.core.credential.TokenCredential; import com.azure.core.util.BinaryData; @@ -54,8 +53,6 @@ import java.time.Duration; import java.time.OffsetDateTime; import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.UUID; @@ -66,11 +63,11 @@ import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; import static com.azure.core.util.tracing.Tracer.DIAGNOSTIC_ID_KEY; -import static com.azure.core.util.tracing.Tracer.PARENT_SPAN_KEY; +import static com.azure.core.util.tracing.Tracer.PARENT_TRACE_CONTEXT_KEY; import static com.azure.core.util.tracing.Tracer.SPAN_BUILDER_KEY; +import static com.azure.core.util.tracing.Tracer.SPAN_CONTEXT_KEY; import static com.azure.messaging.servicebus.ServiceBusSenderAsyncClient.MAX_MESSAGE_LENGTH_BYTES; import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_NAMESPACE_VALUE; -import static com.azure.messaging.servicebus.implementation.ServiceBusConstants.AZ_TRACING_SERVICE_NAME; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; @@ -132,7 +129,7 @@ class ServiceBusSenderAsyncClientTest { private ArgumentCaptor> sequenceNumberCaptor; private final MessageSerializer serializer = new ServiceBusMessageSerializer(); - private final TracerProvider tracerProvider = new TracerProvider(Collections.emptyList()); + private static final ServiceBusSenderTracer DEFAULT_TRACER = new ServiceBusSenderTracer(null, NAMESPACE, ENTITY_NAME); private final AmqpRetryOptions retryOptions = new AmqpRetryOptions() .setDelay(Duration.ofMillis(500)) .setMode(AmqpRetryMode.FIXED) @@ -170,7 +167,7 @@ void setup() { connectionOptions.getRetry())); sender = new ServiceBusSenderAsyncClient(ENTITY_NAME, MessagingEntityType.QUEUE, connectionProcessor, - retryOptions, tracerProvider, serializer, onClientClose, null, CLIENT_IDENTIFIER); + retryOptions, DEFAULT_TRACER, serializer, onClientClose, null, CLIENT_IDENTIFIER); when(connection.getManagementNode(anyString(), any(MessagingEntityType.class))) .thenReturn(just(managementNode)); @@ -327,7 +324,7 @@ void sendMultipleMessagesWithTransaction() { final int count = 4; final byte[] contents = TEST_CONTENTS.toBytes(); final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(256 * 1024, - errorContextProvider, tracerProvider, serializer, null, null); + errorContextProvider, DEFAULT_TRACER, serializer, null, null); IntStream.range(0, count).forEach(index -> { final ServiceBusMessage message = new ServiceBusMessage(BinaryData.fromBytes(contents)); @@ -366,7 +363,7 @@ void sendMultipleMessages() { final int count = 4; final byte[] contents = TEST_CONTENTS.toBytes(); final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(256 * 1024, - errorContextProvider, tracerProvider, serializer, null, null); + errorContextProvider, DEFAULT_TRACER, serializer, null, null); IntStream.range(0, count).forEach(index -> { final ServiceBusMessage message = new ServiceBusMessage(BinaryData.fromBytes(contents)); @@ -398,38 +395,37 @@ void sendMultipleMessagesTracerSpans() { final int count = 4; final byte[] contents = TEST_CONTENTS.toBytes(); final Tracer tracer1 = mock(Tracer.class); - TracerProvider tracerProvider1 = new TracerProvider(Arrays.asList(tracer1)); + String traceparent = "traceparent"; + ServiceBusSenderTracer sbTracer1 = new ServiceBusSenderTracer(tracer1, NAMESPACE, ENTITY_NAME); final ServiceBusMessageBatch batch = new ServiceBusMessageBatch(256 * 1024, - errorContextProvider, tracerProvider1, serializer, null, null); + errorContextProvider, sbTracer1, serializer, null, null); sender = new ServiceBusSenderAsyncClient(ENTITY_NAME, MessagingEntityType.QUEUE, connectionProcessor, - retryOptions, tracerProvider1, serializer, onClientClose, null, CLIENT_IDENTIFIER); + retryOptions, sbTracer1, serializer, onClientClose, null, CLIENT_IDENTIFIER); when(connection.createSendLink(eq(ENTITY_NAME), eq(ENTITY_NAME), eq(retryOptions), isNull(), eq(CLIENT_IDENTIFIER))) .thenReturn(Mono.just(sendLink)); when(sendLink.send(anyList())).thenReturn(Mono.empty()); - when(tracer1.start(eq(AZ_TRACING_SERVICE_NAME + "send"), any(Context.class), eq(ProcessKind.SEND))) + when(tracer1.start(eq("ServiceBus.send"), any(Context.class), eq(ProcessKind.SEND))) .thenAnswer(invocation -> { Context passed = invocation.getArgument(1, Context.class); assertEquals(passed.getData(AZ_TRACING_NAMESPACE_KEY).get(), AZ_TRACING_NAMESPACE_VALUE); - return passed.addData(PARENT_SPAN_KEY, "value"); + return passed.addData(PARENT_TRACE_CONTEXT_KEY, "value"); }); - when(tracer1.start(eq(AZ_TRACING_SERVICE_NAME + "message"), any(Context.class), eq(ProcessKind.MESSAGE))) + when(tracer1.extractContext(eq(traceparent), any(Context.class))).thenAnswer(invocation -> { + Context passed = invocation.getArgument(1, Context.class); + return passed.addData(SPAN_CONTEXT_KEY, "span-context"); + }); + + when(tracer1.start(eq("ServiceBus.message"), any(Context.class), eq(ProcessKind.MESSAGE))) .thenAnswer(invocation -> { Context passed = invocation.getArgument(1, Context.class); assertEquals(passed.getData(AZ_TRACING_NAMESPACE_KEY).get(), AZ_TRACING_NAMESPACE_VALUE); - return passed.addData(PARENT_SPAN_KEY, "value").addData(DIAGNOSTIC_ID_KEY, "value2"); + return passed.addData(PARENT_TRACE_CONTEXT_KEY, "value").addData(DIAGNOSTIC_ID_KEY, traceparent); }); - when(tracer1.getSharedSpanBuilder(eq(AZ_TRACING_SERVICE_NAME + "send"), any(Context.class))).thenAnswer( - invocation -> { - Context passed = invocation.getArgument(1, Context.class); - return passed.addData(SPAN_BUILDER_KEY, "value"); - } - ); - - when(tracer1.getSharedSpanBuilder(eq(AZ_TRACING_SERVICE_NAME + "send"), any(Context.class))).thenAnswer( + when(tracer1.getSharedSpanBuilder(eq("ServiceBus.send"), any(Context.class))).thenAnswer( invocation -> { Context passed = invocation.getArgument(1, Context.class); return passed.addData(SPAN_BUILDER_KEY, "value"); @@ -447,9 +443,9 @@ void sendMultipleMessagesTracerSpans() { // Assert verify(tracer1, times(4)) - .start(eq(AZ_TRACING_SERVICE_NAME + "message"), any(Context.class), eq(ProcessKind.MESSAGE)); + .start(eq("ServiceBus.message"), any(Context.class), eq(ProcessKind.MESSAGE)); verify(tracer1, times(1)) - .start(eq(AZ_TRACING_SERVICE_NAME + "send"), any(Context.class), eq(ProcessKind.SEND)); + .start(eq("ServiceBus.send"), any(Context.class), eq(ProcessKind.SEND)); verify(tracer1, times(5)).end(eq("success"), isNull(), any(Context.class)); } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionManagerTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionManagerTest.java index 9ad4ddfc664d..8312ecfba366 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionManagerTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionManagerTest.java @@ -9,11 +9,11 @@ import com.azure.core.amqp.ProxyOptions; import com.azure.core.amqp.implementation.ConnectionOptions; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.amqp.models.CbsAuthorizationType; import com.azure.core.credential.TokenCredential; import com.azure.core.util.ClientOptions; import com.azure.core.util.logging.ClientLogger; +import com.azure.core.util.tracing.Tracer; import com.azure.messaging.servicebus.implementation.MessagingEntityType; import com.azure.messaging.servicebus.implementation.ServiceBusAmqpConnection; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; @@ -49,7 +49,6 @@ import java.time.Duration; import java.time.OffsetDateTime; -import java.util.Collections; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.TimeUnit; @@ -90,7 +89,7 @@ class ServiceBusSessionManagerTest { private final FluxSink endpointSink = endpointProcessor.sink(FluxSink.OverflowStrategy.BUFFER); private final EmitterProcessor messageProcessor = EmitterProcessor.create(); private final FluxSink messageSink = messageProcessor.sink(FluxSink.OverflowStrategy.BUFFER); - private final TracerProvider tracerProvider = new TracerProvider(Collections.emptyList()); + private final Tracer tracer = null; private ServiceBusConnectionProcessor connectionProcessor; private ServiceBusSessionManager sessionManager; @@ -179,7 +178,7 @@ void properties() { // Arrange ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, MAX_LOCK_RENEWAL, false, null, 5); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); // Act & Assert assertEquals(CLIENT_IDENTIFIER, sessionManager.getIdentifier()); @@ -190,7 +189,7 @@ void receiveNull() { // Arrange ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, MAX_LOCK_RENEWAL, false, null, 5); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); // Act & Assert StepVerifier.create(sessionManager.receive()) @@ -207,7 +206,7 @@ void singleUnnamedSession() { ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, MAX_LOCK_RENEWAL, false, null, 5); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); final String sessionId = "session-1"; final String lockToken = "a-lock-token"; @@ -262,7 +261,7 @@ void singleUnnamedSessionLockRenew() { ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, MAX_LOCK_RENEWAL, false, null, 1); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); final String sessionId = "session-1"; final String lockToken = "a-lock-token"; @@ -320,7 +319,7 @@ void multipleSessions() { final ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, MAX_LOCK_RENEWAL, true, null, 5); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); final int numberOfMessages = 5; final Callable onRenewal = () -> OffsetDateTime.now().plus(Duration.ofSeconds(5)); @@ -449,7 +448,7 @@ void multipleReceiveUnnamedSession() { null, 1); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); final String sessionId = "session-1"; final String linkName = "my-link-name"; @@ -521,7 +520,7 @@ void singleUnnamedSessionCleanupAfterTimeout() { ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, MAX_LOCK_RENEWAL, false, null, 2); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); final String sessionId = "session-1"; final String lockToken = "a-lock-token"; diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClientTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClientTest.java index 29887381876e..1cc3b3613460 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClientTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClientTest.java @@ -9,7 +9,6 @@ import com.azure.core.amqp.ProxyOptions; import com.azure.core.amqp.implementation.ConnectionOptions; import com.azure.core.amqp.implementation.MessageSerializer; -import com.azure.core.amqp.implementation.TracerProvider; import com.azure.core.amqp.models.CbsAuthorizationType; import com.azure.core.credential.TokenCredential; import com.azure.core.util.ClientOptions; @@ -18,6 +17,7 @@ import com.azure.messaging.servicebus.implementation.ServiceBusAmqpConnection; import com.azure.messaging.servicebus.implementation.ServiceBusConnectionProcessor; import com.azure.messaging.servicebus.implementation.ServiceBusConstants; +import com.azure.messaging.servicebus.implementation.ServiceBusReceiverTracer; import com.azure.messaging.servicebus.implementation.ServiceBusManagementNode; import com.azure.messaging.servicebus.implementation.ServiceBusReceiveLink; import com.azure.messaging.servicebus.models.ServiceBusReceiveMode; @@ -44,7 +44,6 @@ import java.time.Duration; import java.time.OffsetDateTime; -import java.util.Collections; import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicInteger; @@ -73,7 +72,7 @@ class ServiceBusSessionReceiverAsyncClientTest { private final FluxSink endpointSink = endpointProcessor.sink(FluxSink.OverflowStrategy.BUFFER); private final EmitterProcessor messageProcessor = EmitterProcessor.create(); private final FluxSink messageSink = messageProcessor.sink(FluxSink.OverflowStrategy.BUFFER); - private final TracerProvider tracerProvider = new TracerProvider(Collections.emptyList()); + private final ServiceBusReceiverTracer tracer = new ServiceBusReceiverTracer(null, NAMESPACE, ENTITY_PATH, false); private ServiceBusConnectionProcessor connectionProcessor; private ServiceBusSessionManager sessionManager; @@ -180,7 +179,7 @@ void acceptSession() { ServiceBusSessionReceiverAsyncClient client = new ServiceBusSessionReceiverAsyncClient( NAMESPACE, ENTITY_PATH, MessagingEntityType.QUEUE, receiverOptions, - connectionProcessor, tracerProvider, + connectionProcessor, tracer, messageSerializer, () -> { }, CLIENT_IDENTIFIER ); @@ -205,7 +204,7 @@ void acceptNextSession() { // Arrange ReceiverOptions receiverOptions = new ReceiverOptions(ServiceBusReceiveMode.PEEK_LOCK, 1, Duration.ZERO, false, null, null); sessionManager = new ServiceBusSessionManager(ENTITY_PATH, ENTITY_TYPE, connectionProcessor, - tracerProvider, messageSerializer, receiverOptions, CLIENT_IDENTIFIER); + messageSerializer, receiverOptions, CLIENT_IDENTIFIER); final int numberOfMessages = 5; final Callable onRenewal = () -> OffsetDateTime.now().plus(Duration.ofSeconds(5)); @@ -272,7 +271,7 @@ void acceptNextSession() { ServiceBusSessionReceiverAsyncClient client = new ServiceBusSessionReceiverAsyncClient( NAMESPACE, ENTITY_PATH, MessagingEntityType.QUEUE, receiverOptions, - connectionProcessor, tracerProvider, + connectionProcessor, tracer, messageSerializer, () -> { }, CLIENT_IDENTIFIER ); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/TracingIntegrationTests.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/TracingIntegrationTests.java new file mode 100644 index 000000000000..d3ee98378b5b --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/TracingIntegrationTests.java @@ -0,0 +1,501 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.logging.ClientLogger; +import com.azure.messaging.servicebus.models.DeferOptions; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.trace.IdGenerator; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.SpanProcessor; +import io.opentelemetry.sdk.trace.data.LinkData; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; +import org.junit.jupiter.api.parallel.Isolated; +import reactor.test.StepVerifier; + +import java.time.Duration; +import java.time.OffsetDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Predicate; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; + +@Isolated +@Execution(ExecutionMode.SAME_THREAD) +public class TracingIntegrationTests extends IntegrationTestBase { + private TestSpanProcessor spanProcessor; + private ServiceBusSenderAsyncClient sender; + ServiceBusReceiverAsyncClient receiver; + ServiceBusReceiverClient receiverSync; + ServiceBusProcessorClient processor; + + public TracingIntegrationTests() { + super(new ClientLogger(TracingIntegrationTests.class)); + } + + @Override + protected void beforeTest() { + spanProcessor = new TestSpanProcessor(getFullyQualifiedDomainName(), getQueueName(0)); + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .addSpanProcessor(spanProcessor) + .build()) + .buildAndRegisterGlobal(); + + sender = new ServiceBusClientBuilder() + .connectionString(getConnectionString()) + .sender() + .queueName(getQueueName(0)) + .buildAsyncClient(); + + receiver = new ServiceBusClientBuilder() + .connectionString(getConnectionString()) + .receiver() + .queueName(getQueueName(0)) + .buildAsyncClient(false, false); + + receiverSync = new ServiceBusClientBuilder() + .connectionString(getConnectionString()) + .receiver() + .queueName(getQueueName(0)) + .buildClient(); + } + + @Override + protected void afterTest() { + GlobalOpenTelemetry.resetForTest(); + sharedBuilder = null; + try { + dispose(receiver, sender, processor, receiverSync); + } catch (Exception e) { + logger.warning("Error occurred when draining queue.", e); + } + } + + @Test + public void sendAndReceive() { + ServiceBusMessage message1 = new ServiceBusMessage(CONTENTS_BYTES); + ServiceBusMessage message2 = new ServiceBusMessage(CONTENTS_BYTES); + List messages = Arrays.asList(message1, message2); + StepVerifier.create(sender.sendMessages(messages)) + .verifyComplete(); + + List received = new ArrayList<>(); + StepVerifier.create(receiver.receiveMessages() + .flatMap(rm -> receiver.complete(rm).thenReturn(rm)) + .take(2)) + .assertNext(rm -> received.add(rm)) + .assertNext(rm -> received.add(rm)) + .verifyComplete(); + + List spans = spanProcessor.getEndedSpans(); + + List messageSpans = findSpans(spans, "ServiceBus.message"); + assertMessageSpan(messageSpans.get(0), message1); + assertMessageSpan(messageSpans.get(1), message2); + + List send = findSpans(spans, "ServiceBus.send"); + assertSendSpan(send.get(0), messages, "ServiceBus.send"); + + List processed = findSpans(spans, "ServiceBus.process"); + assertConsumerSpan(processed.get(0), received.get(0), "ServiceBus.process"); + assertConsumerSpan(processed.get(1), received.get(1), "ServiceBus.process"); + + List completed = findSpans(spans, "ServiceBus.complete"); + assertReceiveSpan(completed.get(0), Collections.singletonList(received.get(0)), "ServiceBus.complete"); + assertParent(completed.get(0), processed.get(0)); + + assertReceiveSpan(completed.get(1), Collections.singletonList(received.get(1)), "ServiceBus.complete"); + assertParent(completed.get(1), processed.get(1)); + } + + @Test + public void sendPeekRenewLockAndDefer() throws InterruptedException { + String traceId = IdGenerator.random().generateTraceId(); + String traceparent = "00-" + traceId + "-" + IdGenerator.random().generateSpanId() + "-01"; + ServiceBusMessage message = new ServiceBusMessage(CONTENTS_BYTES); + AtomicReference receivedMessage = new AtomicReference<>(); + message.getApplicationProperties().put("traceparent", traceparent); + + StepVerifier.create(sender.sendMessage(message)).verifyComplete(); + + CountDownLatch latch = new CountDownLatch(2); + spanProcessor.notifyIfCondition(latch, s -> s.getName().equals("ServiceBus.process") && s.getSpanContext().getTraceId().equals(traceId)); + receiver.receiveMessages() + .skipUntil(m -> traceparent.equals(m.getApplicationProperties().get("traceparent"))) + .flatMap(m -> receiver.renewMessageLock(m, Duration.ofSeconds(1)).thenReturn(m)) + .flatMap(m -> receiver.defer(m, new DeferOptions()).thenReturn(m)) + .flatMap(m -> receiver.receiveDeferredMessage(m.getSequenceNumber()).thenReturn(m)) + .subscribe(m -> { + if (traceparent.equals(m.getApplicationProperties().get("traceparent"))) { + receivedMessage.set(m); + latch.countDown(); + } + }); + + assertTrue(latch.await(50, TimeUnit.SECONDS)); + + List spans = spanProcessor.getEndedSpans(); + + assertEquals(0, findSpans(spans, "ServiceBus.message").size()); + + List send = findSpans(spans, "ServiceBus.send"); + assertSendSpan(send.get(0), Collections.singletonList(message), "ServiceBus.send"); + + List process = findSpans(spans, "ServiceBus.process", traceId); + assertConsumerSpan(process.get(0), receivedMessage.get(), "ServiceBus.process"); + + List renewMessageLock = findSpans(spans, "ServiceBus.renewMessageLock", traceId); + assertReceiveSpan(renewMessageLock.get(0), Collections.singletonList(receivedMessage.get()), "ServiceBus.renewMessageLock"); + assertParent(renewMessageLock.get(0), process.get(0)); + + // for correlation to work after first async call, we need to enable otel rector instrumentations, + // so no correlation beyond this point + List defer = findSpans(spans, "ServiceBus.defer"); + assertReceiveSpan(defer.get(0), Collections.singletonList(receivedMessage.get()), "ServiceBus.defer"); + + List receiveDeferredMessage = findSpans(spans, "ServiceBus.receiveDeferredMessage"); + assertReceiveSpan(receiveDeferredMessage.get(0), Collections.singletonList(receivedMessage.get()), "ServiceBus.receiveDeferredMessage"); + } + + @Test + public void sendReceiveRenewLockAndDeferSync() { + StepVerifier.create(sender.sendMessage(new ServiceBusMessage(CONTENTS_BYTES))) + .verifyComplete(); + + ServiceBusReceivedMessage receivedMessage = receiverSync.receiveMessages(1, Duration.ofSeconds(1)).stream().findFirst().get(); + + receiverSync.renewMessageLock(receivedMessage); + receiverSync.defer(receivedMessage, new DeferOptions()); + receiverSync.receiveDeferredMessage(receivedMessage.getSequenceNumber()); + + List spans = spanProcessor.getEndedSpans(); + + List received = findSpans(spans, "ServiceBus.receiveMessages"); + assertReceiveSpan(received.get(0), Collections.singletonList(receivedMessage), "ServiceBus.receiveMessages"); + + List renewMessageLock = findSpans(spans, "ServiceBus.renewMessageLock"); + assertReceiveSpan(renewMessageLock.get(0), Collections.singletonList(receivedMessage), "ServiceBus.renewMessageLock"); + + List defer = findSpans(spans, "ServiceBus.defer"); + assertReceiveSpan(defer.get(0), Collections.singletonList(receivedMessage), "ServiceBus.defer"); + + List receiveDeferredMessage = findSpans(spans, "ServiceBus.receiveDeferredMessage"); + assertReceiveSpan(receiveDeferredMessage.get(0), Collections.singletonList(receivedMessage), "ServiceBus.receiveDeferredMessage"); + } + + @Test + public void syncReceive() { + List messages = new ArrayList<>(); + messages.add(new ServiceBusMessage(CONTENTS_BYTES)); + messages.add(new ServiceBusMessage(CONTENTS_BYTES)); + + StepVerifier.create(sender.sendMessages(messages)) + .verifyComplete(); + + List receivedMessages = receiverSync.receiveMessages(2, Duration.ofSeconds(10)) + .stream().collect(Collectors.toList()); + receivedMessages.forEach(receiverSync::complete); + + List spans = spanProcessor.getEndedSpans(); + + List received = findSpans(spans, "ServiceBus.receiveMessages"); + assertReceiveSpan(received.get(0), receivedMessages, "ServiceBus.receiveMessages"); + + assertEquals(0, findSpans(spans, "ServiceBus.process").size()); + + List completed = findSpans(spans, "ServiceBus.complete"); + assertReceiveSpan(completed.get(0), Collections.singletonList(receivedMessages.get(0)), "ServiceBus.complete"); + assertReceiveSpan(completed.get(1), Collections.singletonList(receivedMessages.get(1)), "ServiceBus.complete"); + } + + @Test + public void syncReceiveTimeout() { + List receivedMessages = receiverSync.receiveMessages(100, Duration.ofMillis(1)) + .stream().collect(Collectors.toList()); + + List spans = spanProcessor.getEndedSpans(); + + List received = findSpans(spans, "ServiceBus.receiveMessages"); + assertReceiveSpan(received.get(0), receivedMessages, "ServiceBus.receiveMessages"); + assertEquals(StatusCode.OK, received.get(0).toSpanData().getStatus().getStatusCode()); + + assertEquals(0, findSpans(spans, "ServiceBus.process").size()); + } + + @Test + public void peekMessage() { + StepVerifier.create(sender.sendMessage(new ServiceBusMessage(CONTENTS_BYTES))) + .verifyComplete(); + + StepVerifier.create(receiver.peekMessage()) + .assertNext(receivedMessage -> { + List received = findSpans(spanProcessor.getEndedSpans(), "ServiceBus.peekMessage"); + assertReceiveSpan(received.get(0), Collections.singletonList(receivedMessage), "ServiceBus.peekMessage"); + }) + .verifyComplete(); + } + + @Test + public void sendAndProcess() throws InterruptedException { + String messageId = UUID.randomUUID().toString(); + ServiceBusMessage message = new ServiceBusMessage(CONTENTS_BYTES) + .setMessageId(messageId); + + StepVerifier.create(sender.sendMessage(message)) + .verifyComplete(); + + String message1SpanId = message.getApplicationProperties().get("traceparent").toString().substring(36, 52); + CountDownLatch completedFound = new CountDownLatch(1); + spanProcessor.notifyIfCondition(completedFound, span -> { + if (span.getName() != "ServiceBus.complete") { + return false; + } + List links = span.toSpanData().getLinks(); + return links.size() > 0 && links.get(0).getSpanContext().getSpanId().equals(message1SpanId); + }); + + AtomicReference currentInProcess = new AtomicReference<>(Span.getInvalid()); + AtomicReference receivedMessage = new AtomicReference<>(); + processor = new ServiceBusClientBuilder() + .connectionString(getConnectionString()) + .processor() + .queueName(getQueueName(0)) + .processMessage(mc -> { + if (mc.getMessage().getMessageId().equals(messageId)) { + currentInProcess.set(Span.current()); + receivedMessage.set(mc.getMessage()); + } + }) + .processError(e -> { + fail("unexpected error", e.getException()); + }) + .buildProcessorClient(); + + processor.start(); + assertTrue(completedFound.await(10, TimeUnit.SECONDS)); + processor.stop(); + + assertTrue(currentInProcess.get().getSpanContext().isValid()); + List spans = spanProcessor.getEndedSpans(); + + assertMessageSpan(spans.get(0), message); + assertSendSpan(spans.get(1), Collections.singletonList(message), "ServiceBus.send"); + + assertEquals(0, findSpans(spans, "ServiceBus.consume").size()); + List processed = findSpans(spans, "ServiceBus.process") + .stream().filter(p -> p == currentInProcess.get()).collect(Collectors.toList()); + assertEquals(1, processed.size()); + assertConsumerSpan(processed.get(0), receivedMessage.get(), "ServiceBus.process"); + + List completed = findSpans(spans, "ServiceBus.complete").stream() + .filter(c -> c.toSpanData().getLinks().get(0).getSpanContext().getSpanId().equals(message1SpanId)) + .collect(Collectors.toList()); + assertEquals(1, completed.size()); + assertSendSpan(completed.get(0), Collections.singletonList(message), "ServiceBus.complete"); + assertParent(completed.get(0), processed.get(0)); + } + + @Test + public void sendProcessAndFail() throws InterruptedException { + String messageId = UUID.randomUUID().toString(); + ServiceBusMessage message = new ServiceBusMessage(CONTENTS_BYTES) + .setMessageId(messageId); + + StepVerifier.create(sender.sendMessage(message)) + .verifyComplete(); + + String message1SpanId = message.getApplicationProperties().get("traceparent").toString().substring(36, 52); + CountDownLatch messageProcessed = new CountDownLatch(1); + spanProcessor.notifyIfCondition(messageProcessed, span -> + span.getName() == "ServiceBus.process" && span.getParentSpanContext().getSpanId().equals(message1SpanId)); + + AtomicReference receivedMessage = new AtomicReference<>(); + processor = new ServiceBusClientBuilder() + .connectionString(getConnectionString()) + .processor() + .queueName(getQueueName(0)) + .processMessage(mc -> { + if (mc.getMessage().getMessageId().equals(messageId)) { + receivedMessage.set(mc.getMessage()); + throw new RuntimeException("foo"); + } + }) + .processError(e -> { }) + .buildProcessorClient(); + + processor.start(); + assertTrue(messageProcessed.await(10, TimeUnit.SECONDS)); + processor.stop(); + + List spans = spanProcessor.getEndedSpans(); + //assertEquals(0, findSpans(spans, "ServiceBus.consume").size()); + List processed = findSpans(spans, "ServiceBus.process") + .stream().filter(p -> p.getParentSpanContext().isValid()) + .filter(p -> p.toSpanData().getStatus().getStatusCode() == StatusCode.ERROR) + .collect(Collectors.toList()); + assertEquals(1, processed.size()); + assertConsumerSpan(processed.get(0), receivedMessage.get(), "ServiceBus.process"); + + List abandoned = findSpans(spans, "ServiceBus.abandon").stream() + .filter(c -> c.toSpanData().getLinks().get(0).getSpanContext().getSpanId().equals(message1SpanId)) + .collect(Collectors.toList()); + assertEquals(1, abandoned.size()); + assertSendSpan(abandoned.get(0), Collections.singletonList(message), "ServiceBus.abandon"); + assertParent(abandoned.get(0), processed.get(0)); + } + + @Test + public void scheduleAndCancelMessage() { + ServiceBusMessage message = new ServiceBusMessage("m"); + StepVerifier.create( + sender.scheduleMessage(message, OffsetDateTime.now().plusSeconds(100)) + .flatMap(l -> sender.cancelScheduledMessage(l))) + .verifyComplete(); + + List spans = spanProcessor.getEndedSpans(); + assertMessageSpan(spans.get(0), message); + assertSendSpan(spans.get(1), Collections.singletonList(message), "ServiceBus.scheduleMessage"); + assertSendSpan(spans.get(2), Collections.emptyList(), "ServiceBus.cancelScheduledMessage"); + } + + private void assertMessageSpan(ReadableSpan actual, ServiceBusMessage message) { + assertEquals("ServiceBus.message", actual.getName()); + assertEquals(SpanKind.PRODUCER, actual.getKind()); + String traceparent = "00-" + actual.getSpanContext().getTraceId() + "-" + actual.getSpanContext().getSpanId() + "-01"; + assertEquals(message.getApplicationProperties().get("Diagnostic-Id"), traceparent); + assertEquals(message.getApplicationProperties().get("traceparent"), traceparent); + } + + private void assertSendSpan(ReadableSpan actual, List messages, String spanName) { + assertEquals(spanName, actual.getName()); + assertEquals(SpanKind.CLIENT, actual.getKind()); + List links = actual.toSpanData().getLinks(); + assertEquals(messages.size(), links.size()); + for (int i = 0; i < links.size(); i++) { + String messageTraceparent = (String) messages.get(i).getApplicationProperties().get("traceparent"); + SpanContext linkContext = links.get(i).getSpanContext(); + String linkTraceparent = "00-" + linkContext.getTraceId() + "-" + linkContext.getSpanId() + "-01"; + assertEquals(messageTraceparent, linkTraceparent); + } + } + + private void assertReceiveSpan(ReadableSpan actual, List messages, String spanName) { + assertEquals(spanName, actual.getName()); + assertEquals(SpanKind.CLIENT, actual.getKind()); + List links = actual.toSpanData().getLinks(); + assertEquals(messages.size(), links.size()); + for (int i = 0; i < links.size(); i++) { + String messageTraceparent = (String) messages.get(i).getApplicationProperties().get("traceparent"); + SpanContext linkContext = links.get(i).getSpanContext(); + String linkTraceparent = "00-" + linkContext.getTraceId() + "-" + linkContext.getSpanId() + "-01"; + assertEquals(messageTraceparent, linkTraceparent); + // TODO (lmolkova) uncomment after otel 1.0.0-beta.29 ships + // assertNotNull(links.get(i).getAttributes().get(AttributeKey.longKey(Tracer.MESSAGE_ENQUEUED_TIME))); + } + } + + private void assertConsumerSpan(ReadableSpan actual, ServiceBusReceivedMessage message, String spanName) { + assertEquals(spanName, actual.getName()); + assertEquals(SpanKind.CONSUMER, actual.getKind()); + assertEquals(0, actual.toSpanData().getLinks().size()); + + String messageTraceparent = (String) message.getApplicationProperties().get("traceparent"); + String parent = "00-" + actual.getSpanContext().getTraceId() + "-" + actual.getParentSpanContext().getSpanId() + "-01"; + assertEquals(messageTraceparent, parent); + } + + private void assertParent(ReadableSpan child, ReadableSpan parent) { + assertEquals(child.getParentSpanContext().getTraceId(), parent.getSpanContext().getTraceId()); + assertEquals(child.getParentSpanContext().getSpanId(), parent.getSpanContext().getSpanId()); + } + + private List findSpans(List spans, String spanName) { + return spans.stream() + .filter(s -> s.getName().equals(spanName)) + .collect(Collectors.toList()); + } + + private List findSpans(List spans, String spanName, String traceId) { + return spans.stream() + .filter(s -> s.getName().equals(spanName)) + .filter(s -> s.getSpanContext().getTraceId().equals(traceId)) + .collect(Collectors.toList()); + } + + static class TestSpanProcessor implements SpanProcessor { + private final ConcurrentLinkedDeque spans = new ConcurrentLinkedDeque<>(); + private final String entityName; + private final String namespace; + + private AtomicReference> notifier = new AtomicReference<>(); + + TestSpanProcessor(String namespace, String entityName) { + this.namespace = namespace; + this.entityName = entityName; + } + public List getEndedSpans() { + return spans.stream().collect(Collectors.toList()); + } + + @Override + public void onStart(Context context, ReadWriteSpan readWriteSpan) { + } + + @Override + public boolean isStartRequired() { + return false; + } + + @Override + public void onEnd(ReadableSpan readableSpan) { + assertEquals("Microsoft.ServiceBus", readableSpan.getAttribute(AttributeKey.stringKey("az.namespace"))); + assertEquals(entityName, readableSpan.getAttribute(AttributeKey.stringKey("message_bus.destination"))); + assertEquals(namespace, readableSpan.getAttribute(AttributeKey.stringKey("peer.address"))); + + Consumer filter = notifier.get(); + if (filter != null) { + filter.accept(readableSpan); + } + spans.add(readableSpan); + } + + public void notifyIfCondition(CountDownLatch countDownLatch, Predicate filter) { + notifier.set((span) -> { + if (filter.test(span)) { + countDownLatch.countDown(); + } + }); + } + + @Override + public boolean isEndRequired() { + return true; + } + } +}