Add NATS instrumentation

.fossa.yml

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

.github/scripts/check-package-names.sh

@@ -29,6 +29,9 @@ for dir in $(find instrumentation -name "*.java" | grep library/src/main/java |
   if [[ "$dir" == "instrumentation/lettuce/lettuce-5.1/library/src/main/java/io/lettuce/core/protocol" ]]; then
     continue
   fi
+  if [[ "$dir" == "instrumentation/nats/nats-2.17/library/src/main/java/io/nats/client/impl" ]]; then
+    continue
+  fi
 
   # some common modules don't have any base version
   # - lettuce-common

.gitignore

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

docs/instrumentation-list.yaml

@@ -4842,6 +4842,59 @@ libraries:
     target_versions:
       javaagent:
       - org.mybatis:mybatis:[3.2.0,)
+  nats:
+  - name: nats-2.17
+    description: This instrumentation provides messaging spans for NATS
+    disabled_by_default: false
+    source_path: instrumentation/nats/nats-2.17
+    scope:
+      name: io.opentelemetry.nats-2.17
+    target_versions:
+      javaagent:
+      - io.nats:jnats:[2.17.7,)
+      library:
+      - io.nats:jnats:2.17.7
+    configurations:
+    - name: otel.instrumentation.messaging.experimental.receive-telemetry.enabled
+      description: |
+        Enables experimental receive telemetry, which will cause consumers to start a new trace, with only a span link connecting it to the producer trace.
+      type: boolean
+      default: false
+    - name: otel.instrumentation.messaging.experimental.capture-headers
+      description: Allows configuring headers to capture as span attributes.
+      type: list
+      default: ''
+    telemetry:
+    - when: default
+      spans:
+      - span_kind: CONSUMER
+        attributes:
+        - name: messaging.client_id
+          type: STRING
+        - name: messaging.destination.name
+          type: STRING
+        - name: messaging.header.captured_header
+          type: STRING_ARRAY
+        - name: messaging.message.body.size
+          type: LONG
+        - name: messaging.operation
+          type: STRING
+        - name: messaging.system
+          type: STRING
+      - span_kind: PRODUCER
+        attributes:
+        - name: messaging.client_id
+          type: STRING
+        - name: messaging.destination.name
+          type: STRING
+        - name: messaging.header.captured_header
+          type: STRING_ARRAY
+        - name: messaging.message.body.size
+          type: LONG
+        - name: messaging.operation
+          type: STRING
+        - name: messaging.system
+          type: STRING
   netty:
   - name: netty-3.8
     source_path: instrumentation/netty/netty-3.8

docs/supported-libraries.md

@@ -106,6 +106,7 @@ These are the supported libraries and frameworks:
 | [Micrometer](https://micrometer.io/)                                                                                                        | 1.5+ (disabled by default)         | [opentelemetry-micrometer-1.5](../instrumentation/micrometer/micrometer-1.5/library)                                                                                                                                                                                                                                                                                                    | none                                                                                                                                              |
 | [MongoDB Driver](https://mongodb.github.io/mongo-java-driver/)                                                                              | 3.1+                               | [opentelemetry-mongo-3.1](../instrumentation/mongo/mongo-3.1/library)                                                                                                                                                                                                                                                                                                                   | [Database Client Spans], [Database Client Metrics] [6]                                                                                       |
 | [MyBatis](https://mybatis.org/mybatis-3/)                                                                                                   | 3.2+                               | N/A                                                                                                                                                                                                                                                                                                                                                                                     | none                                                                                                                                              |
+| [NATS Client](https://github.com/nats-io/nats.java)                                                                                         | 2.17.2+                            | [nats-2.17](../instrumentation/nats/nats-2.17/library)                                                                                                                                                                                                                                                                                                                                  | [Messaging Spans]                                                                                                                                 |
 | [Netty HTTP codec [5]](https://github.com/netty/netty)                                                                                      | 3.8+                               | [opentelemetry-netty-4.1](../instrumentation/netty/netty-4.1/library)                                                                                                                                                                                                                                                                                                                   | [HTTP Client Spans], [HTTP Client Metrics], [HTTP Server Spans], [HTTP Server Metrics]                                                            |
 | [OpenAI Java SDK](https://github.com/openai/openai-java)                                                                                    | 1.1+                               | [openai-java-1.1](../instrumentation/openai/openai-java-1.1/library)                                                                                                                                                                                                                                                                                                                    | [GenAI Client Spans], [GenAI Client Metrics]                                                                                                      |
 | [OpenSearch Rest Client](https://github.com/opensearch-project/opensearch-java)                                                             | 1.0+                               |                                                                                                                                                                                                                                                                                                                                                                                         | [Database Client Spans], [Database Client Metrics] [6]                                                                                       |

instrumentation-docs/instrumentations.sh

@@ -158,6 +158,8 @@ readonly INSTRUMENTATIONS=(
   "influxdb-2.4:javaagent:test"
   "influxdb-2.4:javaagent:testStableSemconv"
   "java-http-server:javaagent:test"
+  "nats:nats-2.17:javaagent:test"
+  "nats:nats-2.17:javaagent:testExperimental"
 )
 
 #  Some instrumentation test suites don't run ARM, so we use colima to run them in an x86_64

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

@@ -0,0 +1,49 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.nats")
+    module.set("jnats")
+    versions.set("[2.17.2,)")
+
+    // Could not find io.nats:nats-parent:1.0-SNAPSHOT
+    skip("0.5.0", "0.5.1")
+
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  library("io.nats:jnats:2.17.2")
+
+  implementation(project(":instrumentation:nats:nats-2.17:library"))
+  testImplementation(project(":instrumentation:nats:nats-2.17:testing"))
+}
+
+tasks {
+  withType<Test>().configureEach {
+    usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
+    systemProperty("collectMetadata", findProperty("collectMetadata")?.toString() ?: "false")
+  }
+
+  val testExperimental by registering(Test::class) {
+    testClassesDirs = sourceSets.test.get().output.classesDirs
+    classpath = sourceSets.test.get().runtimeClasspath
+    filter {
+      includeTestsMatching("NatsExperimentalTest")
+    }
+    jvmArgs("-Dotel.instrumentation.messaging.experimental.capture-headers=captured-header")
+  }
+
+  test {
+    filter {
+      excludeTestsMatching("NatsExperimentalTest")
+    }
+  }
+
+  check {
+    dependsOn(testExperimental)
+  }
+}

instrumentation/nats/nats-2.17/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_17/CompletableFutureWrapper.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_17;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import java.util.concurrent.CompletableFuture;
+
+public final class CompletableFutureWrapper {
+
+  private CompletableFutureWrapper() {}
+
+  public static <T> CompletableFuture<T> wrap(CompletableFuture<T> future, Context context) {
+    CompletableFuture<T> result = new CompletableFuture<>();
+    future.whenComplete(
+        (T value, Throwable throwable) -> {
+          try (Scope ignored = context.makeCurrent()) {
+            if (throwable != null) {
+              result.completeExceptionally(throwable);
+            } else {
+              result.complete(value);
+            }
+          }
+        });
+
+    return result;
+  }
+}

instrumentation/nats/nats-2.17/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/nats/v2_17/ConnectionPublishInstrumentation.java

@@ -0,0 +1,174 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.nats.v2_17;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.instrumentation.nats.v2_17.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.nats.client.impl.Headers;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.nats.v2_17.internal.NatsMessageWritableHeaders;
+import io.opentelemetry.instrumentation.nats.v2_17.internal.NatsRequest;
+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 ConnectionPublishInstrumentation 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(2))
+            .and(takesArgument(0, String.class))
+            .and(takesArgument(1, byte[].class)),
+        ConnectionPublishInstrumentation.class.getName() + "$PublishBodyAdvice");
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(named("publish"))
+            .and(takesArguments(3))
+            .and(takesArgument(0, String.class))
+            .and(takesArgument(1, named("io.nats.client.impl.Headers")))
+            .and(takesArgument(2, byte[].class)),
+        ConnectionPublishInstrumentation.class.getName() + "$PublishHeadersBodyAdvice");
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(named("publish"))
+            .and(takesArguments(3))
+            .and(takesArgument(0, String.class))
+            .and(takesArgument(1, String.class))
+            .and(takesArgument(2, byte[].class)),
+        ConnectionPublishInstrumentation.class.getName() + "$PublishReplyToBodyAdvice");
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(named("publish"))
+            .and(takesArguments(4))
+            .and(takesArgument(0, String.class))
+            .and(takesArgument(1, String.class))
+            .and(takesArgument(2, named("io.nats.client.impl.Headers")))
+            .and(takesArgument(3, byte[].class)),
+        ConnectionPublishInstrumentation.class.getName() + "$PublishReplyToHeadersBodyAdvice");
+    transformer.applyAdviceToMethod(
+        isPublic()
+            .and(named("publish"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("io.nats.client.Message"))),
+        ConnectionPublishInstrumentation.class.getName() + "$PublishMessageAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class PublishBodyAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static boolean onEnter(
+        @Advice.This Connection connection,
+        @Advice.Argument(0) String subject,
+        @Advice.Argument(1) byte[] body) {
+      // call the instrumented publish method
+      connection.publish(subject, null, null, body);
+      return true;
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class PublishHeadersBodyAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static boolean onEnter(
+        @Advice.This Connection connection,
+        @Advice.Argument(0) String subject,
+        @Advice.Argument(1) Headers headers,
+        @Advice.Argument(2) byte[] body) {
+      // call the instrumented publish method
+      connection.publish(subject, null, headers, body);
+      return true;
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class PublishReplyToBodyAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static boolean onEnter(
+        @Advice.This Connection connection,
+        @Advice.Argument(0) String subject,
+        @Advice.Argument(1) String replyTo,
+        @Advice.Argument(2) byte[] body) {
+      // call the instrumented publish method
+      connection.publish(subject, replyTo, null, body);
+      return true;
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class PublishReplyToHeadersBodyAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This Connection connection,
+        @Advice.Argument(0) String subject,
+        @Advice.Argument(1) String replyTo,
+        @Advice.Argument(value = 2, readOnly = false) Headers headers,
+        @Advice.Argument(3) byte[] body,
+        @Advice.Local("otelContext") Context otelContext,
+        @Advice.Local("otelScope") Scope otelScope,
+        @Advice.Local("natsRequest") NatsRequest natsRequest) {
+      headers = NatsMessageWritableHeaders.create(headers);
+
+      Context parentContext = Context.current();
+      natsRequest = NatsRequest.create(connection, subject, replyTo, headers, body);
+
+      if (!PRODUCER_INSTRUMENTER.shouldStart(parentContext, natsRequest)) {
+        return;
+      }
+
+      otelContext = PRODUCER_INSTRUMENTER.start(parentContext, natsRequest);
+      otelScope = otelContext.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelContext") Context otelContext,
+        @Advice.Local("otelScope") Scope otelScope,
+        @Advice.Local("natsRequest") NatsRequest natsRequest) {
+      if (otelScope == null) {
+        return;
+      }
+
+      otelScope.close();
+      PRODUCER_INSTRUMENTER.end(otelContext, natsRequest, null, throwable);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class PublishMessageAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static boolean onEnter(
+        @Advice.This Connection connection, @Advice.Argument(0) Message message) {
+      if (message == null) {
+        return false;
+      }
+
+      // call the instrumented publish method
+      connection.publish(
+          message.getSubject(), message.getReplyTo(), message.getHeaders(), message.getData());
+      return true;
+    }
+  }
+}