Pulsar: use span links when receive telemetry is enabled

instrumentation/pulsar/pulsar-2.8/javaagent/build.gradle.kts

@@ -19,9 +19,28 @@ dependencies {
   testImplementation("org.apache.pulsar:pulsar-client-admin:2.8.0")
 }
 
+tasks {
+  val testReceiveSpanDisabled by registering(Test::class) {
+    filter {
+      includeTestsMatching("PulsarClientSuppressReceiveSpansTest")
+    }
+    include("**/PulsarClientSuppressReceiveSpansTest.*")
+  }
+
+  test {
+    filter {
+      excludeTestsMatching("PulsarClientSuppressReceiveSpansTest")
+    }
+    jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
+  }
+
+  check {
+    dependsOn(testReceiveSpanDisabled)
+  }
+}
+
 tasks.withType<Test>().configureEach {
   // TODO run tests both with and without experimental span attributes
   jvmArgs("-Dotel.instrumentation.pulsar.experimental-span-attributes=true")
-  jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
   usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
 }

instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/ConsumerBaseInstrumentation.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.pulsar.v2_8;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.MessageListenerContext;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ConsumerBaseInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.apache.pulsar.client.impl.ConsumerBase")
+        .or(named("org.apache.pulsar.client.impl.MultiTopicsConsumerImpl"));
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    // these methods receive a message and pass it on to a message listener
+    // we instrument them so that the span for the receive operation could be suppressed
+    transformer.applyAdviceToMethod(
+        named("triggerListener").and(takesArguments(0)).or(named("receiveMessageFromConsumer")),
+        this.getClass().getName() + "$TriggerListenerAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class TriggerListenerAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      MessageListenerContext.startProcessing();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      MessageListenerContext.endProcessing();
+    }
+  }
+}

instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/PulsarInstrumentationModule.java

