unit tests for records(TopicPartition) method of kafka-clients instrumentation. introduce safeguards when instrumenting iterator() so that duplication of traces is not allowed. do not allow subList() to be instrumented
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
@@ -7,6 +7,7 @@ public class TracingIterable implements Iterable<ConsumerRecord> {
private final Iterable<ConsumerRecord> delegate;
private final String operationName;
private final KafkaDecorator decorator;
+ private boolean firstIterator = true;
public TracingIterable(
final Iterable<ConsumerRecord> delegate,
@@ -19,6 +20,17 @@ public class TracingIterable implements Iterable<ConsumerRecord> {
@Override
public Iterator<ConsumerRecord> iterator() {
- return new TracingIterator(delegate.iterator(), operationName, decorator);
+ Iterator<ConsumerRecord> it;
+ // We should only return one iterator with tracing.
+ // However, this is not thread-safe, but usually the first (hopefully only) traversal of
+ // ConsumerRecords is performed in the same thread that called poll()
+ if (this.firstIterator) {
+ it = new TracingIterator(delegate.iterator(), operationName, decorator);
+ firstIterator = false;
+ } else {
+ it = delegate.iterator();
+ }
+
+ return it;
}
}

MODIFY
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingList.java
dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingList.java
@@ -6,18 +6,15 @@ import java.util.List;
import java.util.ListIterator;
import org.apache.kafka.clients.consumer.ConsumerRecord;
-public class TracingList implements List<ConsumerRecord> {
+public class TracingList extends TracingIterable 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) {
+ super(delegate, operationName, decorator);
this.delegate = delegate;
- this.operationName = operationName;
- this.decorator = decorator;
}
@Override
@@ -35,11 +32,6 @@ public class TracingList implements List<ConsumerRecord> {
return delegate.contains(o);
}
- @Override
- public Iterator<ConsumerRecord> iterator() {
- return new TracingIterator(delegate.iterator(), operationName, decorator);
- }
-
@Override
public Object[] toArray() {
return delegate.toArray();
@@ -121,6 +113,11 @@ public class TracingList implements List<ConsumerRecord> {
return delegate.lastIndexOf(o);
}
+ @Override
+ public Iterator<ConsumerRecord> iterator() {
+ return super.iterator();
+ }
+
@Override
public ListIterator<ConsumerRecord> listIterator() {
// TODO: the API for ListIterator is not really good to instrument it in context of Kafka
@@ -137,6 +134,10 @@ public class TracingList implements List<ConsumerRecord> {
@Override
public List<ConsumerRecord> subList(final int fromIndex, final int toIndex) {
- return new TracingList(delegate.subList(fromIndex, toIndex), operationName, decorator);
+ // TODO: the API for subList is not really good to instrument it in context of Kafka
+ // Consumer so we will not do that for now
+ // Kafka is essentially a sequential commit log. We should only enable tracing when traversing
+ // sequentially with an iterator
+ return delegate.subList(fromIndex, toIndex);
}
}

MODIFY
dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/KafkaClientTest.groovy
dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/KafkaClientTest.groovy
import datadog.trace.agent.test.AgentTestRunner
+import org.apache.kafka.clients.consumer.ConsumerConfig
import org.apache.kafka.clients.consumer.ConsumerRecord
+import org.apache.kafka.clients.consumer.KafkaConsumer
+import org.apache.kafka.clients.producer.KafkaProducer
+import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.kafka.common.TopicPartition
import org.junit.ClassRule
import org.springframework.kafka.core.DefaultKafkaConsumerFactory
import org.springframework.kafka.core.DefaultKafkaProducerFactory
@@ -119,6 +124,94 @@ class KafkaClientTest extends AgentTestRunner {
cleanup:
producerFactory.stop()
container?.stop()
+ embeddedKafka.after()
+ }
+
+ def "test records(TopicPartition) kafka consume"() {
+ setup:
+ embeddedKafka.before()
+
+ // set up the Kafka consumer properties
+ def kafkaPartition = 0
+ def consumerProperties = KafkaTestUtils.consumerProps("sender", "false", embeddedKafka)
+ consumerProperties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+ def consumer = new KafkaConsumer<String,String>(consumerProperties)
+
+ def senderProps = KafkaTestUtils.senderProps(embeddedKafka.getBrokersAsString())
+ def producer = new KafkaProducer(senderProps)
+
+ consumer.assign(Arrays.asList(new TopicPartition(SHARED_TOPIC, kafkaPartition)))
+
+ when:
+ def greeting = "Hello from MockConsumer!"
+ producer.send(new ProducerRecord<Integer, String>(SHARED_TOPIC, kafkaPartition, null, greeting))
+ TEST_WRITER.waitForTraces(1)
+
+ then:
+
+ def records = new LinkedBlockingQueue<ConsumerRecord<String, String>>()
+ def pollResult = KafkaTestUtils.getRecords(consumer)
+
+ def recs = pollResult.records(new TopicPartition(SHARED_TOPIC, kafkaPartition)).iterator()
+
+ def isFirst = true
+ while (recs.hasNext()) {
+ if(!isFirst) {
+ TEST_WRITER.waitForTraces(1) // ensure consistent ordering of traces
+ }
+ records.add(recs.next())
+ isFirst = false
+ }
+ TEST_WRITER.waitForTraces(1) // ensure consistent ordering of traces
+
+ then:
+ def first = records.poll(5, TimeUnit.SECONDS)
+ first.value() == greeting
+ first.key() == null
+
+ assertTraces(2) {
+ trace(0, 1) {
+ // PRODUCER span 0
+ span(0) {
+ serviceName "kafka"
+ operationName "kafka.produce"
+ resourceName "Produce Topic $SHARED_TOPIC"
+ spanType "queue"
+ errored false
+ parent()
+ tags {
+ "component" "java-kafka"
+ "span.kind" "producer"
+ "kafka.partition" { it >= 0 }
+ defaultTags(true)
+ }
+ }
+ }
+ trace(1, 1) {
+ // CONSUMER span 0
+ span(0) {
+ serviceName "kafka"
+ operationName "kafka.consume"
+ resourceName "Consume Topic $SHARED_TOPIC"
+ spanType "queue"
+ errored false
+ childOf TEST_WRITER[0][0]
+ tags {
+ "component" "java-kafka"
+ "span.kind" "consumer"
+ "partition" { it >= 0 }
+ "offset" 0
+ defaultTags(true)
+ }
+ }
+ }
+ }
+
+ cleanup:
+ consumer.close()
+ producer.close()
+ embeddedKafka.after()
+
}
}

unit tests for records(TopicPartition) method of kafka-clients instrumentation. introduce safeguards when instrumenting iterator() so that duplication of traces is not allowed. do not allow subList() to be instrumented