Skip to content

Commit 5a2f529

Browse files
lauritlaurit
authored
Make kafka library and javaagent instrumentations more similar (#9738)
1 parent 1500647 commit 5a2f529

File tree

26 files changed

+933
-298
lines changed

26 files changed

+933
-298
lines changed

‎instrumentation/kafka/kafka-clients/kafka-clients-0.11/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/kafka/KafkaClientsConsumerProcessTracing.java

+6
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
package io.opentelemetry.javaagent.bootstrap.kafka;
77

8+
import java.util.function.BooleanSupplier;
9+
810
// Classes used by multiple instrumentations should be in a bootstrap module to ensure that all
911
// instrumentations see the same class. Helper classes are injected into each class loader that
1012
// contains an instrumentation that uses them, so instrumentations in different class loaders will
@@ -23,4 +25,8 @@ public static boolean setEnabled(boolean enabled) {
2325
public static boolean wrappingEnabled() {
2426
return wrappingEnabled.get();
2527
}
28+
29+
public static BooleanSupplier wrappingEnabledSupplier() {
30+
return KafkaClientsConsumerProcessTracing::wrappingEnabled;
31+
}
2632
}

‎instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/ConsumerRecordsInstrumentation.java

+14-3
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
package io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11;
77

8+
import static io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing.wrappingEnabledSupplier;
9+
import static io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11.KafkaSingletons.consumerProcessInstrumenter;
810
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
911
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
1012
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -14,6 +16,9 @@
1416

1517
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
1618
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
19+
import io.opentelemetry.instrumentation.kafka.internal.TracingIterable;
20+
import io.opentelemetry.instrumentation.kafka.internal.TracingIterator;
21+
import io.opentelemetry.instrumentation.kafka.internal.TracingList;
1722
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
1823
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
1924
import java.util.Iterator;
@@ -70,7 +75,9 @@ public static <K, V> void wrap(
7075
// case it's important to overwrite the leaked span instead of suppressing the correct span
7176
// (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
7277
KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
73-
iterable = TracingIterable.wrap(iterable, consumerContext);
78+
iterable =
79+
TracingIterable.wrap(
80+
iterable, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
7481
}
7582
}
7683

@@ -88,7 +95,9 @@ public static <K, V> void wrap(
8895
// case it's important to overwrite the leaked span instead of suppressing the correct span
8996
// (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
9097
KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
91-
list = TracingList.wrap(list, consumerContext);
98+
list =
99+
TracingList.wrap(
100+
list, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
92101
}
93102
}
94103

@@ -106,7 +115,9 @@ public static <K, V> void wrap(
106115
// case it's important to overwrite the leaked span instead of suppressing the correct span
107116
// (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
108117
KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
109-
iterator = TracingIterator.wrap(iterator, consumerContext);
118+
iterator =
119+
TracingIterator.wrap(
120+
iterator, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
110121
}
111122
}
112123
}

‎instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/TracingIterable.java

-47
This file was deleted.

‎instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/build.gradle.kts

+20
Original file line numberDiff line numberDiff line change
@@ -20,4 +20,24 @@ tasks {
2020
withType<Test>().configureEach {
2121
usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
2222
}
23+
24+
val testReceiveSpansDisabled by registering(Test::class) {
25+
filter {
26+
includeTestsMatching("InterceptorsSuppressReceiveSpansTest")
27+
includeTestsMatching("WrapperSuppressReceiveSpansTest")
28+
}
29+
include("**/InterceptorsSuppressReceiveSpansTest.*", "**/WrapperSuppressReceiveSpansTest.*")
30+
}
31+
32+
test {
33+
filter {
34+
excludeTestsMatching("InterceptorsSuppressReceiveSpansTest")
35+
excludeTestsMatching("WrapperSuppressReceiveSpansTest")
36+
}
37+
jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
38+
}
39+
40+
check {
41+
dependsOn(testReceiveSpansDisabled)
42+
}
2343
}

‎instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/KafkaTelemetry.java