@@ -20,6 +20,7 @@ public PulsarInstrumentationModule() {
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return Arrays.asList(
+        new ConsumerBaseInstrumentation(),
         new ConsumerImplInstrumentation(),
         new ProducerImplInstrumentation(),
         new MessageInstrumentation(),

instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/telemetry/MessageListenerContext.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry;
+
+/**
+ * Helper class used to determine whether message is going to be processed by a listener. If we know
+ * that message is going to be passed to a message listener, that would produce a span for the
+ * "process" operation, we are going to suppress the span from the message "receive" operation.
+ */
+public final class MessageListenerContext {
+  private static final ThreadLocal<Boolean> processing = new ThreadLocal<>();
+
+  private MessageListenerContext() {}
+
+  /** Call on entry to a method that will pass the received message to a message listener. */
+  public static void startProcessing() {
+    processing.set(Boolean.TRUE);
+  }
+
+  public static void endProcessing() {
+    processing.remove();
+  }
+
+  /** Returns true if we expect a received message to be passed to a listener. */
+  public static boolean isProcessing() {
+    return processing.get() != null;
+  }
+}

instrumentation/pulsar/pulsar-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pulsar/v2_8/telemetry/PulsarSingletons.java

@@ -18,7 +18,9 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.InstrumenterBuilder;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanLinksExtractor;
 import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
+import io.opentelemetry.instrumentation.api.internal.PropagatorBasedSpanLinksExtractor;
 import io.opentelemetry.instrumentation.api.internal.Timer;
 import io.opentelemetry.instrumentation.api.semconv.network.ServerAttributesExtractor;
 import io.opentelemetry.javaagent.bootstrap.internal.ExperimentalConfig;
@@ -38,6 +40,8 @@ public final class PulsarSingletons {
       TELEMETRY.getPropagators().getTextMapPropagator();
   private static final List<String> capturedHeaders =
       ExperimentalConfig.get().getMessagingHeaders();
+  private static final boolean receiveInstrumentationEnabled =
+      ExperimentalConfig.get().messagingReceiveInstrumentationEnabled();
 
   private static final Instrumenter<PulsarRequest, Void> CONSUMER_PROCESS_INSTRUMENTER =
       createConsumerProcessInstrumenter();
@@ -64,15 +68,23 @@ private static Instrumenter<PulsarRequest, Void> createConsumerReceiveInstrument
     MessagingAttributesGetter<PulsarRequest, Void> getter =
         PulsarMessagingAttributesGetter.INSTANCE;
 
-    return Instrumenter.<PulsarRequest, Void>builder(
-            TELEMETRY,
-            INSTRUMENTATION_NAME,
-            MessagingSpanNameExtractor.create(getter, MessageOperation.RECEIVE))
-        .addAttributesExtractor(
-            createMessagingAttributesExtractor(getter, MessageOperation.RECEIVE))
-        .addAttributesExtractor(
-            ServerAttributesExtractor.create(new PulsarNetClientAttributesGetter()))
-        .buildConsumerInstrumenter(MessageTextMapGetter.INSTANCE);
+    InstrumenterBuilder<PulsarRequest, Void> instrumenterBuilder =
+        Instrumenter.<PulsarRequest, Void>builder(
+                TELEMETRY,
+                INSTRUMENTATION_NAME,
+                MessagingSpanNameExtractor.create(getter, MessageOperation.RECEIVE))
+            .addAttributesExtractor(
+                createMessagingAttributesExtractor(getter, MessageOperation.RECEIVE))
+            .addAttributesExtractor(
+                ServerAttributesExtractor.create(new PulsarNetClientAttributesGetter()));
+
+    if (receiveInstrumentationEnabled) {
+      return instrumenterBuilder
+          .addSpanLinksExtractor(
+              new PropagatorBasedSpanLinksExtractor<>(PROPAGATOR, MessageTextMapGetter.INSTANCE))
+          .buildInstrumenter(SpanKindExtractor.alwaysConsumer());
+    }
+    return instrumenterBuilder.buildConsumerInstrumenter(MessageTextMapGetter.INSTANCE);
   }
 
   private static Instrumenter<PulsarBatchRequest, Void> createConsumerBatchReceiveInstrumenter() {
@@ -87,24 +99,29 @@ private static Instrumenter<PulsarBatchRequest, Void> createConsumerBatchReceive
             createMessagingAttributesExtractor(getter, MessageOperation.RECEIVE))
         .addAttributesExtractor(
             ServerAttributesExtractor.create(new PulsarNetClientAttributesGetter()))
-        .setEnabled(ExperimentalConfig.get().messagingReceiveInstrumentationEnabled())
-        .addSpanLinksExtractor(
-            new PulsarBatchRequestSpanLinksExtractor(
-                GlobalOpenTelemetry.getPropagators().getTextMapPropagator()))
+        .addSpanLinksExtractor(new PulsarBatchRequestSpanLinksExtractor(PROPAGATOR))
         .buildInstrumenter(SpanKindExtractor.alwaysConsumer());
   }
 
   private static Instrumenter<PulsarRequest, Void> createConsumerProcessInstrumenter() {
     MessagingAttributesGetter<PulsarRequest, Void> getter =
         PulsarMessagingAttributesGetter.INSTANCE;
 
-    return Instrumenter.<PulsarRequest, Void>builder(
-            TELEMETRY,
-            INSTRUMENTATION_NAME,
-            MessagingSpanNameExtractor.create(getter, MessageOperation.PROCESS))
-        .addAttributesExtractor(
-            createMessagingAttributesExtractor(getter, MessageOperation.PROCESS))
-        .buildInstrumenter();
+    InstrumenterBuilder<PulsarRequest, Void> instrumenterBuilder =
+        Instrumenter.<PulsarRequest, Void>builder(
+                TELEMETRY,
+                INSTRUMENTATION_NAME,
+                MessagingSpanNameExtractor.create(getter, MessageOperation.PROCESS))
+            .addAttributesExtractor(
+                createMessagingAttributesExtractor(getter, MessageOperation.PROCESS));
+
+    if (receiveInstrumentationEnabled) {
+      SpanLinksExtractor<PulsarRequest> spanLinksExtractor =
+          new PropagatorBasedSpanLinksExtractor<>(PROPAGATOR, MessageTextMapGetter.INSTANCE);
+      instrumenterBuilder.addSpanLinksExtractor(spanLinksExtractor);
+      return instrumenterBuilder.buildInstrumenter(SpanKindExtractor.alwaysConsumer());
+    }
+    return instrumenterBuilder.buildConsumerInstrumenter(MessageTextMapGetter.INSTANCE);
   }
 
   private static Instrumenter<PulsarRequest, Void> createProducerInstrumenter() {
@@ -146,12 +163,17 @@ public static Context startAndEndConsumerReceive(
     if (!CONSUMER_RECEIVE_INSTRUMENTER.shouldStart(parent, request)) {
       return null;
     }
-    // startAndEnd not supports extract trace context from carrier
-    // start not supports custom startTime
-    // extract trace context by using TEXT_MAP_PROPAGATOR here.
+    if (!receiveInstrumentationEnabled) {
+      // suppress receive span when receive telemetry is not enabled and message is going to be
+      // processed by a listener
+      if (MessageListenerContext.isProcessing()) {
+        return null;
+      }
+      parent = PROPAGATOR.extract(parent, request, MessageTextMapGetter.INSTANCE);
+    }
     return InstrumenterUtil.startAndEnd(
         CONSUMER_RECEIVE_INSTRUMENTER,
-        PROPAGATOR.extract(parent, request, MessageTextMapGetter.INSTANCE),
+        parent,
         request,
         null,
         throwable,
@@ -185,11 +207,17 @@ private static Context startAndEndConsumerReceive(
 
   public static CompletableFuture<Message<?>> wrap(
       CompletableFuture<Message<?>> future, Timer timer, Consumer<?> consumer) {
+    boolean listenerContextActive = MessageListenerContext.isProcessing();
     Context parent = Context.current();
     CompletableFuture<Message<?>> result = new CompletableFuture<>();
     future.whenComplete(
         (message, throwable) -> {
-          Context context = startAndEndConsumerReceive(parent, message, timer, consumer, throwable);
+          // we create a "receive" span when receive telemetry is enabled or when we know that
+          // this message will not be passed to a listener that would create the "process" span
+          Context context =
+              receiveInstrumentationEnabled || !listenerContextActive
+                  ? startAndEndConsumerReceive(parent, message, timer, consumer, throwable)
+                  : parent;
           runWithContext(
               context,
               () -> {