Bootstrap NATS instrumentation with Connection.publish(Message) tracing

Author: AlixBa

.fossa.yml

@@ -691,6 +691,12 @@ targets:
     - type: gradle
       path: ./
       target: ':instrumentation:mongo:mongo-async-3.3:javaagent'
+    - type: gradle
+      path: ./
+      target: ':instrumentation:nats:nats-2.21:javaagent'
+    - type: gradle
+      path: ./
+      target: ':instrumentation:nats:nats-2.21:library'
     - type: gradle
       path: ./
       target: ':instrumentation:netty:netty-3.8:javaagent'

.gitignore

@@ -45,6 +45,7 @@ out/
 ######################
 .vscode
 **/bin/
+.metals
 
 # Others #
 ##########

instrumentation/nats/nats-2.21/javaagent/build.gradle.kts

@@ -0,0 +1,24 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.nats")
+    module.set("jnats")
+    versions.set("[2.21.0,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  library("io.nats:jnats:2.21.0")
+
+  implementation(project(":instrumentation:nats:nats-2.21:library"))
+}
+
+tasks {
+  test {
+    usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
+  }
+}

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

@@ -0,0 +1,76 @@
+/*
+ * 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.PRODUCER_INSTRUMENTER;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.nats.client.Connection;
+import io.nats.client.Message;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ConnectionInstrumentation 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("publish"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("io.nats.client.Message"))),
+        ConnectionInstrumentation.class.getName() + "$PublishAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class PublishAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This Connection connection,
+        @Advice.Argument(0) Message message,
+        @Advice.Local("otelContext") Context otelContext,
+        @Advice.Local("otelScope") Scope otelScope) {
+      Context parentContext = Context.current();
+
+      if (!PRODUCER_INSTRUMENTER.shouldStart(parentContext, message)) {
+        return;
+      }
+
+      otelContext = PRODUCER_INSTRUMENTER.start(parentContext, message);
+      otelScope = otelContext.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.This Connection connection,
+        @Advice.Argument(0) Message message,
+        @Advice.Local("otelContext") Context otelContext,
+        @Advice.Local("otelScope") Scope otelScope) {
+      if (otelScope == null) {
+        return;
+      }
+
+      otelScope.close();
+      PRODUCER_INSTRUMENTER.end(otelContext, message, null, throwable);
+    }
+  }
+}

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_21;
+
+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.List;
+
+@AutoService(InstrumentationModule.class)
+public class NatsInstrumentationModule extends InstrumentationModule {
+
+  public NatsInstrumentationModule() {
+    super("nats", "nats-2.21");
+  }
+
+  // TODO classLoaderMatcher
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Collections.singletonList(new ConnectionInstrumentation());
+  }
+}

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

@@ -0,0 +1,20 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_21;
+
+import static io.opentelemetry.instrumentation.nats.v2_21.internal.NatsInstrumenterFactory.createProducerInstrumenter;
+
+import io.nats.client.Message;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+
+public final class NatsSingletons {
+
+  public static final Instrumenter<Message, Void> PRODUCER_INSTRUMENTER =
+      createProducerInstrumenter(GlobalOpenTelemetry.get());
+
+  private NatsSingletons() {}
+}

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

@@ -0,0 +1,61 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_21;
+
+import io.nats.client.Connection;
+import io.nats.client.Nats;
+import io.nats.client.impl.NatsMessage;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import java.io.IOException;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.utility.DockerImageName;
+
+class NatsInstrumentationTest {
+
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  static final DockerImageName natsImage = DockerImageName.parse("nats:2.11.2-alpine3.21");
+
+  static final GenericContainer<?> natsContainer =
+      new GenericContainer<>(natsImage).withExposedPorts(4222);
+
+  static String natsUrl;
+  static Connection connection;
+
+  @BeforeAll
+  static void beforeAll() throws IOException, InterruptedException {
+    natsContainer.start();
+    natsUrl = "nats://" + natsContainer.getHost() + ":" + natsContainer.getMappedPort(4222);
+    connection = Nats.connect(natsUrl);
+  }
+
+  @AfterAll
+  static void afterAll() throws InterruptedException {
+    connection.close();
+    natsContainer.close();
+  }
+
+  @Test
+  void testConnection() {
+    // given
+    NatsMessage message = NatsMessage.builder().subject("sub").build();
+
+    // when
+    testing.runWithSpan("testConnection", () -> connection.publish(message));
+
+    // then
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("testConnection"), span -> span.hasName("sub publish")));
+  }
+}

instrumentation/nats/nats-2.21/library/build.gradle.kts

@@ -0,0 +1,9 @@
+plugins {
+  id("otel.library-instrumentation")
+}
+
+dependencies {
+  library("io.nats:jnats:2.21.0")
+
+  testImplementation(project(":instrumentation:nats:nats-2.21:testing"))
+}

instrumentation/nats/nats-2.21/library/src/main/java/io/opentelemetry/instrumentation/nats/v2_21/NatsTelemetry.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.nats.v2_21;
+
+import io.nats.client.Connection;
+import io.nats.client.Message;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+
+public final class NatsTelemetry {
+
+  public static NatsTelemetry create(OpenTelemetry openTelemetry) {
+    return new NatsTelemetryBuilder(openTelemetry).build();
+  }
+
+  public static NatsTelemetryBuilder builder(OpenTelemetry openTelemetry) {
+    return new NatsTelemetryBuilder(openTelemetry);
+  }
+
+  private final Instrumenter<Message, Void> producerInstrumenter;
+
+  public NatsTelemetry(Instrumenter<Message, Void> producerInstrumenter) {
+    this.producerInstrumenter = producerInstrumenter;
+  }
+
+  public OpenTelemetryConnection wrap(Connection connection) {
+    return new OpenTelemetryConnection(connection, this.producerInstrumenter);
+  }
+}

instrumentation/nats/nats-2.21/library/src/main/java/io/opentelemetry/instrumentation/nats/v2_21/NatsTelemetryBuilder.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.nats.v2_21;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.instrumentation.nats.v2_21.internal.NatsInstrumenterFactory;
+
+public final class NatsTelemetryBuilder {
+
+  private final OpenTelemetry openTelemetry;
+
+  NatsTelemetryBuilder(OpenTelemetry openTelemetry) {
+    this.openTelemetry = openTelemetry;
+  }
+
+  public NatsTelemetry build() {
+    return new NatsTelemetry(NatsInstrumenterFactory.createProducerInstrumenter(openTelemetry));
+  }
+}