+65-14
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,24 @@
1313
import io.opentelemetry.context.propagation.TextMapPropagator;
1414
import io.opentelemetry.context.propagation.TextMapSetter;
1515
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
16+
import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
17+
import io.opentelemetry.instrumentation.api.internal.Timer;
18+
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
19+
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
1620
import io.opentelemetry.instrumentation.kafka.internal.KafkaHeadersSetter;
1721
import io.opentelemetry.instrumentation.kafka.internal.KafkaProcessRequest;
1822
import io.opentelemetry.instrumentation.kafka.internal.KafkaProducerRequest;
23+
import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
1924
import io.opentelemetry.instrumentation.kafka.internal.KafkaUtil;
2025
import io.opentelemetry.instrumentation.kafka.internal.OpenTelemetryMetricsReporter;
2126
import io.opentelemetry.instrumentation.kafka.internal.OpenTelemetrySupplier;
27+
import io.opentelemetry.instrumentation.kafka.internal.TracingList;
2228
import java.lang.reflect.InvocationTargetException;
2329
import java.lang.reflect.Proxy;
2430
import java.util.Collections;
2531
import java.util.HashMap;
32+
import java.util.LinkedHashMap;
33+
import java.util.List;
2634
import java.util.Map;
2735
import java.util.concurrent.Future;
2836
import java.util.function.BiFunction;
@@ -37,6 +45,7 @@
3745
import org.apache.kafka.clients.producer.Producer;
3846
import org.apache.kafka.clients.producer.ProducerRecord;
3947
import org.apache.kafka.clients.producer.RecordMetadata;
48+
import org.apache.kafka.common.TopicPartition;
4049
import org.apache.kafka.common.header.Headers;
4150
import org.apache.kafka.common.metrics.MetricsReporter;
4251

