Implement vertx-kafka-client instrumentation; single record handler

instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/TracingIterator.java

@@ -54,7 +54,9 @@ public ConsumerRecord<K, V> next() {
     closeScopeAndEndSpan();
 
     ConsumerRecord<K, V> next = delegateIterator.next();
-    if (next != null && consumerProcessInstrumenter().shouldStart(parentContext, next)) {
+    if (next != null
+        && consumerProcessInstrumenter().shouldStart(parentContext, next)
+        && KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
       currentRequest = next;
       currentContext = consumerProcessInstrumenter().start(parentContext, currentRequest);
       currentScope = currentContext.makeCurrent();

instrumentation/vertx/vertx-kafka-client-3.5/javaagent/build.gradle.kts

@@ -0,0 +1,69 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.vertx")
+    module.set("vertx-kafka-client")
+    versions.set("[3.5.0,)")
+    assertInverse.set(true)
+    extraDependency("org.apache.kafka:kafka-clients:0.11.0.0")
+  }
+}
+
+dependencies {
+  bootstrap(project(":instrumentation:kafka:kafka-clients:kafka-clients-0.11:bootstrap"))
+  implementation(project(":instrumentation:kafka:kafka-clients:kafka-clients-common:library"))
+
+  library("io.vertx:vertx-kafka-client:3.5.0")
+
+  // vertx-codegen and vertx-docgen dependencies are needed for Xlint's annotation checking
+  library("io.vertx:vertx-codegen:3.0.0")
+  testLibrary("io.vertx:vertx-docgen:3.0.0")
+
+  // vertx-kafka-client 3.5 uses kafka-clients 0.10.2.1 by default, need to bump it to make instrumentation work
+  testImplementation("org.apache.kafka:kafka-clients:0.11.0.0")
+
+  testImplementation(project(":instrumentation:vertx:vertx-kafka-client-3.5:testing"))
+
+  testInstrumentation(project(":instrumentation:kafka:kafka-clients:kafka-clients-0.11:javaagent"))
+}
+
+testing {
+  suites {
+    val testNoReceiveTelemetry by registering(JvmTestSuite::class) {
+      dependencies {
+        implementation("io.vertx:vertx-kafka-client:3.5.0")
+        implementation("io.vertx:vertx-codegen:3.0.0")
+        implementation("io.vertx:vertx-docgen:3.0.0")
+        implementation("org.apache.kafka:kafka-clients:0.11.0.0")
+        implementation(project(":instrumentation:vertx:vertx-kafka-client-3.5:testing"))
+      }
+
+      targets {
+        all {
+          testTask.configure {
+            usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
+
+            jvmArgs("-Dotel.instrumentation.kafka.experimental-span-attributes=false")
+            jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=false")
+          }
+        }
+      }
+    }
+  }
+}
+
+tasks {
+  test {
+    usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
+
+    jvmArgs("-Dotel.instrumentation.kafka.experimental-span-attributes=true")
+    jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
+  }
+
+  check {
+    dependsOn(testing.suites)
+  }
+}

instrumentation/vertx/vertx-kafka-client-3.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/kafka/v3_5/InstrumentedSingleRecordHandler.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.kafka.v3_5;
+
+import static io.opentelemetry.javaagent.instrumentation.vertx.kafka.v3_5.VertxKafkaSingletons.processInstrumenter;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
+import io.vertx.core.Handler;
+import javax.annotation.Nullable;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+public final class InstrumentedSingleRecordHandler<K, V> implements Handler<ConsumerRecord<K, V>> {
+
+  private final VirtualField<ConsumerRecord<K, V>, Context> receiveContextField;
+  @Nullable private final Handler<ConsumerRecord<K, V>> delegate;
+
+  public InstrumentedSingleRecordHandler(
+      VirtualField<ConsumerRecord<K, V>, Context> receiveContextField,
+      @Nullable Handler<ConsumerRecord<K, V>> delegate) {
+    this.receiveContextField = receiveContextField;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void handle(ConsumerRecord<K, V> record) {
+    Context parentContext = getParentContext(record);
+
+    if (!processInstrumenter().shouldStart(parentContext, record)) {
+      callDelegateHandler(record);
+      return;
+    }
+
+    Context context = processInstrumenter().start(parentContext, record);
+    Throwable error = null;
+    try (Scope ignored = context.makeCurrent()) {
+      callDelegateHandler(record);
+    } catch (Throwable t) {
+      error = t;
+      throw t;
+    } finally {
+      processInstrumenter().end(context, record, null, error);
+    }
+  }
+
+  private Context getParentContext(ConsumerRecord<K, V> records) {
+    Context receiveContext = receiveContextField.get(records);
+
+    // use the receive CONSUMER span as parent if it's available
+    return receiveContext != null ? receiveContext : Context.current();
+  }
+
+  private void callDelegateHandler(ConsumerRecord<K, V> record) {
+    if (delegate != null) {
+      delegate.handle(record);
+    }
+  }
+}

instrumentation/vertx/vertx-kafka-client-3.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/kafka/v3_5/KafkaConsumerRecordsImplInstrumentation.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.kafka.v3_5;
+
+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.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
+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 KafkaConsumerRecordsImplInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.vertx.kafka.client.consumer.impl.KafkaConsumerRecordsImpl");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("recordAt")
+            .and(isPublic())
+            .and(takesArguments(int.class))
+            .and(returns(named("io.vertx.kafka.client.consumer.KafkaConsumerRecord"))),
+        this.getClass().getName() + "$RecordAtAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class RecordAtAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static boolean onEnter() {
+      return KafkaClientsConsumerProcessTracing.setEnabled(false);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void onExit(@Advice.Enter boolean previousValue) {
+      KafkaClientsConsumerProcessTracing.setEnabled(previousValue);
+    }
+  }
+}

