Skip to content

Commit 905f423

Browse files
piochelepiotrpiochelepiotr
committed
data streams: Add Tracking of Kafka commit and produce calls
1 parent b4d2c64 commit 905f423

16 files changed

Lines changed: 383 additions & 13 deletions

File tree

‎dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/ConsumerCoordinatorInstrumentation.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
package datadog.trace.instrumentation.kafka_clients;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.*;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import datadog.trace.bootstrap.InstrumentationContext;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
10+
import java.util.HashMap;
11+
import java.util.Map;
12+
import net.bytebuddy.asm.Advice;
13+
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
14+
import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
15+
import org.apache.kafka.common.TopicPartition;
16+
17+
/** This instrumentation saves the co */
18+
@AutoService(Instrumenter.class)
19+
public final class ConsumerCoordinatorInstrumentation extends Instrumenter.Tracing
20+
implements Instrumenter.ForSingleType {
21+
22+
public ConsumerCoordinatorInstrumentation() {
23+
super("kafka");
24+
}
25+
26+
@Override
27+
public Map<String, String> contextStore() {
28+
Map<String, String> contextStores = new HashMap<>();
29+
// for each ConsumerCoordinator, we store the consumer group
30+
contextStores.put(
31+
"org.apache.kafka.clients.consumer.internals.ConsumerCoordinator", "java.lang.String");
32+
return contextStores;
33+
}
34+
35+
@Override
36+
public String instrumentedType() {
37+
return "org.apache.kafka.clients.consumer.internals.ConsumerCoordinator";
38+
}
39+
40+
@Override
41+
public String[] helperClassNames() {
42+
return new String[] {};
43+
}
44+
45+
@Override
46+
public void adviceTransformations(AdviceTransformation transformation) {
47+
transformation.applyAdvice(
48+
isMethod().and(named("sendOffsetCommitRequest")).and(takesArguments(1)),
49+
ConsumerCoordinatorInstrumentation.class.getName() + "$CommitOffsetAdvice");
50+
}
51+
52+
public static class CommitOffsetAdvice {
53+
@Advice.OnMethodEnter(suppress = Throwable.class)
54+
public static void trackCommitOffset(
55+
@Advice.This ConsumerCoordinator coordinator,
56+
@Advice.Argument(0) final Map<TopicPartition, OffsetAndMetadata> offsets) {
57+
String consumerGroup =
58+
InstrumentationContext.get(ConsumerCoordinator.class, String.class).get(coordinator);
59+
System.out.printf("___________________________________\n");
60+
System.out.printf("committing offsets %s\n", consumerGroup);
61+
for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : offsets.entrySet()) {
62+
AgentTracer.get()
63+
.trackKafkaCommit(
64+
consumerGroup,
65+
entry.getKey().topic(),
66+
entry.getKey().partition(),
67+
entry.getValue().offset());
68+
System.out.printf(
69+
"partition %d, topic %s, offset %d\n",
70+
entry.getKey().partition(), entry.getKey().topic(), entry.getValue().offset());
71+
}
72+
}
73+
}
74+
}

‎dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerGroupInstrumentation.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.apache.kafka.clients.consumer.ConsumerRecord;
1919
import org.apache.kafka.clients.consumer.ConsumerRecords;
2020
import org.apache.kafka.clients.consumer.KafkaConsumer;
21+
import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
2122

2223
/** This instrumentation saves the co */
2324
@AutoService(Instrumenter.class)
@@ -33,7 +34,8 @@ public Map<String, String> contextStore() {
3334
Map<String, String> contextStores = new HashMap<>();
3435
contextStores.put("org.apache.kafka.clients.consumer.KafkaConsumer", "java.lang.String");
3536
contextStores.put("org.apache.kafka.clients.consumer.ConsumerRecords", "java.lang.String");
36-
37+
contextStores.put(
38+
"org.apache.kafka.clients.consumer.internals.ConsumerCoordinator", "java.lang.String");
3739
return contextStores;
3840
}
3941

