-
Notifications
You must be signed in to change notification settings - Fork 281
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add Kafka poll span when DSM is enabled #6969
Changes from 2 commits
24f4532
eaba756
d50befb
46afc8a
1252812
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,10 @@ | ||
package datadog.trace.instrumentation.kafka_clients; | ||
|
||
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; | ||
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; | ||
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; | ||
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.KAFKA_RECORDS_COUNT; | ||
import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.KAFKA_POLL; | ||
import static net.bytebuddy.matcher.ElementMatchers.isConstructor; | ||
import static net.bytebuddy.matcher.ElementMatchers.isMethod; | ||
import static net.bytebuddy.matcher.ElementMatchers.isPublic; | ||
|
@@ -13,6 +17,8 @@ | |
import datadog.trace.agent.tooling.InstrumenterModule; | ||
import datadog.trace.api.Config; | ||
import datadog.trace.bootstrap.InstrumentationContext; | ||
import datadog.trace.bootstrap.instrumentation.api.AgentScope; | ||
import datadog.trace.bootstrap.instrumentation.api.AgentSpan; | ||
import java.util.HashMap; | ||
import java.util.List; | ||
import java.util.Map; | ||
|
@@ -130,10 +136,28 @@ public static void muzzleCheck(ConsumerRecord record) { | |
* KafkaConsumer class. | ||
*/ | ||
public static class RecordsAdvice { | ||
@Advice.OnMethodEnter(suppress = Throwable.class) | ||
public static AgentScope onEnter() { | ||
if (Config.get().isDataStreamsEnabled()) { | ||
final AgentSpan span = startSpan(KAFKA_POLL); | ||
return activateSpan(span); | ||
} | ||
return null; | ||
} | ||
|
||
@Advice.OnMethodExit(suppress = Throwable.class) | ||
public static void captureGroup( | ||
@Advice.This KafkaConsumer consumer, @Advice.Return ConsumerRecords records) { | ||
@Advice.Enter final AgentScope scope, | ||
@Advice.This KafkaConsumer consumer, | ||
@Advice.Return ConsumerRecords records) { | ||
if (records == null) { | ||
if (scope != null) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This feels a bit messy to me. On first read, it looked like we were closing scope & spans twice. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I updated the code, please let me know what you think of the new version. |
||
AgentSpan span = scope.span(); | ||
if (span != null) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is this scope span null check needed at all? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You're right, is scope is not null, there must be a span (as far as I understand) |
||
span.finish(); | ||
} | ||
scope.close(); | ||
} | ||
return; | ||
} | ||
KafkaConsumerInfo kafkaConsumerInfo = | ||
|
@@ -142,6 +166,15 @@ public static void captureGroup( | |
InstrumentationContext.get(ConsumerRecords.class, KafkaConsumerInfo.class) | ||
.put(records, kafkaConsumerInfo); | ||
} | ||
if (scope == null) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Maybe do this scope null check once at the beginning of the method? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I updated a bit the code, the hard part is that I want to attach the kafka consumer info to records even if the scope is null. |
||
return; | ||
} | ||
AgentSpan span = scope.span(); | ||
if (span != null) { | ||
span.setTag(KAFKA_RECORDS_COUNT, records.count()); | ||
span.finish(); | ||
} | ||
scope.close(); | ||
} | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,5 @@ | ||
import datadog.trace.common.writer.ListWriter | ||
|
||
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan | ||
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace | ||
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope | ||
|
@@ -53,13 +55,57 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase { | |
|
||
public static final LinkedHashMap<String, String> PRODUCER_PATHWAY_EDGE_TAGS | ||
|
||
// filter out Kafka poll, since the function is called in a loop, giving inconsistent results | ||
final ListWriter.Filter dropKafkaPoll = new ListWriter.Filter() { | ||
@Override | ||
boolean accept(List<DDSpan> trace) { | ||
return !(trace.size() == 1 && | ||
trace.get(0).getResourceName().toString().equals("kafka.poll")) | ||
} | ||
} | ||
|
||
final ListWriter.Filter dropEmptyKafkaPoll = new ListWriter.Filter() { | ||
@Override | ||
boolean accept(List<DDSpan> trace) { | ||
return !(trace.size() == 1 && | ||
trace.get(0).getResourceName().toString().equals("kafka.poll") && | ||
trace.get(0).getTag(InstrumentationTags.KAFKA_RECORDS_COUNT).equals(0)) | ||
} | ||
} | ||
|
||
// TraceID, start times & names changed based on the configuration, so overriding the sort to give consistent test results | ||
private static class SortKafkaTraces implements Comparator<List<DDSpan>> { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Are the usual test sorting mechanisms not sufficient? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Unfortunately, it doesn't look like it 😢 |
||
@Override | ||
int compare(List<DDSpan> o1, List<DDSpan> o2) { | ||
return rootSpanTrace(o1) - rootSpanTrace(o2) | ||
} | ||
|
||
int rootSpanTrace(List<DDSpan> trace) { | ||
assert !trace.isEmpty() | ||
def rootSpan = trace.get(0).localRootSpan | ||
switch (rootSpan.operationName.toString()) { | ||
case "parent": | ||
return 3 | ||
case "kafka.poll": | ||
return 2 | ||
default: | ||
return 1 | ||
} | ||
} | ||
} | ||
|
||
|
||
static { | ||
PRODUCER_PATHWAY_EDGE_TAGS = new LinkedHashMap<>(3) | ||
PRODUCER_PATHWAY_EDGE_TAGS.put("direction", "out") | ||
PRODUCER_PATHWAY_EDGE_TAGS.put("topic", SHARED_TOPIC) | ||
PRODUCER_PATHWAY_EDGE_TAGS.put("type", "kafka") | ||
} | ||
|
||
def setup() { | ||
TEST_WRITER.setFilter(dropKafkaPoll) | ||
} | ||
|
||
@Override | ||
int version() { | ||
0 | ||
|
@@ -124,9 +170,6 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase { | |
} | ||
} | ||
|
||
if (isDataStreamsEnabled()) { | ||
} | ||
|
||
cleanup: | ||
producer.close() | ||
} | ||
|
@@ -137,6 +180,7 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase { | |
if (isDataStreamsEnabled()) { | ||
senderProps.put(ProducerConfig.METADATA_MAX_AGE_CONFIG, 1000) | ||
} | ||
TEST_WRITER.setFilter(dropEmptyKafkaPoll) | ||
KafkaProducer<String, String> producer = new KafkaProducer<>(senderProps, new StringSerializer(), new StringSerializer()) | ||
String clusterId = "" | ||
if (isDataStreamsEnabled()) { | ||
|
@@ -203,28 +247,37 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase { | |
received.value() == greeting | ||
received.key() == null | ||
|
||
assertTraces(2, SORT_TRACES_BY_ID) { | ||
trace(3) { | ||
basicSpan(it, "parent") | ||
basicSpan(it, "producer callback", span(0)) | ||
producerSpan(it, senderProps, span(0), false) | ||
} | ||
int nTraces = isDataStreamsEnabled() ? 3 : 2 | ||
int produceTraceIdx = nTraces - 1 | ||
TEST_WRITER.waitForTraces(nTraces) | ||
def traces = (Arrays.asList(TEST_WRITER.toArray()) as List<List<DDSpan>>) | ||
Collections.sort(traces, new SortKafkaTraces()) | ||
assertTraces(nTraces, new SortKafkaTraces()) { | ||
if (hasQueueSpan()) { | ||
trace(2) { | ||
consumerSpan(it, consumerProperties, trace(1)[1]) | ||
queueSpan(it, trace(0)[2]) | ||
consumerSpan(it, consumerProperties, span(1)) | ||
queueSpan(it, trace(produceTraceIdx)[2]) | ||
} | ||
} else { | ||
trace(1) { | ||
consumerSpan(it, consumerProperties, trace(0)[2]) | ||
consumerSpan(it, consumerProperties, trace(produceTraceIdx)[2]) | ||
} | ||
} | ||
if (isDataStreamsEnabled()) { | ||
trace(1, { | ||
pollSpan(it) | ||
}) | ||
} | ||
trace(3) { | ||
basicSpan(it, "parent") | ||
basicSpan(it, "producer callback", span(0)) | ||
producerSpan(it, senderProps, span(0), false) | ||
} | ||
} | ||
|
||
def headers = received.headers() | ||
headers.iterator().hasNext() | ||
new String(headers.headers("x-datadog-trace-id").iterator().next().value()) == "${TEST_WRITER[0][2].traceId}" | ||
new String(headers.headers("x-datadog-parent-id").iterator().next().value()) == "${TEST_WRITER[0][2].spanId}" | ||
new String(headers.headers("x-datadog-trace-id").iterator().next().value()) == "${traces[produceTraceIdx][2].traceId}" | ||
new String(headers.headers("x-datadog-parent-id").iterator().next().value()) == "${traces[produceTraceIdx][2].spanId}" | ||
|
||
if (isDataStreamsEnabled()) { | ||
StatsGroup first = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == 0 } | ||
|
@@ -1069,6 +1122,27 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase { | |
} | ||
} | ||
|
||
def pollSpan( | ||
TraceAssert trace, | ||
int recordCount = 1, | ||
DDSpan parentSpan = null, | ||
Range offset = 0..0, | ||
boolean tombstone = false, | ||
boolean distributedRootSpan = !hasQueueSpan() | ||
) { | ||
trace.span { | ||
serviceName Config.get().getServiceName() | ||
operationName "kafka.poll" | ||
resourceName "kafka.poll" | ||
errored false | ||
measured false | ||
tags { | ||
"$InstrumentationTags.KAFKA_RECORDS_COUNT" recordCount | ||
defaultTags(true) | ||
} | ||
} | ||
} | ||
|
||
def waitForKafkaMetadataUpdate(KafkaTemplate kafkaTemplate) { | ||
kafkaTemplate.flush() | ||
Producer<String, String> wrappedProducer = kafkaTemplate.getTheProducer() | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Extra spans have the potential to create noisy traces.
But I guess this is okay, since it is under "data streams enabled".
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You can also try to see if you have an active span and take the dsm config from it since using
Config.get()
won't let you access to the dynamic configThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sounds good, I updated the code.
Note that this codepath will today probably not trigger, since the reason I'm adding this span, is because there is no active span to attach the schema on.