Customize kafka instrumentation

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

@@ -106,7 +106,7 @@ public static <K, V> void wrap(
       // case it's important to overwrite the leaked span instead of suppressing the correct span
       // (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
       KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
-      iterator = TracingIterator.wrap(iterator, consumerContext);
+      iterator = TracingIterator.wrap(iterator, consumerContext, true);
     }
   }
 }

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

@@ -29,6 +29,8 @@ public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new KafkaProducerInstrumentation(),
         new KafkaConsumerInstrumentation(),
-        new ConsumerRecordsInstrumentation());
+        new ConsumerRecordsInstrumentation(),
+        new ServiceCallExecutorServiceInstrumentation(),
+        new TrackingConsumerWrapperInstrumentation());
   }
 }

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

@@ -7,13 +7,8 @@
 
 import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11.KafkaSingletons.consumerReceiveInstrumenter;
-import static io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11.KafkaSingletons.enhanceConfig;
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 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 static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
@@ -23,10 +18,6 @@
 import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.time.Duration;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Properties;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -38,49 +29,15 @@ public class KafkaConsumerInstrumentation implements TypeInstrumentation {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    return named("org.apache.kafka.clients.consumer.KafkaConsumer");
+    return named("com.linkedin.kafka.clients.consumer.LiKafkaConsumerImpl");
   }
 
   @Override
   public void transform(TypeTransformer transformer) {
-    transformer.applyAdviceToMethod(
-        isConstructor().and(takesArgument(0, Map.class)),
-        this.getClass().getName() + "$ConstructorMapAdvice");
-    transformer.applyAdviceToMethod(
-        isConstructor().and(takesArgument(0, Properties.class)),
-        this.getClass().getName() + "$ConstructorPropertiesAdvice");
-    transformer.applyAdviceToMethod(
-        named("poll")
-            .and(isPublic())
-            .and(takesArguments(1))
-            .and(takesArgument(0, long.class).or(takesArgument(0, Duration.class)))
-            .and(returns(named("org.apache.kafka.clients.consumer.ConsumerRecords"))),
+    transformer.applyAdviceToMethod(named("poll").and(isPublic()),
         this.getClass().getName() + "$PollAdvice");
   }
 
-  @SuppressWarnings("unused")
-  public static class ConstructorMapAdvice {
-
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(value = 0, readOnly = false) Map<String, Object> config) {
-      // ensure config is a mutable map
-      if (config.getClass() != HashMap.class) {
-        config = new HashMap<>(config);
-      }
-      enhanceConfig(config);
-    }
-  }
-
-  @SuppressWarnings("unused")
-  public static class ConstructorPropertiesAdvice {
-
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(@Advice.Argument(0) Properties config) {
-      enhanceConfig(config);
-    }
-  }
-
   @SuppressWarnings("unused")
   public static class PollAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)

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

@@ -0,0 +1,44 @@
+package io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+
+import java.util.concurrent.ExecutorService;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+import net.bytebuddy.matcher.ElementMatcher;
+
+
+public class ServiceCallExecutorServiceInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.linkedin.container.servicecall.ServiceCallExecutorService");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod().and(isConstructor()),
+        ServiceCallExecutorServiceInstrumentation.class.getName() + "$ConstructorAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class ConstructorAdvice {
+    @Advice.OnMethodExit
+    public static void wrap(
+        @Advice.FieldValue(value = "_threadType") Object threadType,
+        @Advice.This(
+            typing = Assigner.Typing.DYNAMIC, readOnly = false) ExecutorService executorService) {
+      if (((Enum) threadType).name().equals("QUEUE")) {
+        executorService = Context.taskWrapping(executorService);
+      }
+    }
+  }
+}

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

@@ -21,6 +21,7 @@ public class TracingIterator<K, V> implements Iterator<ConsumerRecord<K, V>> {
   private final Iterator<ConsumerRecord<K, V>> delegateIterator;
   private final Context parentContext;
   private final KafkaConsumerContext consumerContext;
+  private final boolean endTraceImmediately;
 
   /*
    * Note: this may potentially create problems if this iterator is used from different threads. But
@@ -31,19 +32,28 @@ public class TracingIterator<K, V> implements Iterator<ConsumerRecord<K, V>> {
   @Nullable private Scope currentScope;
 
   private TracingIterator(
-      Iterator<ConsumerRecord<K, V>> delegateIterator, KafkaConsumerContext consumerContext) {
+      Iterator<ConsumerRecord<K, V>> delegateIterator, KafkaConsumerContext consumerContext,
+      boolean endTraceImmediately) {
     this.delegateIterator = delegateIterator;
 
     Context receiveContext = consumerContext.getContext();
     // use the receive CONSUMER as parent if it's available
     this.parentContext = receiveContext != null ? receiveContext : Context.current();
     this.consumerContext = consumerContext;
+    this.endTraceImmediately = endTraceImmediately;
   }
 
-  public static <K, V> Iterator<ConsumerRecord<K, V>> wrap(
+
+  public static <K,V> Iterator<ConsumerRecord<K, V>> wrap(
       Iterator<ConsumerRecord<K, V>> delegateIterator, KafkaConsumerContext consumerContext) {
+    return wrap(delegateIterator, consumerContext, false);
+  }
+
+  public static <K, V> Iterator<ConsumerRecord<K, V>> wrap(
+      Iterator<ConsumerRecord<K, V>> delegateIterator, KafkaConsumerContext consumerContext,
+      boolean endTraceImmediately) {
     if (KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
-      return new TracingIterator<>(delegateIterator, consumerContext);
+      return new TracingIterator<>(delegateIterator, consumerContext, endTraceImmediately);
     }
     return delegateIterator;
   }
@@ -70,6 +80,9 @@ public ConsumerRecord<K, V> next() {
       currentContext = consumerProcessInstrumenter().start(parentContext, currentRequest);
       currentScope = currentContext.makeCurrent();
     }
+    if (endTraceImmediately) {
+      closeScopeAndEndSpan();
+    }
     return next;
   }