Skip to content

Commit c8d1107

Browse files
stankiewiczstankiewicz
committed
OTEL in kafka.
1 parent 8d252c4 commit c8d1107

2 files changed

Lines changed: 154 additions & 4 deletions

File tree

‎sdks/java/io/kafka/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,8 @@ dependencies {
6767
}
6868
testImplementation library.java.kafka_clients
6969
testImplementation project(path: ":runners:core-java")
70+
implementation library.java.opentelemetry_api
71+
implementation library.java.opentelemetry_context
7072
implementation library.java.slf4j_api
7173
implementation library.java.joda_time
7274
implementation library.java.jackson_annotations

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java‎

Lines changed: 152 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
*/
1818
package org.apache.beam.sdk.io.kafka;
1919

20+
import static java.nio.charset.StandardCharsets.UTF_8;
2021
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
2122
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
2223
import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
@@ -25,6 +26,13 @@
2526
import com.google.auto.value.AutoValue;
2627
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
2728
import io.confluent.kafka.serializers.KafkaAvroDeserializer;
29+
import io.opentelemetry.api.trace.Span;
30+
import io.opentelemetry.api.trace.Tracer;
31+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
32+
import io.opentelemetry.context.Context;
33+
import io.opentelemetry.context.Scope;
34+
import io.opentelemetry.context.propagation.TextMapGetter;
35+
import io.opentelemetry.context.propagation.TextMapSetter;
2836
import java.io.InputStream;
2937
import java.io.OutputStream;
3038
import java.lang.reflect.Method;
@@ -40,6 +48,7 @@
4048
import java.util.Set;
4149
import java.util.regex.Pattern;
4250
import java.util.stream.Collectors;
51+
import java.util.stream.StreamSupport;
4352
import org.apache.beam.sdk.annotations.Internal;
4453
import org.apache.beam.sdk.coders.AtomicCoder;
4554
import org.apache.beam.sdk.coders.ByteArrayCoder;
@@ -61,6 +70,7 @@
6170
import org.apache.beam.sdk.options.Default;
6271
import org.apache.beam.sdk.options.ExperimentalOptions;
6372
import org.apache.beam.sdk.options.PipelineOptions;
73+
import org.apache.beam.sdk.options.SdkHarnessOptions;
6474
import org.apache.beam.sdk.options.StreamingOptions;
6575
import org.apache.beam.sdk.options.ValueProvider;
6676
import org.apache.beam.sdk.runners.AppliedPTransform;
@@ -125,6 +135,7 @@
125135
import org.apache.kafka.common.TopicPartition;
126136
import org.apache.kafka.common.config.SaslConfigs;
127137
import org.apache.kafka.common.header.Header;
138+
import org.apache.kafka.common.header.Headers;
128139
import org.apache.kafka.common.header.internals.RecordHeader;
129140
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
130141
import org.apache.kafka.common.serialization.Deserializer;
@@ -614,6 +625,7 @@ public static <K, V> Read<K, V> read() {
614625
.setTimestampPolicyFactory(TimestampPolicyFactory.withProcessingTime())
615626
.setConsumerPollingTimeout(2L)
616627
.setRedistributed(false)
628+
.setEnableOpenTelemetryTracing(false)
617629
.setAllowDuplicates(false)
618630
.setRedistributeNumKeys(0)
619631
.build();
@@ -653,6 +665,7 @@ public static <K, V> WriteRecords<K, V> writeRecords() {
653665
.setEosTriggerNumElements(1) // keep default numElements
654666
.setEosTriggerTimeout(null) // keep default trigger (timeout)
655667
.setNumShards(0)
668+
.setEnableOpenTelemetryTracing(false)
656669
.setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
657670
.setBadRecordRouter(BadRecordRouter.THROWING_ROUTER)
658671
.setBadRecordErrorHandler(new DefaultErrorHandler<>())
@@ -742,6 +755,9 @@ public abstract static class Read<K, V>
742755
@Pure
743756
public abstract @Nullable Duration getWatchTopicPartitionDuration();
744757

758+
@Pure
759+
public abstract boolean isEnableOpenTelemetryTracing();
760+
745761
@Pure
746762
public abstract TimestampPolicyFactory<K, V> getTimestampPolicyFactory();
747763

@@ -832,6 +848,8 @@ Builder<K, V> setCheckStopReadingFn(
832848
return setCheckStopReadingFn(CheckStopReadingFnWrapper.of(checkStopReadingFn));
833849
}
834850

851+
abstract Builder<K, V> setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing);
852+
835853
abstract Builder<K, V> setConsumerPollingTimeout(long consumerPollingTimeout);
836854

837855
abstract Builder<K, V> setLogTopicVerification(@Nullable Boolean logTopicVerification);
@@ -865,6 +883,7 @@ static <K, V> void setupExternalBuilder(
865883

866884
// Set required defaults
867885
builder.setTopicPartitions(Collections.emptyList());
886+
builder.setEnableOpenTelemetryTracing(false);
868887
builder.setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN);
869888
if (config.maxReadTime != null) {
870889
builder.setMaxReadTime(Duration.standardSeconds(config.maxReadTime));
@@ -1302,6 +1321,10 @@ public Read<K, V> withValueDeserializer(DeserializerProvider<V> deserializerProv
13021321
return toBuilder().setValueDeserializerProvider(deserializerProvider).build();
13031322
}
13041323

1324+
public Read<K, V> withEnableOpenTelemetryTracing() {
1325+
return toBuilder().setEnableOpenTelemetryTracing(true).build();
1326+
}
1327+
13051328
public Read<K, V> withValueDeserializerProviderAndCoder(
13061329
DeserializerProvider<V> deserializerProvider, Coder<V> valueCoder) {
13071330
return toBuilder()
@@ -1920,6 +1943,14 @@ public PCollection<KafkaRecord<K, V>> expand(PBegin input) {
19201943
.withMaxNumRecords(kafkaRead.getMaxNumRecords());
19211944
}
19221945
PCollection<KafkaRecord<K, V>> output = input.getPipeline().apply(transform);
1946+
1947+
if (kafkaRead.isEnableOpenTelemetryTracing()) {
1948+
output =
1949+
output.apply(
1950+
"Extract OpenTelemetry context from Header",
1951+
ParDo.of(new OpenTelemetryHeaderConsumer<>()));
1952+
}
1953+
19231954
if (kafkaRead.getOffsetDeduplication() != null && kafkaRead.getOffsetDeduplication()) {
19241955
output =
19251956
output.apply(
@@ -2041,9 +2072,15 @@ public PCollection<KafkaRecord<K, V>> expand(PBegin input) {
20412072
.apply(ParDo.of(new GenerateKafkaSourceDescriptor(kafkaRead)));
20422073
}
20432074
}
2075+
PCollection<KafkaRecord<K, V>> pcol =
2076+
output.apply(readTransform).setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
2077+
if (kafkaRead.isEnableOpenTelemetryTracing()) {
2078+
pcol =
2079+
pcol.apply(
2080+
"Extract OpenTelemetry context from Header",
2081+
ParDo.of(new OpenTelemetryHeaderConsumer<>()));
2082+
}
20442083
if (kafkaRead.isRedistributed()) {
2045-
PCollection<KafkaRecord<K, V>> pcol =
2046-
output.apply(readTransform).setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
20472084
if (kafkaRead.getRedistributeNumKeys() == 0) {
20482085
return pcol.apply(
20492086
"Insert Redistribute",
@@ -2057,7 +2094,7 @@ public PCollection<KafkaRecord<K, V>> expand(PBegin input) {
20572094
.withNumBuckets((int) kafkaRead.getRedistributeNumKeys()));
20582095
}
20592096
}
2060-
return output.apply(readTransform).setCoder(KafkaRecordCoder.of(keyCoder, valueCoder));
2097+
return pcol;
20612098
}
20622099
}
20632100

@@ -2218,6 +2255,101 @@ public void populateDisplayData(DisplayData.Builder builder) {
22182255
}
22192256
}
22202257

2258+
static class OpenTelemetryHeaderConsumer<K, V>
2259+
extends DoFn<KafkaRecord<K, V>, KafkaRecord<K, V>> {
2260+
@Nullable Tracer tracer = null;
2261+
2262+
@Setup
2263+
public void setup(PipelineOptions options) {
2264+
// inject tracer via options
2265+
io.opentelemetry.api.OpenTelemetry openTelemetry =
2266+
options.as(SdkHarnessOptions.class).getOpenTelemetry();
2267+
if (openTelemetry != null) {
2268+
tracer = openTelemetry.getTracer("KafkaIO");
2269+
}
2270+
}
2271+
2272+
Context extractSpanContext(KafkaRecord<K, V> message) {
2273+
TextMapGetter<KafkaRecord<K, V>> extractMessageAttributes =
2274+
new TextMapGetter<KafkaRecord<K, V>>() {
2275+
2276+
@Override
2277+
public @Nullable String get(@Nullable KafkaRecord<K, V> carrier, String key) {
2278+
if (carrier == null) {
2279+
return null;
2280+
}
2281+
Headers headers = carrier.getHeaders();
2282+
if (headers == null) {
2283+
return null;
2284+
}
2285+
Header header = headers.lastHeader(key);
2286+
if (header == null) {
2287+
return null;
2288+
}
2289+
return new String(header.value(), UTF_8);
2290+
}
2291+
2292+
@Override
2293+
public Iterable<String> keys(@Nullable KafkaRecord<K, V> carrier) {
2294+
if (carrier == null || carrier.getHeaders() == null) {
2295+
return ImmutableList.of();
2296+
}
2297+
return StreamSupport.stream(carrier.getHeaders().spliterator(), false)
2298+
.map(Header::key)
2299+
.collect(Collectors.toList());
2300+
}
2301+
};
2302+
return W3CTraceContextPropagator.getInstance()
2303+
.extract(Context.current(), message, extractMessageAttributes);
2304+
}
2305+
2306+
@ProcessElement
2307+
public void processElement(
2308+
@Element KafkaRecord<K, V> element, OutputReceiver<KafkaRecord<K, V>> receiver) {
2309+
Context context = extractSpanContext(element);
2310+
Span span =
2311+
Preconditions.checkArgumentNotNull(tracer)
2312+
.spanBuilder("KafkaIO.Read")
2313+
.setParent(context)
2314+
.startSpan();
2315+
try (Scope ignored = span.makeCurrent()) {
2316+
receiver.output(element);
2317+
} finally {
2318+
span.end();
2319+
}
2320+
}
2321+
}
2322+
2323+
static class OpenTelemetryHeaderPropagator<K, V>
2324+
extends DoFn<ProducerRecord<K, V>, ProducerRecord<K, V>> {
2325+
ProducerRecord<K, V> injectTraceContext(ProducerRecord<K, V> message) {
2326+
org.apache.kafka.common.header.internals.RecordHeaders headers =
2327+
new org.apache.kafka.common.header.internals.RecordHeaders(message.headers());
2328+
TextMapSetter<org.apache.kafka.common.header.internals.RecordHeaders>
2329+
injectMessageAttributes =
2330+
(carrier, key, value) -> {
2331+
if (carrier != null) {
2332+
carrier.add(key, value.getBytes(UTF_8));
2333+
}
2334+
};
2335+
W3CTraceContextPropagator.getInstance()
2336+
.inject(Context.current(), headers, injectMessageAttributes);
2337+
return new ProducerRecord<>(
2338+
message.topic(),
2339+
message.partition(),
2340+
message.timestamp(),
2341+
message.key(),
2342+
message.value(),
2343+
headers);
2344+
}
2345+
2346+
@ProcessElement
2347+
public void processElement(
2348+
@Element ProducerRecord<K, V> element, OutputReceiver<ProducerRecord<K, V>> receiver) {
2349+
receiver.output(injectTraceContext(element));
2350+
}
2351+
}
2352+
22212353
/**
22222354
* A {@link PTransform} to read from Kafka topics. Similar to {@link KafkaIO.Read}, but removes
22232355
* Kafka metatdata and returns a {@link PCollection} of {@link KV}. See {@link KafkaIO} for more
@@ -3162,6 +3294,8 @@ public abstract static class WriteRecords<K, V>
31623294
// we shouldn't have to duplicate the same API for similar transforms like {@link Write} and
31633295
// {@link WriteRecords}. See example at {@link PubsubIO.Write}.
31643296

3297+
public abstract boolean isEnableOpenTelemetryTracing();
3298+
31653299
@Pure
31663300
public abstract @Nullable String getTopic();
31673301

@@ -3212,6 +3346,8 @@ public abstract static class WriteRecords<K, V>
32123346
abstract static class Builder<K, V> {
32133347
abstract Builder<K, V> setTopic(String topic);
32143348

3349+
abstract Builder<K, V> setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing);
3350+
32153351
abstract Builder<K, V> setProducerConfig(Map<String, Object> producerConfig);
32163352

32173353
abstract Builder<K, V> setProducerFactoryFn(
@@ -3277,6 +3413,10 @@ public WriteRecords<K, V> withValueSerializer(Class<? extends Serializer<V>> val
32773413
return toBuilder().setValueSerializer(valueSerializer).build();
32783414
}
32793415

3416+
public WriteRecords<K, V> withEnableOpenTelemetryTracing() {
3417+
return toBuilder().setEnableOpenTelemetryTracing(true).build();
3418+
}
3419+
32803420
/**
32813421
* Adds the given producer properties, overriding old values of properties with the same key.
32823422
*
@@ -3413,7 +3553,11 @@ public PDone expand(PCollection<ProducerRecord<K, V>> input) {
34133553

34143554
checkArgument(getKeySerializer() != null, "withKeySerializer() is required");
34153555
checkArgument(getValueSerializer() != null, "withValueSerializer() is required");
3416-
3556+
if (this.isEnableOpenTelemetryTracing()) {
3557+
input =
3558+
input.apply(
3559+
"Propagate OpenTelemetry Tracing", ParDo.of(new OpenTelemetryHeaderPropagator<>()));
3560+
}
34173561
if (isEOS()) {
34183562
checkArgument(getTopic() != null, "withTopic() is required when isEOS() is true");
34193563
checkArgument(
@@ -3653,6 +3797,10 @@ public Write<K, V> withInputTimestamp() {
36533797
return withWriteRecordsTransform(getWriteRecordsTransform().withInputTimestamp());
36543798
}
36553799

3800+
public Write<K, V> withEnableOpenTelemetryTracing() {
3801+
return withWriteRecordsTransform(getWriteRecordsTransform().withEnableOpenTelemetryTracing());
3802+
}
3803+
36563804
/**
36573805
* Wrapper method over {@link
36583806
* WriteRecords#withPublishTimestampFunction(KafkaPublishTimestampFunction)}, used to keep the

0 commit comments

Comments
 (0)