Instrument records(TopicPartition) in kafka consumer
MODIFY
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInstrumentation.java
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInstrumentation.java
package datadog.trace.instrumentation.kafka_clients;
-import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.CONSUMER_DECORATE;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@ import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
import java.util.HashMap;
import java.util.Iterator;
+import java.util.List;
import java.util.Map;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
@@ -42,7 +42,8 @@ public final class KafkaConsumerInstrumentation extends Instrumenter.Default {
packageName + ".KafkaDecorator$2",
packageName + ".TextMapExtractAdapter",
packageName + ".TracingIterable",
- packageName + ".TracingIterable$TracingIterator",
+ packageName + ".TracingIterator",
+ packageName + ".TracingList",
};
}
@@ -53,11 +54,16 @@ public final class KafkaConsumerInstrumentation extends Instrumenter.Default {
isMethod()
.and(isPublic())
.and(named("records"))
- .and(
- takesArgument(0, String.class)
- .or(takesArgument(0, named("org.apache.kafka.common.TopicPartition"))))
- .and(returns(safeHasSuperType(named(Iterable.class.getName())))),
+ .and(takesArgument(0, String.class))
+ .and(returns(Iterable.class)),
IterableAdvice.class.getName());
+ transformers.put(
+ isMethod()
+ .and(isPublic())
+ .and(named("records"))
+ .and(takesArgument(0, named("org.apache.kafka.common.TopicPartition")))
+ .and(returns(List.class)),
+ ListAdvice.class.getName());
transformers.put(
isMethod()
.and(isPublic())
@@ -78,13 +84,22 @@ public final class KafkaConsumerInstrumentation extends Instrumenter.Default {
}
}
+ public static class ListAdvice {
+
+ @Advice.OnMethodExit(suppress = Throwable.class)
+ public static void wrap(@Advice.Return(readOnly = false) List<ConsumerRecord> iterable) {
+ if (iterable != null) {
+ iterable = new TracingList(iterable, "kafka.consume", CONSUMER_DECORATE);
+ }
+ }
+ }
+
public static class IteratorAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void wrap(@Advice.Return(readOnly = false) Iterator<ConsumerRecord> iterator) {
if (iterator != null) {
- iterator =
- new TracingIterable.TracingIterator(iterator, "kafka.consume", CONSUMER_DECORATE);
+ iterator = new TracingIterator(iterator, "kafka.consume", CONSUMER_DECORATE);
}
}
}

MODIFY
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterable.java
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterable.java
package datadog.trace.instrumentation.kafka_clients;
-import io.opentracing.Scope;
-import io.opentracing.SpanContext;
-import io.opentracing.propagation.Format;
-import io.opentracing.util.GlobalTracer;
import java.util.Iterator;
-import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerRecord;
public class TracingIterable implements Iterable<ConsumerRecord> {
- private final Iterable<ConsumerRecord> delegateIterable;
+ private final Iterable<ConsumerRecord> delegate;
private final String operationName;
private final KafkaDecorator decorator;
public TracingIterable(
- final Iterable<ConsumerRecord> delegateIterable,
+ final Iterable<ConsumerRecord> delegate,
final String operationName,
final KafkaDecorator decorator) {
- this.delegateIterable = delegateIterable;
+ this.delegate = delegate;
this.operationName = operationName;
this.decorator = decorator;
}
@Override
public Iterator<ConsumerRecord> iterator() {
- return new TracingIterator(delegateIterable.iterator(), operationName, decorator);
- }
-
- @Slf4j
- public static class TracingIterator implements Iterator<ConsumerRecord> {
- private final Iterator<ConsumerRecord> delegateIterator;
- private final String operationName;
- private final KafkaDecorator decorator;
-
- /**
- * Note: this may potentially create problems if this iterator is used from different threads.
- * But at the moment we cannot do much about this.
- */
- private Scope currentScope;
-
- public TracingIterator(
- final Iterator<ConsumerRecord> delegateIterator,
- final String operationName,
- final KafkaDecorator decorator) {
- this.delegateIterator = delegateIterator;
- this.operationName = operationName;
- this.decorator = decorator;
- }
-
- @Override
- public boolean hasNext() {
- if (currentScope != null) {
- currentScope.close();
- currentScope = null;
- }
- return delegateIterator.hasNext();
- }
-
- @Override
- public ConsumerRecord next() {
- if (currentScope != null) {
- // in case they didn't call hasNext()...
- currentScope.close();
- currentScope = null;
- }
-
- final ConsumerRecord next = delegateIterator.next();
-
- try {
- if (next != null) {
- final SpanContext spanContext =
- GlobalTracer.get()
- .extract(Format.Builtin.TEXT_MAP, new TextMapExtractAdapter(next.headers()));
- currentScope =
- GlobalTracer.get().buildSpan(operationName).asChildOf(spanContext).startActive(true);
- decorator.afterStart(currentScope);
- decorator.onConsume(currentScope, next);
- }
- } catch (final Exception e) {
- log.debug("Error during decoration", e);
- }
- return next;
- }
-
- @Override
- public void remove() {
- delegateIterator.remove();
- }
+ return new TracingIterator(delegate.iterator(), operationName, decorator);
}
}

