Instrument reactor-kafka (#8439)

Author: Mateusz Rzeszutek

instrumentation/kafka/kafka-clients/kafka-clients-common/library/src/main/java/io/opentelemetry/instrumentation/kafka/internal/KafkaConsumerContextUtil.java

@@ -51,5 +51,10 @@ public static void set(ConsumerRecords<?, ?> records, Context context, Consumer<
     recordsConsumerField.set(records, consumer);
   }
 
+  public static void copy(ConsumerRecord<?, ?> from, ConsumerRecord<?, ?> to) {
+    recordContextField.set(to, recordContextField.get(from));
+    recordConsumerField.set(to, recordConsumerField.get(from));
+  }
+
   private KafkaConsumerContextUtil() {}
 }

instrumentation/reactor/reactor-kafka-1.0/javaagent/build.gradle.kts

@@ -0,0 +1,48 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.projectreactor.kafka")
+    module.set("reactor-kafka")
+    // TODO: add support for 1.3
+    versions.set("[1.0.0,1.3.0)")
+  }
+}
+
+dependencies {
+  compileOnly("com.google.auto.value:auto-value-annotations")
+  annotationProcessor("com.google.auto.value:auto-value")
+
+  bootstrap(project(":instrumentation:kafka:kafka-clients:kafka-clients-0.11:bootstrap"))
+
+  implementation(project(":instrumentation:kafka:kafka-clients:kafka-clients-common:library"))
+  implementation(project(":instrumentation:reactor:reactor-3.1:library"))
+
+  library("io.projectreactor.kafka:reactor-kafka:1.0.0.RELEASE")
+
+  testInstrumentation(project(":instrumentation:kafka:kafka-clients:kafka-clients-0.11:javaagent"))
+  testInstrumentation(project(":instrumentation:reactor:reactor-3.1:javaagent"))
+
+  testImplementation("org.testcontainers:kafka")
+
+  testLibrary("io.projectreactor:reactor-test:3.1.0.RELEASE")
+
+  latestDepTestLibrary("io.projectreactor:reactor-core:3.4.+")
+  // TODO: add support for 1.3
+  latestDepTestLibrary("io.projectreactor.kafka:reactor-kafka:1.2.+")
+}
+
+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/reactor/reactor-kafka-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/DefaultKafkaReceiverInstrumentation.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+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;
+import reactor.core.publisher.Flux;
+
+// handles versions 1.0.0 - 1.2.+
+public class DefaultKafkaReceiverInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("reactor.kafka.receiver.internals.DefaultKafkaReceiver");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("createConsumerFlux").and(returns(named("reactor.core.publisher.Flux"))),
+        this.getClass().getName() + "$CreateConsumerFluxAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class CreateConsumerFluxAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Return(readOnly = false) Flux<?> flux) {
+      if (!(flux instanceof TracingDisablingKafkaFlux)) {
+        flux = new TracingDisablingKafkaFlux<>(flux);
+      }
+    }
+  }
+}

instrumentation/reactor/reactor-kafka-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/InstrumentedKafkaFlux.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0;
+
+import static io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0.ReactorKafkaSingletons.processInstrumenter;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaProcessRequest;
+import io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.Scannable;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.FluxOperator;
+import reactor.core.publisher.Operators;
+
+final class InstrumentedKafkaFlux<R extends ConsumerRecord<?, ?>> extends FluxOperator<R, R> {
+
+  InstrumentedKafkaFlux(Flux<R> source) {
+    super(source);
+  }
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public void subscribe(CoreSubscriber<? super R> actual) {
+    source.subscribe(new InstrumentedSubscriber((CoreSubscriber<ConsumerRecord<?, ?>>) actual));
+  }
+
+  static final class InstrumentedSubscriber
+      implements CoreSubscriber<ConsumerRecord<?, ?>>, Subscription, Scannable {
+
+    private final CoreSubscriber<ConsumerRecord<?, ?>> actual;
+    private final Context currentContext;
+    private Subscription subscription;
+
+    InstrumentedSubscriber(CoreSubscriber<ConsumerRecord<?, ?>> actual) {
+      this.actual = actual;
+      currentContext =
+          ContextPropagationOperator.getOpenTelemetryContext(
+              actual.currentContext(), Context.current());
+    }
+
+    @Override
+    public void onSubscribe(Subscription s) {
+      if (Operators.validate(this.subscription, s)) {
+        this.subscription = s;
+
+        actual.onSubscribe(this);
+      }
+    }
+
+    @Override
+    public reactor.util.context.Context currentContext() {
+      return actual.currentContext();
+    }
+
+    @Override
+    public void onNext(ConsumerRecord<?, ?> record) {
+      KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(record);
+      Context receiveContext = consumerContext.getContext();
+      // use the receive CONSUMER span as parent if it's available
+      Context parentContext = receiveContext != null ? receiveContext : currentContext;
+
+      KafkaProcessRequest request = KafkaProcessRequest.create(consumerContext, record);
+      if (!processInstrumenter().shouldStart(parentContext, request)) {
+        actual.onNext(record);
+        return;
+      }
+
+      Context context = processInstrumenter().start(parentContext, request);
+      Throwable error = null;
+      try (Scope ignored = context.makeCurrent()) {
+        actual.onNext(record);
+      } catch (Throwable t) {
+        error = t;
+        throw t;
+      } finally {
+        processInstrumenter().end(context, request, null, error);
+      }
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+      try (Scope ignored = currentContext.makeCurrent()) {
+        actual.onError(throwable);
+      }
+    }
+
+    @Override
+    public void onComplete() {
+      try (Scope ignored = currentContext.makeCurrent()) {
+        actual.onComplete();
+      }
+    }
+
+    @Override
+    public void request(long l) {
+      subscription.request(l);
+    }
+
+    @Override
+    public void cancel() {
+      subscription.cancel();
+    }
+
+    @SuppressWarnings("rawtypes") // that's how the method is defined
+    @Override
+    public Object scanUnsafe(Attr key) {
+      if (key == Attr.ACTUAL) {
+        return actual;
+      }
+      if (key == Attr.PARENT) {
+        return subscription;
+      }
+      return null;
+    }
+  }
+}