@@ -47,16 +56,19 @@ public final class KafkaTelemetry {
4756

4857
private final OpenTelemetry openTelemetry;
4958
private final Instrumenter<KafkaProducerRequest, RecordMetadata> producerInstrumenter;
59+
private final Instrumenter<KafkaReceiveRequest, Void> consumerReceiveInstrumenter;
5060
private final Instrumenter<KafkaProcessRequest, Void> consumerProcessInstrumenter;
5161
private final boolean producerPropagationEnabled;
5262

5363
KafkaTelemetry(
5464
OpenTelemetry openTelemetry,
5565
Instrumenter<KafkaProducerRequest, RecordMetadata> producerInstrumenter,
66+
Instrumenter<KafkaReceiveRequest, Void> consumerReceiveInstrumenter,
5667
Instrumenter<KafkaProcessRequest, Void> consumerProcessInstrumenter,
5768
boolean producerPropagationEnabled) {
5869
this.openTelemetry = openTelemetry;
5970
this.producerInstrumenter = producerInstrumenter;
71+
this.consumerReceiveInstrumenter = consumerReceiveInstrumenter;
6072
this.consumerProcessInstrumenter = consumerProcessInstrumenter;
6173
this.producerPropagationEnabled = producerPropagationEnabled;
6274
}
@@ -115,6 +127,7 @@ public <K, V> Consumer<K, V> wrap(Consumer<K, V> consumer) {
115127
new Class<?>[] {Consumer.class},
116128
(proxy, method, args) -> {
117129
Object result;
130+
Timer timer = "poll".equals(method.getName()) ? Timer.start() : null;
118131
try {
119132
result = method.invoke(consumer, args);
120133
} catch (InvocationTargetException exception) {
@@ -123,12 +136,36 @@ public <K, V> Consumer<K, V> wrap(Consumer<K, V> consumer) {
123136
// ConsumerRecords<K, V> poll(long timeout)
124137
// ConsumerRecords<K, V> poll(Duration duration)
125138
if ("poll".equals(method.getName()) && result instanceof ConsumerRecords) {
126-
buildAndFinishSpan((ConsumerRecords) result, consumer);
139+
ConsumerRecords<K, V> consumerRecords = (ConsumerRecords<K, V>) result;
140+
Context receiveContext = buildAndFinishSpan(consumerRecords, consumer, timer);
141+
if (receiveContext == null) {
142+
receiveContext = Context.current();
143+
}
144+
KafkaConsumerContext consumerContext =
145+
KafkaConsumerContextUtil.create(receiveContext, consumer);
146+
result = addTracing(consumerRecords, consumerContext);
127147
}
128148
return result;
129149
});
130150
}
131151

152+
<K, V> ConsumerRecords<K, V> addTracing(
153+
ConsumerRecords<K, V> consumerRecords, KafkaConsumerContext consumerContext) {
154+
if (consumerRecords.isEmpty()) {
155+
return consumerRecords;
156+
}
157+
158+
Map<TopicPartition, List<ConsumerRecord<K, V>>> records = new LinkedHashMap<>();
159+
for (TopicPartition partition : consumerRecords.partitions()) {
160+
List<ConsumerRecord<K, V>> list = consumerRecords.records(partition);
161+
if (list != null && !list.isEmpty()) {
162+
list = TracingList.wrap(list, consumerProcessInstrumenter, () -> true, consumerContext);
163+
}
164+
records.put(partition, list);
165+
}
166+
return new ConsumerRecords<>(records);
167+
}
168+
132169
/**
133170
* Produces a set of kafka client config properties (consumer or producer) to register a {@link
134171
* MetricsReporter} that records metrics to an {@code openTelemetry} instance. Add these resulting
@@ -221,23 +258,37 @@ <K, V> Future<RecordMetadata> buildAndInjectSpan(
221258
}
222259
}
223260

224-
private <K, V> void buildAndFinishSpan(ConsumerRecords<K, V> records, Consumer<K, V> consumer) {
225-
buildAndFinishSpan(
226-
records, KafkaUtil.getConsumerGroup(consumer), KafkaUtil.getClientId(consumer));
261+
private <K, V> Context buildAndFinishSpan(
262+
ConsumerRecords<K, V> records, Consumer<K, V> consumer, Timer timer) {
263+
return buildAndFinishSpan(
264+
records, KafkaUtil.getConsumerGroup(consumer), KafkaUtil.getClientId(consumer), timer);
227265
}
228266

229-
<K, V> void buildAndFinishSpan(
230-
ConsumerRecords<K, V> records, String consumerGroup, String clientId) {
267+
<K, V> Context buildAndFinishSpan(
268+
ConsumerRecords<K, V> records, String consumerGroup, String clientId, Timer timer) {
269+
if (records.isEmpty()) {
270+
return null;
271+
}
231272
Context parentContext = Context.current();
232-
for (ConsumerRecord<K, V> record : records) {
233-
KafkaProcessRequest request = KafkaProcessRequest.create(record, consumerGroup, clientId);
234-
if (!consumerProcessInstrumenter.shouldStart(parentContext, request)) {
235-
continue;
236-
}
237-
238-
Context context = consumerProcessInstrumenter.start(parentContext, request);
239-
consumerProcessInstrumenter.end(context, request, null, null);
273+
KafkaReceiveRequest request = KafkaReceiveRequest.create(records, consumerGroup, clientId);
274+
Context context = null;
275+
if (consumerReceiveInstrumenter.shouldStart(parentContext, request)) {
276+
context =
277+
InstrumenterUtil.startAndEnd(
278+
consumerReceiveInstrumenter,
279+
parentContext,
280+
request,
281+
null,
282+
null,
283+
timer.startTime(),
284+
timer.now());
240285
}
286+
287+
// we're returning the context of the receive span so that process spans can use it as
288+
// parent context even though the span has ended
289+
// this is the suggested behavior according to the spec batch receive scenario:
290+
// https://github.com/open-telemetry/semantic-conventions/blob/main/docs/messaging/messaging-spans.md#batch-receiving
291+
return context;
241292
}
242293

243294
private class ProducerCallback implements Callback {

‎instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/KafkaTelemetryBuilder.java

+23-6
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@
1010
import com.google.errorprone.annotations.CanIgnoreReturnValue;
1111
import io.opentelemetry.api.OpenTelemetry;
1212
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
13-
import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessageOperation;
1413
import io.opentelemetry.instrumentation.kafka.internal.KafkaInstrumenterFactory;
1514
import io.opentelemetry.instrumentation.kafka.internal.KafkaProcessRequest;
1615
import io.opentelemetry.instrumentation.kafka.internal.KafkaProducerRequest;
16+
import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
1717
import java.util.ArrayList;
1818
import java.util.List;
1919
import java.util.Objects;
@@ -25,8 +25,10 @@ public final class KafkaTelemetryBuilder {
2525
private final OpenTelemetry openTelemetry;
2626
private final List<AttributesExtractor<KafkaProducerRequest, RecordMetadata>>
2727
producerAttributesExtractors = new ArrayList<>();
28-
private final List<AttributesExtractor<KafkaProcessRequest, Void>> consumerAttributesExtractors =
29-
new ArrayList<>();
28+
private final List<AttributesExtractor<KafkaProcessRequest, Void>>
29+
consumerProcessAttributesExtractors = new ArrayList<>();
30+
private final List<AttributesExtractor<KafkaReceiveRequest, Void>>
31+
consumerReceiveAttributesExtractors = new ArrayList<>();
3032
private List<String> capturedHeaders = emptyList();
3133
private boolean captureExperimentalSpanAttributes = false;
3234
private boolean propagationEnabled = true;
@@ -43,10 +45,25 @@ public KafkaTelemetryBuilder addProducerAttributesExtractors(
4345
return this;
4446
}
4547

48+
/** Use {@link #addConsumerProcessAttributesExtractors(AttributesExtractor)} instead. */
49+
@Deprecated
4650
@CanIgnoreReturnValue
4751
public KafkaTelemetryBuilder addConsumerAttributesExtractors(
4852
AttributesExtractor<KafkaProcessRequest, Void> extractor) {
49-
consumerAttributesExtractors.add(extractor);
53+
return addConsumerProcessAttributesExtractors(extractor);
54+
}
55+
56+
@CanIgnoreReturnValue
57+
public KafkaTelemetryBuilder addConsumerProcessAttributesExtractors(
58+
AttributesExtractor<KafkaProcessRequest, Void> extractor) {
59+
consumerProcessAttributesExtractors.add(extractor);
60+
return this;
61+
}
62+
63+
@CanIgnoreReturnValue
64+
public KafkaTelemetryBuilder addConsumerReceiveAttributesExtractors(
65+
AttributesExtractor<KafkaReceiveRequest, Void> extractor) {
66+
consumerReceiveAttributesExtractors.add(extractor);
5067
return this;
5168
}
5269

@@ -109,8 +126,8 @@ public KafkaTelemetry build() {
109126
return new KafkaTelemetry(
110127
openTelemetry,
111128
instrumenterFactory.createProducerInstrumenter(producerAttributesExtractors),
112-
instrumenterFactory.createConsumerOperationInstrumenter(
113-
MessageOperation.RECEIVE, consumerAttributesExtractors),
129+
instrumenterFactory.createConsumerReceiveInstrumenter(consumerReceiveAttributesExtractors),
130+
instrumenterFactory.createConsumerProcessInstrumenter(consumerProcessAttributesExtractors),
114131
propagationEnabled);
115132
}
116133
}

‎instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/TracingConsumerInterceptor.java

+20-3
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,11 @@
77

88
import com.google.errorprone.annotations.CanIgnoreReturnValue;
99
import io.opentelemetry.api.GlobalOpenTelemetry;
10+
import io.opentelemetry.context.Context;
11+
import io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil;
12+
import io.opentelemetry.instrumentation.api.internal.Timer;
13+
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
14+
import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
1015
import java.util.Map;
1116
import java.util.Objects;
1217
import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -22,16 +27,28 @@
2227
*/
2328
public class TracingConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
2429

25-
private static final KafkaTelemetry telemetry = KafkaTelemetry.create(GlobalOpenTelemetry.get());
30+
private static final KafkaTelemetry telemetry =
31+
KafkaTelemetry.builder(GlobalOpenTelemetry.get())
32+
.setMessagingReceiveInstrumentationEnabled(
33+
ConfigPropertiesUtil.getBoolean(
34+
"otel.instrumentation.messaging.experimental.receive-telemetry.enabled", false))
35+
.build();
2636

2737
private String consumerGroup;
2838
private String clientId;
2939

3040
@Override
3141
@CanIgnoreReturnValue
3242
public ConsumerRecords<K, V> onConsume(ConsumerRecords<K, V> records) {
33-
telemetry.buildAndFinishSpan(records, consumerGroup, clientId);
34-
return records;
43+
// timer should be started before fetching ConsumerRecords, but there is no callback for that
44+
Timer timer = Timer.start();
45+
Context receiveContext = telemetry.buildAndFinishSpan(records, consumerGroup, clientId, timer);
46+
if (receiveContext == null) {
47+
receiveContext = Context.current();
48+
}
49+
KafkaConsumerContext consumerContext =
50+
KafkaConsumerContextUtil.create(receiveContext, consumerGroup, clientId);
51+
return telemetry.addTracing(records, consumerContext);
3552
}
3653

3754
@Override

0 commit comments

Comments
 (0)