Skip to content

Commit 58d9872

Browse files
amarzialiamarziali
committed
Fix Kafka async consumer commit callback registration
1 parent 134a70e commit 58d9872

5 files changed

Lines changed: 175 additions & 17 deletions

File tree

‎dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java/datadog/trace/instrumentation/kafka_clients38/OffsetCommitCallbackInvokerInstrumentation.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,16 @@
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
55
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
66

7+
import com.google.auto.service.AutoService;
78
import datadog.trace.agent.tooling.Instrumenter;
89
import datadog.trace.agent.tooling.InstrumenterModule;
10+
import java.util.HashMap;
11+
import java.util.Map;
912
import net.bytebuddy.matcher.ElementMatcher;
1013

1114
// new - this instrumentation is completely new.
1215
// the purpose of this class is to provide us with information on consumer group and cluster ID
16+
@AutoService(InstrumenterModule.class)
1317
public class OffsetCommitCallbackInvokerInstrumentation extends InstrumenterModule.Tracing
1418
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
1519
public OffsetCommitCallbackInvokerInstrumentation() {
@@ -26,6 +30,26 @@ public String instrumentedType() {
2630
return "org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker";
2731
}
2832

33+
@Override
34+
public String[] helperClassNames() {
35+
return new String[] {
36+
"datadog.trace.instrumentation.kafka_common.PendingConfig",
37+
"datadog.trace.instrumentation.kafka_common.MetadataState",
38+
packageName + ".KafkaConsumerInfo",
39+
packageName + ".DDOffsetCommitCallback"
40+
};
41+
}
42+
43+
@Override
44+
public Map<String, String> contextStore() {
45+
Map<String, String> stores = new HashMap<>(2);
46+
stores.put(instrumentedType(), packageName + ".KafkaConsumerInfo");
47+
stores.put(
48+
"org.apache.kafka.clients.Metadata",
49+
"datadog.trace.instrumentation.kafka_common.MetadataState");
50+
return stores;
51+
}
52+
2953
@Override
3054
public void methodAdvice(MethodTransformer transformer) {
3155
transformer.applyAdvice(

‎dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/DDOffsetCommitCallback.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package datadog.trace.instrumentation.kafka_clients38;
22

33
import datadog.trace.api.datastreams.DataStreamsTags;
4-
import datadog.trace.bootstrap.InstrumentationContext;
4+
import datadog.trace.bootstrap.ContextStore;
55
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
66
import datadog.trace.instrumentation.kafka_common.MetadataState;
77
import java.util.Map;
@@ -13,11 +13,15 @@
1313
public class DDOffsetCommitCallback implements OffsetCommitCallback {
1414
OffsetCommitCallback callback;
1515
KafkaConsumerInfo kafkaConsumerInfo;
16+
private final ContextStore<Metadata, MetadataState> metadataStore;
1617

1718
public DDOffsetCommitCallback(
18-
OffsetCommitCallback callback, KafkaConsumerInfo kafkaConsumerInfo) {
19+
OffsetCommitCallback callback,
20+
KafkaConsumerInfo kafkaConsumerInfo,
21+
ContextStore<Metadata, MetadataState> metadataStore) {
1922
this.callback = callback;
2023
this.kafkaConsumerInfo = kafkaConsumerInfo;
24+
this.metadataStore = metadataStore;
2125
}
2226

2327
@Override
@@ -36,8 +40,7 @@ public void onComplete(Map<TopicPartition, OffsetAndMetadata> map, Exception e)
3640
consumerGroup = kafkaConsumerInfo.getConsumerGroup().orElse(null);
3741
Metadata consumerMetadata = kafkaConsumerInfo.getmetadata().orElse(null);
3842
if (consumerMetadata != null) {
39-
MetadataState metadataState =
40-
InstrumentationContext.get(Metadata.class, MetadataState.class).get(consumerMetadata);
43+
MetadataState metadataState = metadataStore.get(consumerMetadata);
4144
clusterId = metadataState != null ? metadataState.clusterId : null;
4245
}
4346
}

‎dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/OffsetCommitCallbackInvokerAdvice.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package datadog.trace.instrumentation.kafka_clients38;
22

33
import datadog.trace.bootstrap.InstrumentationContext;
4+
import datadog.trace.instrumentation.kafka_common.MetadataState;
45
import net.bytebuddy.asm.Advice;
6+
import org.apache.kafka.clients.Metadata;
57
import org.apache.kafka.clients.consumer.OffsetCommitCallback;
68
import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
79

@@ -13,6 +15,10 @@ public static void onEnqueue(
1315
KafkaConsumerInfo kafkaConsumerInfo =
1416
InstrumentationContext.get(OffsetCommitCallbackInvoker.class, KafkaConsumerInfo.class)
1517
.get(callbackInvoker);
16-
callback = new DDOffsetCommitCallback(callback, kafkaConsumerInfo);
18+
callback =
19+
new DDOffsetCommitCallback(
20+
callback,
21+
kafkaConsumerInfo,
22+
InstrumentationContext.get(Metadata.class, MetadataState.class));
1723
}
1824
}

‎dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/DDOffsetCommitCallbackNullGroupForkedTest.groovy‎

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import datadog.trace.instrumentation.kafka_clients38.KafkaConsumerInfo
44
import org.apache.kafka.clients.consumer.OffsetAndMetadata
55
import org.apache.kafka.clients.consumer.OffsetCommitCallback
66
import org.apache.kafka.common.TopicPartition
7+
import spock.util.concurrent.PollingConditions
78

89
/**
910
* Unit test to verify that DSM backlog tracking works even when consumer group is null.
@@ -51,15 +52,15 @@ class DDOffsetCommitCallbackNullGroupForkedTest extends VersionedNamingTestBase
5152
// Verify the consumer group is indeed empty
5253
assert !kafkaConsumerInfo.getConsumerGroup().isPresent() : "Consumer group should be empty for this test"
5354

54-
def callbackInvoked = false
55+
def callbacks = []
5556
def innerCallback = new OffsetCommitCallback() {
5657
@Override
5758
void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
58-
callbackInvoked = true
59+
callbacks.add([offsets, exception, Thread.currentThread()])
5960
}
6061
}
6162

62-
def ddCallback = new DDOffsetCommitCallback(innerCallback, kafkaConsumerInfo)
63+
def ddCallback = new DDOffsetCommitCallback(innerCallback, kafkaConsumerInfo, null)
6364

6465
def offsets = [
6566
(new TopicPartition("test-topic", 0)): new OffsetAndMetadata(100L)
@@ -73,24 +74,25 @@ class DDOffsetCommitCallbackNullGroupForkedTest extends VersionedNamingTestBase
7374

7475
then:
7576
// Inner callback should have been invoked
76-
callbackInvoked
77+
callbacks == [[offsets, null, Thread.currentThread()]]
7778

7879
// No exception was thrown (implicit - test would fail otherwise)
7980
noExceptionThrown()
81+
awaitBacklog(null)
8082
}
8183

8284
def "test DDOffsetCommitCallback handles null KafkaConsumerInfo"() {
8385
setup:
84-
def callbackInvoked = false
86+
def callbacks = []
8587
def innerCallback = new OffsetCommitCallback() {
8688
@Override
8789
void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
88-
callbackInvoked = true
90+
callbacks.add([offsets, exception, Thread.currentThread()])
8991
}
9092
}
9193

9294
// KafkaConsumerInfo is completely null
93-
def ddCallback = new DDOffsetCommitCallback(innerCallback, null)
95+
def ddCallback = new DDOffsetCommitCallback(innerCallback, null, null)
9496

9597
def offsets = [
9698
(new TopicPartition("test-topic", 0)): new OffsetAndMetadata(100L)
@@ -100,8 +102,9 @@ class DDOffsetCommitCallbackNullGroupForkedTest extends VersionedNamingTestBase
100102
ddCallback.onComplete(offsets, null)
101103

102104
then:
103-
callbackInvoked
105+
callbacks == [[offsets, null, Thread.currentThread()]]
104106
noExceptionThrown()
107+
awaitBacklog(null)
105108
}
106109

107110
def "test DDOffsetCommitCallback with valid consumer group still works"() {
@@ -111,15 +114,15 @@ class DDOffsetCommitCallbackNullGroupForkedTest extends VersionedNamingTestBase
111114
assert kafkaConsumerInfo.getConsumerGroup().isPresent()
112115
assert kafkaConsumerInfo.getConsumerGroup().get() == "test-group"
113116

114-
def callbackInvoked = false
117+
def callbacks = []
115118
def innerCallback = new OffsetCommitCallback() {
116119
@Override
117120
void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
118-
callbackInvoked = true
121+
callbacks.add([offsets, exception, Thread.currentThread()])
119122
}
120123
}
121124

122-
def ddCallback = new DDOffsetCommitCallback(innerCallback, kafkaConsumerInfo)
125+
def ddCallback = new DDOffsetCommitCallback(innerCallback, kafkaConsumerInfo, null)
123126

124127
def offsets = [
125128
(new TopicPartition("test-topic", 0)): new OffsetAndMetadata(100L)
@@ -129,7 +132,20 @@ class DDOffsetCommitCallbackNullGroupForkedTest extends VersionedNamingTestBase
129132
ddCallback.onComplete(offsets, null)
130133

131134
then:
132-
callbackInvoked
135+
callbacks == [[offsets, null, Thread.currentThread()]]
133136
noExceptionThrown()
137+
awaitBacklog("test-group")
138+
}
139+
140+
private void awaitBacklog(String group) {
141+
new PollingConditions(timeout: 15).eventually {
142+
assert TEST_DATA_STREAMS_WRITER.payloads.any { bucket ->
143+
bucket.backlogs.any { entry ->
144+
entry.value == 100L && entry.key.kafkaClusterId == null &&
145+
entry.key.consumerGroup == (group == null ? null : "consumer_group:$group".toString()) &&
146+
entry.key.hasAllTags("type:kafka_commit", "topic:test-topic", "partition:0")
147+
}
148+
}
149+
}
134150
}
135151
}

‎dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaAsyncCommitCallbackForkedTest.groovy‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import datadog.trace.agent.test.naming.VersionedNamingTestBase
2+
import org.apache.kafka.clients.consumer.ConsumerConfig
3+
import org.apache.kafka.clients.consumer.KafkaConsumer
4+
import org.apache.kafka.clients.consumer.OffsetAndMetadata
5+
import org.apache.kafka.clients.consumer.OffsetCommitCallback
6+
import org.apache.kafka.common.TopicPartition
7+
import org.apache.kafka.common.serialization.StringDeserializer
8+
import org.springframework.kafka.test.EmbeddedKafkaKraftBroker
9+
import org.springframework.kafka.test.utils.KafkaTestUtils
10+
import spock.lang.Shared
11+
import spock.util.concurrent.PollingConditions
12+
13+
import java.time.Duration
14+
15+
class KafkaAsyncCommitCallbackForkedTest extends VersionedNamingTestBase {
16+
static final String TOPIC = "async-commit-callback"
17+
static final TopicPartition PARTITION = new TopicPartition(TOPIC, 0)
18+
static final Duration TIMEOUT = Duration.ofSeconds(15)
19+
20+
@Shared
21+
EmbeddedKafkaKraftBroker broker
22+
23+
def setupSpec() {
24+
broker = new EmbeddedKafkaKraftBroker(1, 1, TOPIC)
25+
broker.brokerProperty("group.coordinator.rebalance.protocols", "classic,consumer")
26+
broker.brokerProperty("group.coordinator.new.enable", "true")
27+
broker.afterPropertiesSet()
28+
}
29+
30+
def cleanupSpec() {
31+
broker?.destroy()
32+
}
33+
34+
@Override
35+
int version() {
36+
0
37+
}
38+
39+
@Override
40+
String operation() {
41+
null
42+
}
43+
44+
@Override
45+
String service() {
46+
"kafka"
47+
}
48+
49+
@Override
50+
protected boolean isDataStreamsEnabled() {
51+
true
52+
}
53+
54+
def "async commit preserves callback and reports backlog during #completion"() {
55+
setup:
56+
def group = "async-${UUID.randomUUID()}"
57+
def properties = KafkaTestUtils.consumerProps(group, "false", broker)
58+
properties.remove(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG)
59+
properties.remove(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG)
60+
properties.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, "consumer")
61+
properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer)
62+
properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer)
63+
def consumer = new KafkaConsumer<String, String>(properties)
64+
assert consumer.@delegate.class.simpleName == "AsyncKafkaConsumer"
65+
consumer.assign([PARTITION])
66+
consumer.seek(PARTITION, 0L)
67+
// Resolve the coordinator before the one-shot asynchronous commit under test.
68+
consumer.committed(Collections.singleton(PARTITION), TIMEOUT)
69+
def callbacks = []
70+
def applicationThread = Thread.currentThread()
71+
def offsets = [(PARTITION): new OffsetAndMetadata(2L)]
72+
73+
when:
74+
consumer.commitAsync(offsets, { Map committed, Exception error ->
75+
callbacks.add([committed, error, Thread.currentThread()])
76+
} as OffsetCommitCallback)
77+
if (completion == "close") {
78+
consumer.close(TIMEOUT)
79+
} else if (completion == "sync commit") {
80+
consumer.commitSync(Collections.emptyMap(), TIMEOUT)
81+
} else {
82+
long deadline = System.nanoTime() + TIMEOUT.toNanos()
83+
while (callbacks.empty && System.nanoTime() < deadline) {
84+
consumer.poll(Duration.ofMillis(100))
85+
}
86+
}
87+
88+
then:
89+
callbacks.size() == 1
90+
callbacks[0][0] == offsets
91+
callbacks[0][1] == null
92+
callbacks[0][2].is(applicationThread)
93+
new PollingConditions(timeout: 15).eventually {
94+
assert TEST_DATA_STREAMS_WRITER.payloads.any { bucket ->
95+
bucket.backlogs.any { entry ->
96+
entry.value == 2L && entry.key.kafkaClusterId &&
97+
entry.key.hasAllTags("type:kafka_commit", "consumer_group:$group".toString(),
98+
"topic:$TOPIC".toString(), "partition:0")
99+
}
100+
}
101+
}
102+
103+
cleanup:
104+
consumer?.close(TIMEOUT)
105+
106+
where:
107+
completion << ["poll", "sync commit", "close"]
108+
}
109+
}

0 commit comments

Comments
 (0)