Tracing for eventhubs consumer and batch processor (Azure#31197)

* Trace eventhubs consumer and batch processor

Author: Liudmila Molkova

eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml

@@ -100,6 +100,7 @@
   <suppress checks="IllegalImport" files=".*[/\\]com[/\\]azure[/\\]identity[/\\]*"/>
   <suppress checks="IllegalImport" files="com.azure.messaging.servicebus.TracingIntegrationTests.java"/>
   <suppress checks="IllegalImport" files="com.azure.messaging.eventhubs.TracingIntegrationTests.java"/>
+  <suppress checks="IllegalImport" files="com.azure.messaging.eventhubs.PublishEventsTracingWithCustomContextSample.java"/>
 
   <!-- Suppress warnings for Event Processor until the usage of "Client" is discussed and resolved:
     https://github.com/Azure/azure-sdk/issues/321 -->

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/AmqpMetricsProvider.java

@@ -146,6 +146,14 @@ public boolean isSendDeliveryEnabled() {
         return isEnabled && sendDuration.isEnabled();
     }
 
+    /**
+     * Checks if request-response duration metric is enabled (for micro-optimizations).
+     */
+    public boolean isRequestResponseDurationEnabled() {
+        return isEnabled && sendDuration.isEnabled();
+    }
+
+
     /**
      * Checks if prefetched sequence number is enabled (for micro-optimizations).
      */

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/RequestResponseChannel.java

@@ -504,7 +504,7 @@ private void terminateUnconfirmedSends(Throwable error) {
      * Captures current time in mono context - used to report send metric
      */
     private Mono<Message> captureStartTime(Message toSend, Mono<Message> publisher) {
-        if (metricsProvider.isSendDeliveryEnabled()) {
+        if (metricsProvider.isRequestResponseDurationEnabled()) {
             String operationName = "unknown";
             if (toSend != null && toSend.getApplicationProperties() != null && toSend.getApplicationProperties().getValue() != null) {
                 Map<String, Object> properties = toSend.getApplicationProperties().getValue();
@@ -532,7 +532,7 @@ private static ContextView getSinkContext(MonoSink<?> sink) {
      * Records send call duration metric.
      **/
     private void recordDelivery(ContextView context, Message response) {
-        if (metricsProvider.isSendDeliveryEnabled()) {
+        if (metricsProvider.isRequestResponseDurationEnabled()) {
             Object startTimestamp = context.getOrDefault(START_SEND_TIME_CONTEXT_KEY, null);
             Object operationName = context.getOrDefault(OPERATION_CONTEXT_KEY, null);
             AmqpResponseCode responseCode = response == null ? null : RequestResponseUtils.getStatusCode(response);

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/TracerProvider.java

@@ -11,6 +11,10 @@
 
 import java.util.Objects;
 
+@Deprecated
+/**
+ * @deprecated use EventHubs*Tracer and ServiceBus*Tracer instead.
+ */
 public class TracerProvider {
     private static final ClientLogger LOGGER = new ClientLogger(TracerProvider.class);
     private Tracer tracer;

sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/TracerProviderTest.java

@@ -31,6 +31,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+@SuppressWarnings("deprecation")
 public class TracerProviderTest {
     private static final String SERVICE_BASE_NAME = "serviceBaseName";
     private static final String METHOD_NAME = SERVICE_BASE_NAME + "send";

sdk/core/azure-core-tracing-opentelemetry/README.md

@@ -122,9 +122,13 @@ try {
 
 Send a single event/message using [azure-messaging-eventhubs][azure-messaging-eventhubs] with tracing enabled.
 
-Users can additionally pass the value of the current tracing span to the EventData object with key **PARENT_TRACE_CONTEXT_KEY** on the [Context][context] object:
+Users can additionally pass custom value of the trace context to the EventData object with key **PARENT_TRACE_CONTEXT_KEY** on the [Context][context] object.
+
+Please refer to [Event Hubs samples][event_hubs_samples]
+for more information.
+
+```java
 
-```java readme-sample-context-manual-propagation-amqp
 Flux<EventData> events = Flux.just(
     new EventData("EventData Sample 1"),
     new EventData("EventData Sample 2"));
@@ -151,6 +155,7 @@ events.collect(batchRef::get, (b, e) ->
         return ctx.put(PARENT_TRACE_CONTEXT_KEY, traceContextRef.updateAndGet(traceContext -> traceContext.with(span)));
     })
     .block();
+
 ```
 
 ## Troubleshooting
@@ -203,9 +208,9 @@ This project has adopted the [Microsoft Open Source Code of Conduct](https://ope
 [OpenTelemetry]: https://github.com/open-telemetry/opentelemetry-java#opentelemetry-for-java
 [sample_app_config]: https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/core/azure-core-tracing-opentelemetry/src/samples/java/com/azure/core/tracing/opentelemetry/CreateConfigurationSettingLoggingExporterSample.java
 [sample_async_key_vault]: https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/core/azure-core-tracing-opentelemetry/src/samples/java/com/azure/core/tracing/opentelemetry/ListKeyVaultSecretsAutoConfigurationSample.java
-[sample_eventhubs]: https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/core/azure-core-tracing-opentelemetry/src/samples/java/com/azure/core/tracing/opentelemetry/PublishEventsJaegerExporterSample.java
 [sample_key_vault]: https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/core/azure-core-tracing-opentelemetry/src/samples/java/com/azure/core/tracing/opentelemetry/ListKeyVaultSecretsJaegerExporterSample.java
 [samples]: https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/core/azure-core-tracing-opentelemetry/src/samples/
 [source_code]: https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/core/azure-core-tracing-opentelemetry/src
+[event_hubs_samples](https://github.com/Azure/azure-sdk-for-java/blob/10a18ccc2f20cad6004ae90d64f22009d65e9ef7/sdk/eventhubs/azure-messaging-eventhubs/src/samples/java/com/azure/messaging/eventhubs/PublishEventsTracingWithCustomContextSample.java)
 
 ![Impressions](https://azure-sdk-impressions.azurewebsites.net/api/impressions/azure-sdk-for-java%2Fsdk%2Fcore%2Fazure-core-tracing-opentelemetry%2FREADME.png)

sdk/core/azure-core-tracing-opentelemetry/pom.xml

@@ -134,12 +134,6 @@
       <version>1.6.1</version> <!-- {x-version-update;com.azure:azure-identity;dependency} -->
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>com.azure</groupId>
-      <artifactId>azure-messaging-eventhubs</artifactId>
-      <version>5.14.0</version> <!-- {x-version-update;com.azure:azure-messaging-eventhubs;dependency} -->
-      <scope>test</scope>
-    </dependency>
     <dependency>
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-sdk-extension-autoconfigure</artifactId>

sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java

@@ -22,6 +22,7 @@
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Scope;
 
+import java.time.Instant;
 import java.time.OffsetDateTime;
 import java.util.Map;
 import java.util.Objects;
@@ -36,6 +37,8 @@
  */
 public class OpenTelemetryTracer implements com.azure.core.util.tracing.Tracer {
     private static final StartSpanOptions DEFAULT_OPTIONS = new StartSpanOptions(com.azure.core.util.tracing.SpanKind.INTERNAL);
+    private static final String SPAN_KIND_KEY = "span-kind";
+    private static final String START_TIME_KEY = "span-start-time";
     private final Tracer tracer;
 
     /**
@@ -65,7 +68,7 @@ public OpenTelemetryTracer() {
 
     private static final ClientLogger LOGGER = new ClientLogger(OpenTelemetryTracer.class);
     private static final AutoCloseable NOOP_CLOSEABLE = () -> { };
-    private static final SpanKind SHARED_SPAN_BUILDER_KIND = SpanKind.CLIENT;
+    private static final SpanKind DEFAULT_SHARED_SPAN_BUILDER_KIND = SpanKind.CLIENT;
     private static final String SUPPRESSED_SPAN_FLAG = "suppressed-span-flag";
     private static final String CLIENT_METHOD_CALL_FLAG = "client-method-call-flag";
 
@@ -114,7 +117,7 @@ public Context start(String spanName, Context context, ProcessKind processKind)
         context = unsuppress(context);
         switch (processKind) {
             case SEND:
-                // use previously created span builder from the LINK process.
+                // use previously created span builder with the links
                 spanBuilder = getOrNull(context, SPAN_BUILDER_KEY, SpanBuilder.class);
                 if (spanBuilder == null) {
                     // we can't return context here, because caller would not know that span was not created.
@@ -123,17 +126,23 @@ public Context start(String spanName, Context context, ProcessKind processKind)
                         .addKeyValue("spanName", spanName)
                         .addKeyValue("processKind", processKind)
                         .log("Start span is called without builder on the context, creating default builder.");
-                    spanBuilder = createSpanBuilder(spanName, null, SHARED_SPAN_BUILDER_KIND, null, context);
+                    spanBuilder = createSpanBuilder(spanName, null, SpanKind.CLIENT, null, context);
                 }
 
-                return startSpanInternal(spanBuilder, isClientCall(SHARED_SPAN_BUILDER_KIND), this::addMessagingAttributes, context);
+                return startSpanInternal(spanBuilder, true, this::addMessagingAttributes, context);
             case MESSAGE:
                 spanBuilder = createSpanBuilder(spanName, null, SpanKind.PRODUCER, null, context);
                 context = startSpanInternal(spanBuilder, false, this::addMessagingAttributes, context);
                 return setDiagnosticId(context);
             case PROCESS:
-                SpanContext remoteParentContext = getOrNull(context, SPAN_CONTEXT_KEY, SpanContext.class);
-                spanBuilder = createSpanBuilder(spanName, remoteParentContext, SpanKind.CONSUMER, null, context);
+                // use previously created span builder with the links
+                spanBuilder = getOrNull(context, SPAN_BUILDER_KEY, SpanBuilder.class);
+                if (spanBuilder == null) {
+                    // if there is no builder, create new one from parent in context
+                    SpanContext remoteParentContext = getOrNull(context, SPAN_CONTEXT_KEY, SpanContext.class);
+                    spanBuilder = createSpanBuilder(spanName, remoteParentContext, SpanKind.CONSUMER, null, context);
+                }
+
                 context = startSpanInternal(spanBuilder, false, this::addMessagingAttributes, context);
 
                 // TODO (limolkova) we should do this in the EventHub/ServiceBus SDK instead to make sure scope is
@@ -227,7 +236,14 @@ public void addLink(Context context) {
         if (spanContext == null) {
             return;
         }
-        spanBuilder.addLink(spanContext);
+
+        Attributes linkAttributes = Attributes.empty();
+        Long messageEnqueuedTime = getOrNull(context, MESSAGE_ENQUEUED_TIME, Long.class);
+        if (messageEnqueuedTime != null) {
+            linkAttributes = Attributes.of(AttributeKey.longKey(MESSAGE_ENQUEUED_TIME), messageEnqueuedTime);
+        }
+
+        spanBuilder.addLink(spanContext, linkAttributes);
     }
 
     /**
@@ -243,8 +259,18 @@ public Context extractContext(String diagnosticId, Context context) {
      */
     @Override
     public Context getSharedSpanBuilder(String spanName, Context context) {
-        // this is used to create messaging send spanBuilder, and it's a CLIENT span
-        return context.addData(SPAN_BUILDER_KEY, createSpanBuilder(spanName, null, SHARED_SPAN_BUILDER_KIND, null, context));
+        com.azure.core.util.tracing.SpanKind spanKind = getOrNull(context, SPAN_KIND_KEY, com.azure.core.util.tracing.SpanKind.class);
+        if (spanKind == null) {
+            spanKind = com.azure.core.util.tracing.SpanKind.CLIENT;
+        }
+
+        SpanBuilder builder = createSpanBuilder(spanName, null, convertToOtelKind(spanKind), null, context);
+        Instant startTime = getOrNull(context, START_TIME_KEY, Instant.class);
+        if (startTime != null) {
+            builder.setStartTimestamp(startTime);
+        }
+
+        return context.addData(SPAN_BUILDER_KEY, builder);
     }
 
     /**
@@ -563,7 +589,7 @@ private Span getSpanOrNull(Context azContext) {
     private SpanKind processKindToSpanKind(ProcessKind processKind) {
         switch (processKind) {
             case SEND:
-                return SHARED_SPAN_BUILDER_KIND;
+                return SpanKind.CLIENT;
             case MESSAGE:
                 return SpanKind.PRODUCER;
             case PROCESS:

sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracerTest.java

@@ -323,6 +323,58 @@ public void startProcessSpanWithRemoteParent() {
         assertTrue(updatedContext.getData(SCOPE_KEY).isPresent());
     }
 
+    @Test
+    public void startProcessSpanWithLinks() {
+        // Arrange
+        final Context spanBuilder = openTelemetryTracer.getSharedSpanBuilder("span", Context.NONE);
+
+        Span link1 = tracer.spanBuilder("link1").startSpan();
+        Span link2 = tracer.spanBuilder("link2").startSpan();
+
+        openTelemetryTracer.addLink(spanBuilder.addData(SPAN_CONTEXT_KEY, link1.getSpanContext()));
+        openTelemetryTracer.addLink(spanBuilder
+            .addData(SPAN_CONTEXT_KEY, link2.getSpanContext())
+            .addData(MESSAGE_ENQUEUED_TIME, MESSAGE_ENQUEUED_VALUE));
+
+        // Act
+        final Context spanCtx = openTelemetryTracer.start(METHOD_NAME, spanBuilder, ProcessKind.PROCESS);
+        openTelemetryTracer.end(null, null, spanCtx);
+
+        // Assert
+        ReadableSpan span = getSpan(spanCtx);
+        List<LinkData> links = span.toSpanData().getLinks();
+        assertEquals(2, links.size());
+        assertEquals(link1.getSpanContext().getTraceId(), links.get(0).getSpanContext().getTraceId());
+        assertEquals(link1.getSpanContext().getSpanId(),  links.get(0).getSpanContext().getSpanId());
+        assertEquals(0, links.get(0).getAttributes().size());
+
+        assertEquals(link2.getSpanContext().getTraceId(), links.get(1).getSpanContext().getTraceId());
+        assertEquals(link2.getSpanContext().getSpanId(),  links.get(1).getSpanContext().getSpanId());
+        Attributes linkAttributes = links.get(1).getAttributes();
+        assertEquals(1, linkAttributes.size());
+        assertEquals(MESSAGE_ENQUEUED_VALUE, linkAttributes.get(AttributeKey.longKey(MESSAGE_ENQUEUED_TIME)));
+    }
+
+    @Test
+    public void startConsumeSpanWitStartTimeInContext() {
+        // Arrange
+        final Context spanBuilder = openTelemetryTracer.getSharedSpanBuilder("span",
+            new Context("span-start-time", Instant.now().minusSeconds(1000)));
+
+        Span link = tracer.spanBuilder("link1").startSpan();
+
+        openTelemetryTracer.addLink(spanBuilder.addData(SPAN_CONTEXT_KEY, link.getSpanContext()));
+
+        // Act
+        final Context spanCtx = openTelemetryTracer.start(METHOD_NAME, spanBuilder, ProcessKind.PROCESS);
+        openTelemetryTracer.end(null, null, spanCtx);
+
+        // Assert
+        ReadableSpan span = getSpan(spanCtx);
+        assertEquals(1, span.toSpanData().getLinks().size());
+        assertEquals(span.getLatencyNanos() / 1000_000_000d, 1000d, 10);
+    }
+
     @Test
     public void startSpanOverloadNullPointerException() {
 
@@ -930,7 +982,6 @@ private static Stream<Arguments> spanKinds() {
             Arguments.of(com.azure.core.util.tracing.SpanKind.SERVER, com.azure.core.util.tracing.SpanKind.PRODUCER, false),
             Arguments.of(com.azure.core.util.tracing.SpanKind.SERVER, com.azure.core.util.tracing.SpanKind.CONSUMER, false),
             Arguments.of(com.azure.core.util.tracing.SpanKind.SERVER, com.azure.core.util.tracing.SpanKind.SERVER, false));
-
     }
 
     @Test

sdk/eventhubs/azure-messaging-eventhubs/CHANGELOG.md

@@ -14,6 +14,8 @@
 
 ### Features Added
 
+- Enabled metrics for sent events, consumer lag, checkpointing. ([#31024](https://github.com/Azure/azure-sdk-for-java/pull/31024))
+- Enabled distributed tracing for consumer and batch processor. ([#31197](https://github.com/Azure/azure-sdk-for-java/pull/31197)) 
 - Added algorithm for mapping partition keys to partition ids.
 - Added EventHubBufferedProducerAsyncClient and EventHubBufferedProducerClient