instrumentation/reactor/reactor-kafka-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/InstrumentedKafkaReceiver.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0;
+
+import java.util.function.Function;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.kafka.receiver.KafkaReceiver;
+import reactor.kafka.receiver.ReceiverRecord;
+import reactor.kafka.sender.TransactionManager;
+
+public final class InstrumentedKafkaReceiver<K, V> implements KafkaReceiver<K, V> {
+
+  private final KafkaReceiver<K, V> actual;
+
+  public InstrumentedKafkaReceiver(KafkaReceiver<K, V> actual) {
+    this.actual = actual;
+  }
+
+  @Override
+  public Flux<ReceiverRecord<K, V>> receive() {
+    return new InstrumentedKafkaFlux<>(actual.receive());
+  }
+
+  @Override
+  public Flux<Flux<ConsumerRecord<K, V>>> receiveAutoAck() {
+    return actual.receiveAutoAck().map(InstrumentedKafkaFlux::new);
+  }
+
+  @Override
+  public Flux<ConsumerRecord<K, V>> receiveAtmostOnce() {
+    return new InstrumentedKafkaFlux<>(actual.receiveAtmostOnce());
+  }
+
+  @Override
+  public Flux<Flux<ConsumerRecord<K, V>>> receiveExactlyOnce(
+      TransactionManager transactionManager) {
+    return actual.receiveAutoAck().map(InstrumentedKafkaFlux::new);
+  }
+
+  @Override
+  public <T> Mono<T> doOnConsumer(Function<Consumer<K, V>, ? extends T> function) {
+    return actual.doOnConsumer(function);
+  }
+}

instrumentation/reactor/reactor-kafka-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/KafkaReceiverInstrumentation.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0;
+
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+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;
+import reactor.kafka.receiver.KafkaReceiver;
+
+public class KafkaReceiverInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("reactor.kafka.receiver.KafkaReceiver");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("create").and(isStatic()).and(returns(named("reactor.kafka.receiver.KafkaReceiver"))),
+        this.getClass().getName() + "$CreateAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class CreateAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Return(readOnly = false) KafkaReceiver<?, ?> receiver) {
+      if (!(receiver instanceof InstrumentedKafkaReceiver)) {
+        receiver = new InstrumentedKafkaReceiver<>(receiver);
+      }
+    }
+  }
+}

instrumentation/reactor/reactor-kafka-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/ReactorKafkaInstrumentationModule.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0;
+
+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 ReactorKafkaInstrumentationModule extends InstrumentationModule {
+
+  public ReactorKafkaInstrumentationModule() {
+    super("reactor-kafka", "reactor-kafka-1.0");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return asList(
+        new KafkaReceiverInstrumentation(),
+        new ReceiverRecordInstrumentation(),
+        new DefaultKafkaReceiverInstrumentation());
+  }
+}

instrumentation/reactor/reactor-kafka-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/ReactorKafkaSingletons.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.kafka.v1_0;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaInstrumenterFactory;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaProcessRequest;
+import io.opentelemetry.javaagent.bootstrap.internal.ExperimentalConfig;
+import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
+
+final class ReactorKafkaSingletons {
+
+  private static final String INSTRUMENTATION_NAME = "io.opentelemetry.reactor-kafka-1.0";
+
+  private static final Instrumenter<KafkaProcessRequest, Void> PROCESS_INSTRUMENTER =
+      new KafkaInstrumenterFactory(GlobalOpenTelemetry.get(), INSTRUMENTATION_NAME)
+          .setCapturedHeaders(ExperimentalConfig.get().getMessagingHeaders())
+          .setCaptureExperimentalSpanAttributes(
+              InstrumentationConfig.get()
+                  .getBoolean("otel.instrumentation.kafka.experimental-span-attributes", false))
+          .setMessagingReceiveInstrumentationEnabled(
+              ExperimentalConfig.get().messagingReceiveInstrumentationEnabled())
+          .createConsumerProcessInstrumenter();
+
+  public static Instrumenter<KafkaProcessRequest, Void> processInstrumenter() {
+    return PROCESS_INSTRUMENTER;
+  }
+
+  private ReactorKafkaSingletons() {}
+}