@@ -70,11 +72,16 @@ public void adviceTransformations(AdviceTransformation transformation) {
7072
public static class ConstructorAdvice {
7173
@Advice.OnMethodExit(suppress = Throwable.class)
7274
public static void captureGroup(
73-
@Advice.This KafkaConsumer consumer, @Advice.Argument(0) ConsumerConfig consumerConfig) {
75+
@Advice.This KafkaConsumer consumer,
76+
@Advice.FieldValue("coordinator") ConsumerCoordinator coordinator,
77+
@Advice.Argument(0) ConsumerConfig consumerConfig) {
7478
String consumerGroup = consumerConfig.getString(ConsumerConfig.GROUP_ID_CONFIG);
79+
System.out.printf("getting consumer group %s\n", consumerGroup);
7580

7681
if (consumerGroup != null && !consumerGroup.isEmpty()) {
7782
InstrumentationContext.get(KafkaConsumer.class, String.class).put(consumer, consumerGroup);
83+
InstrumentationContext.get(ConsumerCoordinator.class, String.class)
84+
.put(coordinator, consumerGroup);
7885
}
7986
}
8087

@@ -85,6 +92,9 @@ public static void muzzleCheck(ConsumerRecord record) {
8592
}
8693
}
8794

95+
// this method transfers the consumer group from the KafkaConsumer class key to the
96+
// ConsumerRecords key. This is because in the poll method, we actually don't have access to the
97+
// KafkaConsumer class.
8898
public static class RecordsAdvice {
8999
@Advice.OnMethodExit(suppress = Throwable.class)
90100
public static void captureGroup(

‎dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaDecorator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ public void onConsume(final AgentSpan span, final ConsumerRecord record, String
9292
final String topic = record.topic() == null ? "kafka" : record.topic();
9393
span.setResourceName(CONSUMER_RESOURCE_NAME_CACHE.computeIfAbsent(topic, CONSUMER_PREFIX));
9494
span.setTag(PARTITION, record.partition());
95+
span.setTag("partition2", record.partition());
9596
span.setTag(OFFSET, record.offset());
9697
if (consumerGroup != null) {
9798
span.setTag(CONSUMER_GROUP, consumerGroup);

‎dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerCallback.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,17 @@
55

66
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
77
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
8+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
89
import org.apache.kafka.clients.producer.Callback;
910
import org.apache.kafka.clients.producer.RecordMetadata;
11+
import org.slf4j.Logger;
12+
import org.slf4j.LoggerFactory;
1013

1114
public class KafkaProducerCallback implements Callback {
1215
private final Callback callback;
1316
private final AgentSpan parent;
1417
private final AgentSpan span;
18+
private static final Logger log = LoggerFactory.getLogger(KafkaProducerCallback.class);
1519

1620
public KafkaProducerCallback(
1721
final Callback callback, final AgentSpan parent, final AgentSpan span) {
@@ -24,6 +28,11 @@ public KafkaProducerCallback(
2428
public void onCompletion(final RecordMetadata metadata, final Exception exception) {
2529
PRODUCER_DECORATE.onError(span, exception);
2630
PRODUCER_DECORATE.beforeFinish(span);
31+
AgentTracer.get().trackKafkaProduce(metadata.topic(), metadata.partition(), metadata.offset());
32+
span.setTag("partition", metadata.partition());
33+
span.setTag("topic", metadata.topic());
34+
span.setTag("offset", metadata.offset());
35+
log.warn("Produced a message.");
2736
span.finish();
2837
if (callback != null) {
2938
if (parent != null) {

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -850,6 +850,16 @@ public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {
850850
LambdaHandler.notifyEndInvocation(span, result, isError);
851851
}
852852

853+
@Override
854+
public void trackKafkaProduce(String topic, int partition, long offset) {
855+
dataStreamsCheckpointer.trackKafkaProduce(topic, partition, offset);
856+
}
857+
858+
@Override
859+
public void trackKafkaCommit(String consumerGroup, String topic, int partition, long offset) {
860+
dataStreamsCheckpointer.trackKafkaCommit(consumerGroup, topic, partition, offset);
861+
}
862+
853863
private final RatelimitedLogger rlLog = new RatelimitedLogger(log, 1, MINUTES);
854864

855865
/**

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DataStreamsCheckpointer.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@ <C> PathwayContext extractBinaryPathwayContext(
1515

1616
<C> PathwayContext extractPathwayContext(C carrier, AgentPropagation.ContextVisitor<C> getter);
1717

18+
void trackKafkaProduce(String topic, int partition, long offset);
19+
20+
void trackKafkaCommit(String consumerGroup, String topic, int partition, long offset);
21+
1822
@Override
1923
void close();
2024

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsCheckpointer.java‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,7 @@
1010
import datadog.trace.api.Config;
1111
import datadog.trace.api.WellKnownTags;
1212
import datadog.trace.api.time.TimeSource;
13-
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
14-
import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
15-
import datadog.trace.bootstrap.instrumentation.api.StatsPoint;
13+
import datadog.trace.bootstrap.instrumentation.api.*;
1614
import datadog.trace.common.metrics.EventListener;
1715
import datadog.trace.common.metrics.OkHttpSink;
1816
import datadog.trace.common.metrics.Sink;
@@ -43,7 +41,7 @@ public class DefaultDataStreamsCheckpointer
4341
new StatsPoint(Collections.<String>emptyList(), 0, 0, 0, 0, 0);
4442

4543
private final Map<Long, StatsBucket> timeToBucket = new HashMap<>();
46-
private final BlockingQueue<StatsPoint> inbox = new MpscBlockingConsumerArrayQueue<>(1024);
44+
private final BlockingQueue<StatsPayload> inbox = new MpscBlockingConsumerArrayQueue<>(1024);
4745
private final DatastreamsPayloadWriter payloadWriter;
4846
private final DDAgentFeaturesDiscovery features;
4947
private final TimeSource timeSource;
@@ -143,6 +141,17 @@ public <C> PathwayContext extractPathwayContext(
143141
return DefaultPathwayContext.extract(carrier, getter, timeSource, wellKnownTags);
144142
}
145143

144+
public void trackKafkaProduce(String topic, int partition, long offset) {
145+
System.out.println("tracking produce");
146+
inbox.offer(new KafkaOffset(topic, partition, offset, timeSource.getCurrentTimeNanos()));
147+
}
148+
149+
public void trackKafkaCommit(String consumerGroup, String topic, int partition, long offset) {
150+
System.out.println("tracking commit");
151+
inbox.offer(
152+
new KafkaOffset(consumerGroup, topic, partition, offset, timeSource.getCurrentTimeNanos()));
153+
}
154+
146155
@Override
147156
public void close() {
148157
if (null != cancellation) {
@@ -163,25 +172,30 @@ public void run() {
163172
Thread currentThread = Thread.currentThread();
164173
while (!currentThread.isInterrupted()) {
165174
try {
166-
StatsPoint statsPoint = inbox.take();
175+
StatsPayload payload = inbox.take();
167176

168-
if (statsPoint == REPORT) {
177+
if (payload == REPORT) {
169178
if (supportsDataStreams) {
170179
flush(timeSource.getCurrentTimeNanos());
171180
} else if (timeSource.getCurrentTimeNanos() >= nextFeatureCheck) {
172181
checkFeatures();
173182
}
174-
} else if (statsPoint == POISON_PILL) {
183+
} else if (payload == POISON_PILL) {
175184
if (supportsDataStreams) {
176185
flush(Long.MAX_VALUE);
177186
}
178187
break;
179188
} else if (supportsDataStreams) {
180-
Long bucket = currentBucket(statsPoint.getTimestampNanos());
189+
Long bucket = currentBucket(payload.getTimestampNanos());
181190
StatsBucket statsBucket =
182191
timeToBucket.computeIfAbsent(
183192
bucket, startTime -> new StatsBucket(startTime, bucketDurationNanos));
184-
statsBucket.addPoint(statsPoint);
193+
if (payload.type() == StatsPayload.Type.StatsGroup) {
194+
StatsPoint statsPoint = (StatsPoint) payload;
195+
statsBucket.addPoint(statsPoint);
196+
} else {
197+
statsBucket.addOffset((KafkaOffset) payload);
198+
}
185199
}
186200
} catch (InterruptedException e) {
187201
currentThread.interrupt();

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import datadog.trace.api.WellKnownTags;
1010
import datadog.trace.common.metrics.Sink;
1111
import java.util.Collection;
12+
import java.util.Map;
1213

1314
public class MsgPackDatastreamsPayloadWriter implements DatastreamsPayloadWriter {
1415
private static final byte[] ENV = "Env".getBytes(ISO_8859_1);
@@ -22,8 +23,15 @@ public class MsgPackDatastreamsPayloadWriter implements DatastreamsPayloadWriter
2223
private static final byte[] EDGE_LATENCY = "EdgeLatency".getBytes(ISO_8859_1);
2324
private static final byte[] SERVICE = "Service".getBytes(ISO_8859_1);
2425
private static final byte[] EDGE_TAGS = "EdgeTags".getBytes(ISO_8859_1);
26+
private static final byte[] KAFKA = "Kafka".getBytes(ISO_8859_1);
27+
private static final byte[] LATEST_PRODUCE_OFFSETS = "LatestProduceOffsets".getBytes(ISO_8859_1);
28+
private static final byte[] LATEST_COMMIT_OFFSETS = "LatestCommitOffsets".getBytes(ISO_8859_1);
2529
private static final byte[] HASH = "Hash".getBytes(ISO_8859_1);
2630
private static final byte[] PARENT_HASH = "ParentHash".getBytes(ISO_8859_1);
31+
private static final byte[] TOPIC = "Topic".getBytes(ISO_8859_1);
32+
private static final byte[] CONSUMER_GROUP = "ConsumerGroup".getBytes(ISO_8859_1);
33+
private static final byte[] PARTITION = "Partition".getBytes(ISO_8859_1);
34+
private static final byte[] OFFSET = "Offset".getBytes(ISO_8859_1);
2735

2836
private static final int INITIAL_CAPACITY = 512 * 1024;
2937

@@ -75,7 +83,10 @@ public void writePayload(Collection<StatsBucket> data) {
7583
writer.writeUTF8(STATS);
7684
writer.startArray(data.size());
7785
for (StatsBucket bucket : data) {
78-
writer.startMap(3);
86+
boolean hasKafka =
87+
!bucket.getLatestKafkaCommitOffsets().isEmpty()
88+
|| !bucket.getLatestKafkaProduceOffsets().isEmpty();
89+
writer.startMap(3 + (hasKafka ? 1 : 0));
7990

8091
/* 1 */
8192
writer.writeUTF8(START);
@@ -88,6 +99,12 @@ public void writePayload(Collection<StatsBucket> data) {
8899
/* 3 */
89100
writer.writeUTF8(STATS);
90101
writeBucket(bucket, writer);
102+
103+
if (hasKafka) {
104+
/* 4 */
105+
writeKafka(
106+
bucket.getLatestKafkaCommitOffsets(), bucket.getLatestKafkaProduceOffsets(), writer);
107+
}
91108
}
92109

93110
buffer.mark();
@@ -96,6 +113,7 @@ public void writePayload(Collection<StatsBucket> data) {
96113
}
97114

98115
private void writeBucket(StatsBucket bucket, Writable packer) {
116+
System.out.println("printing bucket");
99117
Collection<StatsGroup> groups = bucket.getGroups();
100118
packer.startArray(groups.size());
101119
for (StatsGroup group : groups) {
@@ -129,4 +147,45 @@ private void writeBucket(StatsBucket bucket, Writable packer) {
129147
}
130148
}
131149
}
150+
151+
private void writeKafka(
152+
Collection<Map.Entry<TopicPartitionGroup, Long>> latestCommitOffsets,
153+
Collection<Map.Entry<TopicPartition, Long>> latestProduceOffsets,
154+
Writable packer) {
155+
System.out.println("has kafka!");
156+
packer.writeUTF8(KAFKA);
157+
boolean hasCommitOffsets = !latestCommitOffsets.isEmpty();
158+
boolean hasProduceOffsets = !latestProduceOffsets.isEmpty();
159+
packer.startMap((hasCommitOffsets ? 1 : 0) + (hasProduceOffsets ? 1 : 0));
160+
if (hasCommitOffsets) {
161+
System.out.println("has commit");
162+
packer.writeUTF8(LATEST_COMMIT_OFFSETS);
163+
packer.startArray(latestCommitOffsets.size());
164+
for (Map.Entry<TopicPartitionGroup, Long> entry : latestCommitOffsets) {
165+
packer.startMap(4);
166+
packer.writeUTF8(CONSUMER_GROUP);
167+
packer.writeString(entry.getKey().getGroup(), null);
168+
packer.writeUTF8(TOPIC);
169+
packer.writeString(entry.getKey().getTopic(), null);
170+
packer.writeUTF8(PARTITION);
171+
packer.writeInt(entry.getKey().getPartition());
172+
packer.writeUTF8(OFFSET);
173+
packer.writeLong(entry.getValue());
174+
}
175+
}
176+
if (hasProduceOffsets) {
177+
System.out.println("has produce");
178+
packer.writeUTF8(LATEST_PRODUCE_OFFSETS);
179+
packer.startArray(latestProduceOffsets.size());
180+
for (Map.Entry<TopicPartition, Long> entry : latestProduceOffsets) {
181+
packer.startMap(3);
182+
packer.writeUTF8(TOPIC);
183+
packer.writeString(entry.getKey().getTopic(), null);
184+
packer.writeUTF8(PARTITION);
185+
packer.writeInt(entry.getKey().getPartition());
186+
packer.writeUTF8(OFFSET);
187+
packer.writeLong(entry.getValue());
188+
}
189+
}
190+
}
132191
}

0 commit comments

Comments
 (0)