Instrument addition pulsar receive methods

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

@@ -6,10 +6,10 @@
 package io.opentelemetry.javaagent.instrumentation.pulsar.v2_8;
 
 import static io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarSingletons.startAndEndConsumerReceive;
+import static io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.telemetry.PulsarSingletons.wrap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isProtected;
-import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -52,21 +52,16 @@ public void transform(TypeTransformer transformer) {
             .and(takesArguments(2))
             .and(takesArgument(1, named("java.util.concurrent.TimeUnit"))),
         className + "$ConsumerInternalReceiveAdviser");
-    // receive/batchReceive will apply to Consumer#receive()/Consumer#batchReceive()
+    // internalReceive will apply to Consumer#receive()
     transformer.applyAdviceToMethod(
-        isMethod()
-            .and(isPublic())
-            .and(namedOneOf("receive", "batchReceive"))
-            .and(takesArguments(0)),
+        isMethod().and(isProtected()).and(named("internalReceive")).and(takesArguments(0)),
         className + "$ConsumerSyncReceiveAdviser");
-    // receiveAsync/batchReceiveAsync will apply to
-    // Consumer#receiveAsync()/Consumer#batchReceiveAsync()
+    // internalReceiveAsync will apply to Consumer#receiveAsync()
     transformer.applyAdviceToMethod(
-        isMethod()
-            .and(isPublic())
-            .and(namedOneOf("receiveAsync", "batchReceiveAsync"))
-            .and(takesArguments(0)),
+        isMethod().and(isProtected()).and(named("internalReceiveAsync")).and(takesArguments(0)),
         className + "$ConsumerAsyncReceiveAdviser");
+    // TODO batch receiving not implemented (Consumer#batchReceive() and
+    // Consumer#batchReceiveAsync())
   }
 
   @SuppressWarnings("unused")
@@ -97,14 +92,10 @@ public static void after(
         @Advice.Enter Timer timer,
         @Advice.This Consumer<?> consumer,
         @Advice.Return Message<?> message,
-        @Advice.Thrown Throwable t) {
-      if (t != null) {
-        return;
-      }
-
+        @Advice.Thrown Throwable throwable) {
       Context parent = Context.current();
-      Context current = startAndEndConsumerReceive(parent, message, timer, consumer);
-      if (current != null) {
+      Context current = startAndEndConsumerReceive(parent, message, timer, consumer, throwable);
+      if (current != null && throwable == null) {
         // ConsumerBase#internalReceive(long,TimeUnit) will be called before
         // ConsumerListener#receive(Consumer,Message), so, need to inject Context into Message.
         VirtualFieldStore.inject(message, current);
@@ -126,13 +117,9 @@ public static void after(
         @Advice.Enter Timer timer,
         @Advice.This Consumer<?> consumer,
         @Advice.Return Message<?> message,
-        @Advice.Thrown Throwable t) {
-      if (t != null) {
-        return;
-      }
-
+        @Advice.Thrown Throwable throwable) {
       Context parent = Context.current();
-      startAndEndConsumerReceive(parent, message, timer, consumer);
+      startAndEndConsumerReceive(parent, message, timer, consumer, throwable);
       // No need to inject context to message.
     }
   }
@@ -150,17 +137,8 @@ public static Timer before() {
     public static void after(
         @Advice.Enter Timer timer,
         @Advice.This Consumer<?> consumer,
-        @Advice.Return CompletableFuture<Message<?>> future) {
-      future.whenComplete(
-          (message, t) -> {
-            if (t != null) {
-              return;
-            }
-
-            Context parent = Context.current();
-            startAndEndConsumerReceive(parent, message, timer, consumer);
-            // No need to inject context to message.
-          });
+        @Advice.Return(readOnly = false) CompletableFuture<Message<?>> future) {
+      future = wrap(future, timer, consumer);
     }
   }
 }

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

@@ -8,6 +8,7 @@
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
@@ -23,6 +24,7 @@
 import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
 import io.opentelemetry.javaagent.instrumentation.pulsar.v2_8.VirtualFieldStore;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import org.apache.pulsar.client.api.Consumer;
 import org.apache.pulsar.client.api.Message;
 
@@ -109,7 +111,7 @@ private static AttributesExtractor<PulsarRequest, Void> createMessagingAttribute
   }
 
   public static Context startAndEndConsumerReceive(
-      Context parent, Message<?> message, Timer timer, Consumer<?> consumer) {
+      Context parent, Message<?> message, Timer timer, Consumer<?> consumer, Throwable throwable) {
     if (message == null) {
       return null;
     }
@@ -126,10 +128,41 @@ public static Context startAndEndConsumerReceive(
         PROPAGATOR.extract(parent, request, MessageTextMapGetter.INSTANCE),
         request,
         null,
-        null,
+        throwable,
         timer.startTime(),
         timer.now());
   }
 
+  public static CompletableFuture<Message<?>> wrap(
+      CompletableFuture<Message<?>> future, Timer timer, Consumer<?> consumer) {
+    Context parent = Context.current();
+    CompletableFuture<Message<?>> result = new CompletableFuture<>();
+    future.whenComplete(
+        (message, throwable) -> {
+          Context context = startAndEndConsumerReceive(parent, message, timer, consumer, throwable);
+          runWithContext(
+              context,
+              () -> {
+                if (throwable != null) {
+                  result.completeExceptionally(throwable);
+                } else {
+                  result.complete(message);
+                }
+              });
+        });
+
+    return result;
+  }
+
+  private static void runWithContext(Context context, Runnable runnable) {
+    if (context != null) {
+      try (Scope ignored = context.makeCurrent()) {
+        runnable.run();
+      }
+    } else {
+      runnable.run();
+    }
+  }
+
   private PulsarSingletons() {}
 }