Instrument NATS Subscription.nextMessage

Author: AlixBa

instrumentation/nats/nats-2.21/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_21/ConnectionSubscribeInstrumentation.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_21;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.nats.client.Connection;
+import io.nats.client.Subscription;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.nats.v2_21.internal.NatsData;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ConnectionSubscribeInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return implementsInterface(named("io.nats.client.Connection"));
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(named("subscribe"))
+            .and(takesArgument(0, String.class))
+            .and(returns(named("io.nats.client.Subscription"))),
+        ConnectionSubscribeInstrumentation.class.getName() + "$SubscribeAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class SubscribeAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This Connection connection, @Advice.Return Subscription subscription) {
+      NatsData.addSubscription(subscription, connection);
+    }
+  }
+}

instrumentation/nats/nats-2.21/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_21/NatsInstrumentationModule.java

@@ -8,7 +8,7 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
-import java.util.Collections;
+import java.util.Arrays;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
@@ -22,6 +22,9 @@ public NatsInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return Collections.singletonList(new ConnectionPublishInstrumentation());
+    return Arrays.asList(
+        new ConnectionSubscribeInstrumentation(),
+        new ConnectionPublishInstrumentation(),
+        new SubscriptionInstrumentation());
   }
 }

instrumentation/nats/nats-2.21/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_21/NatsSingletons.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.nats.v2_21;
 
+import static io.opentelemetry.instrumentation.nats.v2_21.internal.NatsInstrumenterFactory.createConsumerInstrumenter;
 import static io.opentelemetry.instrumentation.nats.v2_21.internal.NatsInstrumenterFactory.createProducerInstrumenter;
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
@@ -16,5 +17,8 @@ public final class NatsSingletons {
   public static final Instrumenter<NatsRequest, Void> PRODUCER_INSTRUMENTER =
       createProducerInstrumenter(GlobalOpenTelemetry.get());
 
+  public static final Instrumenter<NatsRequest, Void> CONSUMER_INSTRUMENTER =
+      createConsumerInstrumenter(GlobalOpenTelemetry.get());
+
   private NatsSingletons() {}
 }

instrumentation/nats/nats-2.21/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_21/SubscriptionInstrumentation.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_21;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.instrumentation.nats.v2_21.NatsSingletons.CONSUMER_INSTRUMENTER;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.nats.client.Connection;
+import io.nats.client.Message;
+import io.nats.client.Subscription;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
+import io.opentelemetry.instrumentation.api.internal.Timer;
+import io.opentelemetry.instrumentation.nats.v2_21.internal.NatsRequest;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.nats.v2_21.internal.NatsData;
+import java.util.concurrent.TimeoutException;
+import javax.annotation.Nullable;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class SubscriptionInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return implementsInterface(named("io.nats.client.Subscription"));
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(named("nextMessage"))
+            .and(takesArguments(1))
+            .and(returns(named("io.nats.client.Message"))),
+        SubscriptionInstrumentation.class.getName() + "$NextMessageAdvice");
+    transformer.applyAdviceToMethod(
+        isPublic().and(named("unsubscribe")),
+        SubscriptionInstrumentation.class.getName() + "$UnsubscribeAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class NextMessageAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Timer onEnter() {
+      return Timer.start();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.This Subscription subscription,
+        @Advice.Enter Timer timer,
+        @Advice.Return @Nullable Message message) {
+      Context parentContext = Context.current();
+      TimeoutException timeout = null;
+
+      Connection connection = NatsData.getConnection(subscription);
+
+      // connection should always be non-null at this stage
+      if (connection == null) {
+        return;
+      }
+
+      NatsRequest natsRequest = NatsRequest.create(connection, subscription.getSubject());
+      if (message == null) {
+        timeout = new TimeoutException("Timed out waiting for message");
+      } else {
+        natsRequest = NatsRequest.create(connection, message);
+      }
+
+      if (!CONSUMER_INSTRUMENTER.shouldStart(parentContext, natsRequest)) {
+        return;
+      }
+
+      InstrumenterUtil.startAndEnd(
+          CONSUMER_INSTRUMENTER,
+          parentContext,
+          natsRequest,
+          null,
+          timeout,
+          timer.startTime(),
+          timer.now());
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class UnsubscribeAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.This Subscription subscription) {
+      NatsData.removeSubscription(subscription);
+    }
+  }
+}

instrumentation/nats/nats-2.21/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_21/internal/NatsData.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_21.internal;
+
+import io.nats.client.Connection;
+import io.nats.client.Subscription;
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.", or "This class is internal and experimental. Its APIs are unstable and can change at
+ * any time. Its APIs (or a version of them) may be promoted to the public stable API in the future,
+ * but no guarantees are made.
+ */
+public final class NatsData {
+
+  private static final VirtualField<Subscription, Connection> subscriptionConnection =
+      VirtualField.find(Subscription.class, Connection.class);
+
+  public static void addSubscription(Subscription subscription, Connection connection) {
+    subscriptionConnection.set(subscription, connection);
+  }
+
+  public static void removeSubscription(Subscription subscription) {
+    subscriptionConnection.set(subscription, null);
+  }
+
+  public static Connection getConnection(Subscription subscription) {
+    return subscriptionConnection.get(subscription);
+  }
+
+  private NatsData() {}
+}

instrumentation/nats/nats-2.21/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/nats/v2_21/NatsInstrumentationPublishTest.java

@@ -63,7 +63,7 @@ static void afterAll() {
   @BeforeEach
   void beforeEach() throws IOException, InterruptedException {
     connection = Nats.connect(natsUrl);
-    subscription = connection.subscribe("*");
+    subscription = connection.subscribe("sub");
     clientId = connection.getServerInfo().getClientId();
   }