instrumentation/vertx/vertx-kafka-client-3.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/kafka/v3_5/KafkaReadStreamImplInstrumentation.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.kafka.v3_5;
+
+import static net.bytebuddy.matcher.ElementMatchers.isPrivate;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
+import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.vertx.core.Handler;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+
+public class KafkaReadStreamImplInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.vertx.kafka.client.consumer.impl.KafkaReadStreamImpl");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("handler")
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("io.vertx.core.Handler"))),
+        this.getClass().getName() + "$HandlerAdvice");
+    transformer.applyAdviceToMethod(
+        named("batchHandler")
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("io.vertx.core.Handler"))),
+        this.getClass().getName() + "$BatchHandlerAdvice");
+    transformer.applyAdviceToMethod(
+        named("run").and(isPrivate()), this.getClass().getName() + "$RunAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class HandlerAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static <K, V> void onEnter(
+        @Advice.Argument(value = 0, readOnly = false) Handler<ConsumerRecord<K, V>> handler) {
+
+      VirtualField<ConsumerRecord<K, V>, Context> receiveContextField =
+          VirtualField.find(ConsumerRecord.class, Context.class);
+      handler = new InstrumentedSingleRecordHandler<>(receiveContextField, handler);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class BatchHandlerAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static <K, V> void onEnter(
+        @Advice.Argument(value = 0, readOnly = false) Handler<ConsumerRecords<K, V>> handler) {
+      // TODO: next PR
+    }
+  }
+
+  // this advice suppresses the CONSUMER spans created by the kafka-clients instrumentation
+  @SuppressWarnings("unused")
+  public static class RunAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static boolean onEnter() {
+      return KafkaClientsConsumerProcessTracing.setEnabled(false);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Enter boolean previousValue) {
+      KafkaClientsConsumerProcessTracing.setEnabled(previousValue);
+    }
+  }
+}

instrumentation/vertx/vertx-kafka-client-3.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/kafka/v3_5/VertxKafkaInstrumentationModule.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.kafka.v3_5;
+
+import static java.util.Arrays.asList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class VertxKafkaInstrumentationModule extends InstrumentationModule {
+
+  public VertxKafkaInstrumentationModule() {
+    super("vertx-kafka-client", "vertx-kafka-client-3.5", "vertx");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return asList(
+        new KafkaReadStreamImplInstrumentation(), new KafkaConsumerRecordsImplInstrumentation());
+  }
+}

instrumentation/vertx/vertx-kafka-client-3.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/kafka/v3_5/VertxKafkaSingletons.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.kafka.v3_5;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaInstrumenterFactory;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+public final class VertxKafkaSingletons {
+
+  private static final String INSTRUMENTATION_NAME = "io.opentelemetry.vertx-kafka-client-3.5";
+
+  private static final Instrumenter<ConsumerRecord<?, ?>, Void> PROCESS_INSTRUMENTER =
+      new KafkaInstrumenterFactory(GlobalOpenTelemetry.get(), INSTRUMENTATION_NAME)
+          .createConsumerProcessInstrumenter();
+
+  public static Instrumenter<ConsumerRecord<?, ?>, Void> processInstrumenter() {
+    return PROCESS_INSTRUMENTER;
+  }
+
+  private VertxKafkaSingletons() {}
+}