Enable Async Propagation of trace in Kafka producer callback (DataDog/dd-trace-java#1727)
MODIFY
instrumentation/kafka-clients-0.11/src/main/java/io/opentelemetry/instrumentation/auto/kafkaclients/KafkaProducerInstrumentation.java
instrumentation/kafka-clients-0.11/src/main/java/io/opentelemetry/instrumentation/auto/kafkaclients/KafkaProducerInstrumentation.java
package io.opentelemetry.instrumentation.auto.kafkaclients;
+import static io.opentelemetry.context.ContextUtils.withScopedContext;
import static io.opentelemetry.instrumentation.auto.kafkaclients.KafkaDecorator.DECORATE;
import static io.opentelemetry.instrumentation.auto.kafkaclients.KafkaDecorator.TRACER;
import static io.opentelemetry.instrumentation.auto.kafkaclients.TextMapInjectAdapter.SETTER;
@@ -86,12 +87,13 @@ public final class KafkaProducerInstrumentation extends Instrumenter.Default {
@Advice.FieldValue("apiVersions") ApiVersions apiVersions,
@Advice.Argument(value = 0, readOnly = false) ProducerRecord record,
@Advice.Argument(value = 1, readOnly = false) Callback callback) {
+ Context parent = Context.current();
Span span =
TRACER.spanBuilder(DECORATE.spanNameOnProduce(record)).setSpanKind(PRODUCER).startSpan();
DECORATE.afterStart(span);
DECORATE.onProduce(span, record);
- callback = new ProducerCallback(callback, span);
+ callback = new ProducerCallback(callback, parent, span);
boolean isTombstone = record.value() == null && !record.headers().iterator().hasNext();
if (isTombstone) {
@@ -147,24 +149,27 @@ public final class KafkaProducerInstrumentation extends Instrumenter.Default {
public static class ProducerCallback implements Callback {
private final Callback callback;
+ private final Context parent;
private final Span span;
- public ProducerCallback(Callback callback, Span span) {
+ public ProducerCallback(Callback callback, Context parent, Span span) {
this.callback = callback;
+ this.parent = parent;
this.span = span;
}
@Override
public void onCompletion(RecordMetadata metadata, Exception exception) {
- try (Scope scope = currentContextWith(span)) {
DECORATE.onError(span, exception);
- try {
+ DECORATE.beforeFinish(span);
+ span.end();
if (callback != null) {
+ if (parent != null) {
+ try (Scope scope = withScopedContext(parent)) {
callback.onCompletion(metadata, exception);
}
- } finally {
- DECORATE.beforeFinish(span);
- span.end();
+ } else {
+ callback.onCompletion(metadata, exception);
}
}
}

MODIFY
instrumentation/kafka-clients-0.11/src/test/groovy/KafkaClientTest.groovy
instrumentation/kafka-clients-0.11/src/test/groovy/KafkaClientTest.groovy
*/
import static io.opentelemetry.auto.test.utils.ConfigUtils.withConfigOverride
+import static io.opentelemetry.auto.test.utils.TraceUtils.basicSpan
+import static io.opentelemetry.auto.test.utils.TraceUtils.runUnderTrace
import static io.opentelemetry.trace.Span.Kind.CONSUMER
import static io.opentelemetry.trace.Span.Kind.PRODUCER
@@ -27,8 +29,10 @@ 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.Producer
import org.apache.kafka.clients.producer.ProducerRecord
import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.serialization.StringSerializer
import org.junit.Rule
import org.springframework.kafka.core.DefaultKafkaConsumerFactory
import org.springframework.kafka.core.DefaultKafkaProducerFactory
@@ -47,6 +51,90 @@ class KafkaClientTest extends AgentTestRunner {
KafkaEmbedded embeddedKafka = new KafkaEmbedded(1, true, SHARED_TOPIC)
def "test kafka produce and consume"() {
+ setup:
+ def senderProps = KafkaTestUtils.senderProps(embeddedKafka.getBrokersAsString())
+ Producer<String, String> producer = new KafkaProducer<>(senderProps, new StringSerializer(), new StringSerializer())
+
+ // set up the Kafka consumer properties
+ def consumerProperties = KafkaTestUtils.consumerProps("sender", "false", embeddedKafka)
+
+ // create a Kafka consumer factory
+ def consumerFactory = new DefaultKafkaConsumerFactory<String, String>(consumerProperties)
+
+ // set the topic that needs to be consumed
+ def containerProperties = containerProperties()
+
+ // create a Kafka MessageListenerContainer
+ def container = new KafkaMessageListenerContainer<>(consumerFactory, containerProperties)
+
+ // create a thread safe queue to store the received message
+ def records = new LinkedBlockingQueue<ConsumerRecord<String, String>>()
+
+ // setup a Kafka message listener
+ container.setupMessageListener(new MessageListener<String, String>() {
+ @Override
+ void onMessage(ConsumerRecord<String, String> record) {
+ TEST_WRITER.waitForTraces(1) // ensure consistent ordering of traces
+ records.add(record)
+ }
+ })
+
+ // start the container and underlying message listener
+ container.start()
+
+ // wait until the container has the required number of assigned partitions
+ ContainerTestUtils.waitForAssignment(container, embeddedKafka.getPartitionsPerTopic())
+
+ when:
+ String greeting = "Hello Spring Kafka Sender!"
+ runUnderTrace("parent") {
+ producer.send(new ProducerRecord(SHARED_TOPIC, greeting)) { meta, ex ->
+ if (ex == null) {
+ runUnderTrace("producer callback") {}
+ } else {
+ runUnderTrace("producer exception: " + ex) {}
+ }
+ }
+ }
+
+ then:
+ // check that the message was received
+ def received = records.poll(5, TimeUnit.SECONDS)
+ received.value() == greeting
+ received.key() == null
+
+ assertTraces(1) {
+ trace(0, 4) {
+ basicSpan(it, 0, "parent")
+ span(1) {
+ operationName SHARED_TOPIC
+ spanKind PRODUCER
+ errored false
+ childOf span(0)
+ attributes {
+ }
+ }
+ span(2) {
+ operationName SHARED_TOPIC
+ spanKind CONSUMER
+ errored false
+ childOf span(1)
+ attributes {
+ "partition" { it >= 0 }
+ "offset" 0
+ "record.queue_time_ms" { it >= 0 }
+ }
+ }
+ basicSpan(it, 3, "producer callback", span(0))
+ }
+ }
+
+ cleanup:
+ producer.close()
+ container?.stop()
+ }
+
+ def "test spring kafka template produce and consume"() {
setup:
def senderProps = KafkaTestUtils.senderProps(embeddedKafka.getBrokersAsString())
def producerFactory = new DefaultKafkaProducerFactory<String, String>(senderProps)
@@ -83,8 +171,13 @@ class KafkaClientTest extends AgentTestRunner {
when:
String greeting = "Hello Spring Kafka Sender!"
- kafkaTemplate.send(SHARED_TOPIC, greeting)
-
+ runUnderTrace("parent") {
+ kafkaTemplate.send(SHARED_TOPIC, greeting).addCallback({
+ runUnderTrace("producer callback") {}
+ }, { ex ->
+ runUnderTrace("producer exception: " + ex) {}
+ })
+ }
then:
// check that the message was received
@@ -93,26 +186,28 @@ class KafkaClientTest extends AgentTestRunner {
received.key() == null
assertTraces(1) {
- trace(0, 2) {
- span(0) {
+ trace(0, 4) {
+ basicSpan(it, 0, "parent")
+ span(1) {
operationName SHARED_TOPIC
spanKind PRODUCER
errored false
- parent()
+ childOf span(0)
attributes {
}
}
- span(1) {
+ span(2) {
operationName SHARED_TOPIC
spanKind CONSUMER
errored false
- childOf span(0)
+ childOf span(1)
attributes {
"partition" { it >= 0 }
"offset" 0
"record.queue_time_ms" { it >= 0 }
}
}
+ basicSpan(it, 3, "producer callback", span(0))
}
}
@@ -121,7 +216,6 @@ class KafkaClientTest extends AgentTestRunner {
container?.stop()
}
-
def "test pass through tombstone"() {
setup:
def senderProps = KafkaTestUtils.senderProps(embeddedKafka.getBrokersAsString())
@@ -160,7 +254,6 @@ class KafkaClientTest extends AgentTestRunner {
when:
kafkaTemplate.send(SHARED_TOPIC, null)
-
then:
// check that the message was received
def received = records.poll(5, TimeUnit.SECONDS)
@@ -270,7 +363,6 @@ class KafkaClientTest extends AgentTestRunner {
cleanup:
consumer.close()
producer.close()
-
}
@Unroll
@@ -329,8 +421,6 @@ class KafkaClientTest extends AgentTestRunner {
value | expected
"false" | false
"true" | true
- String.valueOf(Config.DEFAULT_KAFKA_CLIENT_PROPAGATION_ENABLED) | true
-
}
def "should not read remote context when consuming messages if propagation is disabled"() {
@@ -448,7 +538,6 @@ class KafkaClientTest extends AgentTestRunner {
ConfigUtils.updateConfig {
System.clearProperty("otel." + Config.KAFKA_CLIENT_PROPAGATION_ENABLED)
}
-
}
protected KafkaMessageListenerContainer<Object, Object> startConsumer(String groupId, records) {

Enable Async Propagation of trace in Kafka producer callback (DataDog/dd-trace-java#1727)