ADD
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java
+package datadog.trace.instrumentation.kafka_clients;
+
+import io.opentracing.Scope;
+import io.opentracing.SpanContext;
+import io.opentracing.propagation.Format;
+import io.opentracing.util.GlobalTracer;
+import java.util.Iterator;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+@Slf4j
+public class TracingIterator implements Iterator<ConsumerRecord> {
+ private final Iterator<ConsumerRecord> delegateIterator;
+ private final String operationName;
+ private final KafkaDecorator decorator;
+
+ /**
+ * Note: this may potentially create problems if this iterator is used from different threads. But
+ * at the moment we cannot do much about this.
+ */
+ private Scope currentScope;
+
+ public TracingIterator(
+ final Iterator<ConsumerRecord> delegateIterator,
+ final String operationName,
+ final KafkaDecorator decorator) {
+ this.delegateIterator = delegateIterator;
+ this.operationName = operationName;
+ this.decorator = decorator;
+ }
+
+ @Override
+ public boolean hasNext() {
+ if (currentScope != null) {
+ currentScope.close();
+ currentScope = null;
+ }
+ return delegateIterator.hasNext();
+ }
+
+ @Override
+ public ConsumerRecord next() {
+ if (currentScope != null) {
+ // in case they didn't call hasNext()...
+ currentScope.close();
+ currentScope = null;
+ }
+
+ final ConsumerRecord next = delegateIterator.next();
+
+ try {
+ if (next != null) {
+ final SpanContext spanContext =
+ GlobalTracer.get()
+ .extract(Format.Builtin.TEXT_MAP, new TextMapExtractAdapter(next.headers()));
+ currentScope =
+ GlobalTracer.get().buildSpan(operationName).asChildOf(spanContext).startActive(true);
+ decorator.afterStart(currentScope);
+ decorator.onConsume(currentScope, next);
+ }
+ } catch (final Exception e) {
+ log.debug("Error during decoration", e);
+ }
+ return next;
+ }
+
+ @Override
+ public void remove() {
+ delegateIterator.remove();
+ }
+}

ADD
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingList.java
+package datadog.trace.instrumentation.kafka_clients;
+
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.ListIterator;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+public class TracingList implements List<ConsumerRecord> {
+ private final List<ConsumerRecord> delegate;
+ private final String operationName;
+ private final KafkaDecorator decorator;
+
+ public TracingList(
+ final List<ConsumerRecord> delegate,
+ final String operationName,
+ final KafkaDecorator decorator) {
+ this.delegate = delegate;
+ this.operationName = operationName;
+ this.decorator = decorator;
+ }
+
+ @Override
+ public int size() {
+ return delegate.size();
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return delegate.isEmpty();
+ }
+
+ @Override
+ public boolean contains(final Object o) {
+ return delegate.contains(o);
+ }
+
+ @Override
+ public Iterator<ConsumerRecord> iterator() {
+ return new TracingIterator(delegate.iterator(), operationName, decorator);
+ }
+
+ @Override
+ public Object[] toArray() {
+ return delegate.toArray();
+ }
+
+ @Override
+ public <T> T[] toArray(final T[] a) {
+ return delegate.toArray(a);
+ }
+
+ @Override
+ public boolean add(final ConsumerRecord consumerRecord) {
+ return delegate.add(consumerRecord);
+ }
+
+ @Override
+ public boolean remove(final Object o) {
+ return delegate.remove(o);
+ }
+
+ @Override
+ public boolean containsAll(final Collection<?> c) {
+ return delegate.containsAll(c);
+ }
+
+ @Override
+ public boolean addAll(final Collection<? extends ConsumerRecord> c) {
+ return delegate.addAll(c);
+ }
+
+ @Override
+ public boolean addAll(final int index, final Collection<? extends ConsumerRecord> c) {
+ return delegate.addAll(index, c);
+ }
+
+ @Override
+ public boolean removeAll(final Collection<?> c) {
+ return delegate.removeAll(c);
+ }
+
+ @Override
+ public boolean retainAll(final Collection<?> c) {
+ return delegate.retainAll(c);
+ }
+
+ @Override
+ public void clear() {
+ delegate.clear();
+ }
+
+ @Override
+ public ConsumerRecord get(final int index) {
+ // TODO: should this be instrumented as well?
+ return delegate.get(index);
+ }
+
+ @Override
+ public ConsumerRecord set(final int index, final ConsumerRecord element) {
+ return delegate.set(index, element);
+ }
+
+ @Override
+ public void add(final int index, final ConsumerRecord element) {
+ delegate.add(index, element);
+ }
+
+ @Override
+ public ConsumerRecord remove(final int index) {
+ return delegate.remove(index);
+ }
+
+ @Override
+ public int indexOf(final Object o) {
+ return delegate.indexOf(o);
+ }
+
+ @Override
+ public int lastIndexOf(final Object o) {
+ return delegate.lastIndexOf(o);
+ }
+
+ @Override
+ public ListIterator<ConsumerRecord> listIterator() {
+ // TODO: should this be instrumented as well?
+ return delegate.listIterator();
+ }
+
+ @Override
+ public ListIterator<ConsumerRecord> listIterator(final int index) {
+ // TODO: should this be instrumented as well?
+ return delegate.listIterator(index);
+ }
+
+ @Override
+ public List<ConsumerRecord> subList(final int fromIndex, final int toIndex) {
+ return new TracingList(delegate.subList(fromIndex, toIndex), operationName, decorator);
+ }
+}

Instrument records(TopicPartition) in kafka consumer