From 537503d6e06aa38a2223fc3d3c03c7fec84d8242 Mon Sep 17 00:00:00 2001 From: maciej-gromul <89904409+maciej-gromul@users.noreply.github.com> Date: Thu, 9 Mar 2023 15:01:40 +0100 Subject: [PATCH] New instrumentation for reactive kafka clients (#2268) * New instrumentation for reactive kafka clients * Fix javadocs for auto configuration * Javadocs and copyright doc changes --- docs/src/main/asciidoc/integrations.adoc | 24 +++- ...TracingKafkaReceiverBeanPostProcessor.java | 50 +++++++ .../TracingReactorKafkaAutoConfiguration.java | 33 ++++- ...aceReactorKafkaAutoConfigurationTests.java | 33 ++++- .../kafka/ReactiveKafkaTracingPropagator.java | 62 +++++++++ .../kafka/TracingKafkaConsumerFactory.java | 49 ------- .../kafka/TracingKafkaReceiver.java | 124 ++++++++++++++++++ .../kafka/TracingKafkaReceiverTest.java | 60 +++++++++ ...cingReactorKafkaAutoConfigurationTest.java | 9 -- .../instrument/kafka/KafkaReceiverTest.java | 84 ++++++++++-- 10 files changed, 444 insertions(+), 84 deletions(-) create mode 100644 spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingKafkaReceiverBeanPostProcessor.java create mode 100644 spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/ReactiveKafkaTracingPropagator.java delete mode 100644 spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaConsumerFactory.java create mode 100644 spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiver.java create mode 100644 spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiverTest.java diff --git a/docs/src/main/asciidoc/integrations.adoc b/docs/src/main/asciidoc/integrations.adoc index 9d015425d..583914c8d 100644 --- a/docs/src/main/asciidoc/integrations.adoc +++ b/docs/src/main/asciidoc/integrations.adoc @@ -14,17 +14,33 @@ We decorate the Kafka clients (`KafkaProducer` and `KafkaConsumer`) to create a IMPORTANT: You have to register the `Producer` or `Consumer` as beans in order for Sleuth's auto-configuration to decorate them. When you then inject the beans, the expected type must be `Producer` or `Consumer` (and NOT e.g. `KafkaProducer`). -We also provide `TracingKafkaProducerFactory` and `TracingKafkaConsumerFactory` to be used with the https://projectreactor.io/docs/kafka/release/reference/[Reactor Kafka] clients (`KafkaSender` and `KafkaReceiver`, respectively). See an example in the snippet below: +For use with project reactor we decorate `KafkaReceiver` with `TracingKafkaReceiver` for every bean of that type declared. This will create separate publisher for each element received with its own tracing context propagated. When used with reactor instrumentation you will have access to the context of spans. + +If there's no parent context, it will create just child span with new trace-id. [source,java,indent=0] ---- @Bean -KafkaReceiver reactiveKafkaReceiver(TracingKafkaConsumerFactory tracingKafkaConsumerFactory, KafkaReceiverOptions kafkaReceiverOptions) { - return KafkaReceiver.create(tracingKafkaConsumerFactory, kafkaReceiverOptions); +KafkaReceiver reactiveKafkaReceiver(ReceiverOptions options) { + return KafkaReceiver.create(options); } ---- -Additionally, we decorate any https://docs.spring.io/spring-kafka/docs/current/reference/html/[Spring Kafka] `ProducerFactory` and `ConsumerFactory` available in the context. However, this is disabled if Brave instrumentation is on the classpath. +Later you can simply start receiving your elements to process the context. + +[source,java,indent=0] +---- +@Bean +DisposableBean exampleRunningConsumer(KafkaReceiver receiver){ + reactor.core.Disposable disposable = receiver.receive() + //If you need to read context you can for example use deferContextual + .flatMap(record -> Mono.deferContextual(context -> Mono.just(record))) + .doOnNext(record -> log.info("I will be coorelated to the child span created with parent context from kafka record")) + .subscribe(record -> record.receiverOffset().acknowledge()); + + return disposable::dispose; +} +---- [[sleuth-async-integration]] == Asynchronous Communication diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingKafkaReceiverBeanPostProcessor.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingKafkaReceiverBeanPostProcessor.java new file mode 100644 index 000000000..d1a782a9f --- /dev/null +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingKafkaReceiverBeanPostProcessor.java @@ -0,0 +1,50 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.autoconfig.instrument.kafka; + +import org.jetbrains.annotations.NotNull; +import reactor.kafka.receiver.KafkaReceiver; + +import org.springframework.beans.BeansException; +import org.springframework.beans.factory.config.BeanPostProcessor; +import org.springframework.cloud.sleuth.instrument.kafka.ReactiveKafkaTracingPropagator; +import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaReceiver; + +/** + * Wraps reactors KafkaReceiver in custom TracingKafkaReceiver that provides tracing + * context to tracer and reactor context. Downside operators will keep the context correct + * thanks to + * {@link org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration} + */ +public class TracingKafkaReceiverBeanPostProcessor implements BeanPostProcessor { + + private final ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator; + + public TracingKafkaReceiverBeanPostProcessor(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator) { + this.reactiveKafkaTracingPropagator = reactiveKafkaTracingPropagator; + } + + @Override + public Object postProcessAfterInitialization(@NotNull Object bean, @NotNull String beanName) throws BeansException { + if (bean instanceof KafkaReceiver && !(bean instanceof TracingKafkaReceiver)) { + return new TracingKafkaReceiver<>(reactiveKafkaTracingPropagator, (KafkaReceiver) bean); + } + + return BeanPostProcessor.super.postProcessAfterInitialization(bean, beanName); + } + +} diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingReactorKafkaAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingReactorKafkaAutoConfiguration.java index 49615622f..4ed362f55 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingReactorKafkaAutoConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TracingReactorKafkaAutoConfiguration.java @@ -16,9 +16,13 @@ package org.springframework.cloud.sleuth.autoconfig.instrument.kafka; +import org.apache.kafka.clients.consumer.ConsumerRecord; import reactor.kafka.receiver.KafkaReceiver; +import reactor.kafka.receiver.ReceiverOptions; +import reactor.kafka.receiver.internals.ConsumerFactory; import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.config.BeanPostProcessor; import org.springframework.boot.autoconfigure.AutoConfigureAfter; import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; @@ -26,8 +30,10 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.autoconfig.brave.BraveAutoConfiguration; -import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaConsumerFactory; +import org.springframework.cloud.sleuth.instrument.kafka.ReactiveKafkaTracingPropagator; import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaProducerFactory; +import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaReceiver; +import org.springframework.cloud.sleuth.propagation.Propagator; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -46,16 +52,29 @@ import org.springframework.context.annotation.Configuration; @ConditionalOnProperty(value = "spring.sleuth.kafka.enabled", matchIfMissing = true) public class TracingReactorKafkaAutoConfiguration { + @Bean + ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator(Tracer tracer, Propagator propagator, + Propagator.Getter> extractor) { + return new ReactiveKafkaTracingPropagator(tracer, propagator, extractor); + } + + /** + * This will be wrapping KafkaReceiver beans in tracing wrapper. Can still use it + * manually with + * {@link TracingKafkaReceiver#create(ReactiveKafkaTracingPropagator, ReceiverOptions)} + * {@link TracingKafkaReceiver#create(ReactiveKafkaTracingPropagator, ConsumerFactory, ReceiverOptions)} + */ + @Bean + @ConditionalOnClass({ KafkaReceiver.class }) + static BeanPostProcessor tracingKafkaReceiverBeanPostProcessor( + ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator) { + return new TracingKafkaReceiverBeanPostProcessor(reactiveKafkaTracingPropagator); + } + @Bean @ConditionalOnMissingBean TracingKafkaProducerFactory tracingKafkaProducerFactory(BeanFactory beanFactory) { return new TracingKafkaProducerFactory(beanFactory); } - @Bean - @ConditionalOnMissingBean - TracingKafkaConsumerFactory tracingKafkaConsumerFactory(BeanFactory beanFactory) { - return new TracingKafkaConsumerFactory(beanFactory); - } - } diff --git a/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TraceReactorKafkaAutoConfigurationTests.java b/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TraceReactorKafkaAutoConfigurationTests.java index 4386651e8..cb04d02ab 100644 --- a/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TraceReactorKafkaAutoConfigurationTests.java +++ b/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/kafka/TraceReactorKafkaAutoConfigurationTests.java @@ -16,15 +16,22 @@ package org.springframework.cloud.sleuth.autoconfig.instrument.kafka; +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.MockConsumer; +import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.junit.jupiter.api.Test; import reactor.kafka.receiver.KafkaReceiver; +import reactor.kafka.receiver.ReceiverOptions; +import reactor.kafka.receiver.internals.ConsumerFactory; +import reactor.kafka.receiver.internals.DefaultKafkaReceiver; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.test.context.FilteredClassLoader; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.cloud.sleuth.autoconfig.TraceNoOpAutoConfiguration; -import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaConsumerFactory; +import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaConsumer; import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaProducerFactory; +import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaReceiver; import static org.assertj.core.api.Assertions.assertThat; @@ -38,14 +45,30 @@ class TraceReactorKafkaAutoConfigurationTests { @Test void should_not_create_factories_when_reactor_kafka_not_on_classpath() { this.contextRunner.withClassLoader(new FilteredClassLoader(KafkaReceiver.class)) - .run(context -> assertThat(context).doesNotHaveBean(TracingKafkaProducerFactory.class) - .doesNotHaveBean(TracingKafkaConsumerFactory.class)); + .run(context -> assertThat(context).doesNotHaveBean(TracingKafkaProducerFactory.class)); } @Test void should_create_factories_when_reactor_kafka_on_classpath() { - this.contextRunner.run(context -> assertThat(context).hasSingleBean(TracingKafkaProducerFactory.class) - .hasSingleBean(TracingKafkaConsumerFactory.class)); + this.contextRunner.run(context -> assertThat(context).hasSingleBean(TracingKafkaProducerFactory.class)); + } + + @Test + void should_decorate_kafka_receiver_beans() { + this.contextRunner + .withBean(KafkaReceiver.class, + () -> new DefaultKafkaReceiver<>(new MockConsumerFactory(), ReceiverOptions.create())) + .run(context -> assertThat(context).hasSingleBean(TracingKafkaReceiver.class) + .doesNotHaveBean(TracingKafkaConsumer.class)); + } + + public static class MockConsumerFactory extends ConsumerFactory { + + @Override + public Consumer createConsumer(ReceiverOptions config) { + return new MockConsumer<>(OffsetResetStrategy.NONE); + } + } } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/ReactiveKafkaTracingPropagator.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/ReactiveKafkaTracingPropagator.java new file mode 100644 index 000000000..c6231a673 --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/ReactiveKafkaTracingPropagator.java @@ -0,0 +1,62 @@ +/* + * Copyright 2013-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.kafka; + +import org.apache.kafka.clients.consumer.ConsumerRecord; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.Tracer; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; +import org.springframework.cloud.sleuth.propagation.Propagator; + +/** + * Uses {@link ReactorSleuth} to create separate mono publisher for each element in flux, + * that will be injecting the tracing context to {@link Tracer} and + * {@link reactor.util.context.Context} for each element separately, giving downstream + * operators proper tracing context and span. + * + * @see TracingKafkaReceiver + */ +public class ReactiveKafkaTracingPropagator { + + private final Tracer tracer; + + private final Propagator propagator; + + private final Propagator.Getter> extractor; + + public ReactiveKafkaTracingPropagator(Tracer tracer, Propagator propagator, + Propagator.Getter> extractor) { + this.tracer = tracer; + this.propagator = propagator; + this.extractor = extractor; + } + + public > Flux propagateSpanContextToReactiveContext(Flux publisher) { + return publisher.flatMap(consumerRecord -> Mono.deferContextual((contextView) -> { + Span newSpanWithParent = propagator.extract(consumerRecord, extractor).kind(Span.Kind.CONSUMER) + .name("kafka.consumer").tag("kafka.topic", consumerRecord.topic()) + .tag("kafka.offset", Long.toString(consumerRecord.offset())) + .tag("kafka.partition", Integer.toString(consumerRecord.partition())).start(); + + return ReactorSleuth.tracedMono(tracer, newSpanWithParent, () -> Mono.just(consumerRecord)); + })); + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaConsumerFactory.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaConsumerFactory.java deleted file mode 100644 index 342b669b0..000000000 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaConsumerFactory.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * Copyright 2013-2021 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.cloud.sleuth.instrument.kafka; - -import org.apache.kafka.clients.consumer.Consumer; -import reactor.kafka.receiver.KafkaReceiver; -import reactor.kafka.receiver.ReceiverOptions; -import reactor.kafka.receiver.internals.ConsumerFactory; - -import org.springframework.beans.factory.BeanFactory; - -/** - * This decorates a Reactor Kafka {@link ConsumerFactory} to create decorated consumers of - * type {@link TracingKafkaConsumer}. This can be used by the {@link KafkaReceiver} - * factory methods to create instrumented receivers. - * - * @author Anders Clausen - * @author Flaviu Muresan - * @since 3.1.0 - */ -public class TracingKafkaConsumerFactory extends ConsumerFactory { - - private final BeanFactory beanFactory; - - public TracingKafkaConsumerFactory(BeanFactory beanFactory) { - super(); - this.beanFactory = beanFactory; - } - - @Override - public Consumer createConsumer(ReceiverOptions receiverOptions) { - return new TracingKafkaConsumer<>(super.createConsumer(receiverOptions), beanFactory); - } - -} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiver.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiver.java new file mode 100644 index 000000000..8536363be --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiver.java @@ -0,0 +1,124 @@ +/* + * Copyright 2013-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.kafka; + +import java.util.function.Function; + +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.kafka.receiver.KafkaReceiver; +import reactor.kafka.receiver.ReceiverOptions; +import reactor.kafka.receiver.ReceiverRecord; +import reactor.kafka.receiver.internals.ConsumerFactory; +import reactor.kafka.sender.TransactionManager; + +/** + * Decorator for {@link KafkaReceiver} that delegates most of the work back to original + * consumer, but returns publishers decorated with tracing context per each element. + * + * @author Maciej GromuĊ‚ + * @see ReactiveKafkaTracingPropagator + */ +public class TracingKafkaReceiver implements KafkaReceiver { + + private final ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator; + + private final KafkaReceiver delegate; + + public TracingKafkaReceiver(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator, + KafkaReceiver delegate) { + this.reactiveKafkaTracingPropagator = reactiveKafkaTracingPropagator; + this.delegate = delegate; + } + + /** + * Creates KafkaReceiver that will de decorated by tracing propagator to provide kafka consumer publishing elements + * containing tracing context in their reactor context. + * @param reactiveKafkaTracingPropagator Instance of trace propagation decorator. Should be available in spring application context as a bean. + * @param options Options to pass for underlying {@link KafkaReceiver#create(ReceiverOptions)} + * @param Key of the record + * @param Value of the record + */ + public static KafkaReceiver create(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator, + ReceiverOptions options) { + return new TracingKafkaReceiver<>(reactiveKafkaTracingPropagator, + KafkaReceiver.create(options)); + } + + /** + * Creates KafkaReceiver that will de decorated by tracing propagator to provide kafka consumer publishing elements + * containing tracing context in their reactor context. + * @param reactiveKafkaTracingPropagator Instance of trace propagation decorator. Should be available in spring application context as a bean. + * @param factory Custom factory to provide for underlying {@link KafkaReceiver#create(ConsumerFactory, ReceiverOptions)} + * @param options Options to provide for underlying {@link KafkaReceiver#create(ConsumerFactory, ReceiverOptions)} + * @param Key of the record + * @param Value of the record + */ + public static KafkaReceiver create(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator, + ConsumerFactory factory, ReceiverOptions options) { + return new TracingKafkaReceiver<>(reactiveKafkaTracingPropagator, KafkaReceiver.create(factory, options)); + } + + @Override + public Flux> receive(Integer prefetch) { + return delegate.receive(prefetch) + .transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext); + } + + @Override + public Flux> receive() { + return delegate.receive() + .transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext); + } + + @Override + public Flux>> receiveAutoAck(Integer prefetch) { + return delegate.receiveAutoAck(prefetch) + .map(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext); + } + + @Override + public Flux>> receiveAutoAck() { + return delegate.receiveAutoAck().map(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext); + } + + @Override + public Flux> receiveAtmostOnce(Integer prefetch) { + return delegate.receiveAtmostOnce(prefetch) + .transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext); + } + + @Override + public Flux> receiveAtmostOnce() { + return delegate.receiveAtmostOnce() + .transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext); + } + + @Override + public Flux>> receiveExactlyOnce(TransactionManager transactionManager, + Integer prefetch) { + return delegate.receiveExactlyOnce(transactionManager, prefetch); + } + + @Override + public Mono doOnConsumer(Function, ? extends T> function) { + return delegate.doOnConsumer(function); + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiverTest.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiverTest.java new file mode 100644 index 000000000..272fe36a0 --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaReceiverTest.java @@ -0,0 +1,60 @@ +/* + * Copyright 2013-2021 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.kafka; + +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +import reactor.core.publisher.Flux; +import reactor.kafka.receiver.KafkaReceiver; +import reactor.kafka.receiver.ReceiverOffset; +import reactor.kafka.receiver.ReceiverRecord; +import reactor.test.StepVerifier; + +@ExtendWith(MockitoExtension.class) +public class TracingKafkaReceiverTest { + + @Mock + KafkaReceiver sourceReceiver; + + @Mock + ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator; + + @Test + void should_wrap_delegate_kafka_receiver() { + + ReceiverOffset offset = Mockito.mock(ReceiverOffset.class); + TracingKafkaReceiver tracingReceiverTest = new TracingKafkaReceiver<>( + reactiveKafkaTracingPropagator, sourceReceiver); + + String key = "foo"; + String value = "bar"; + + Flux> recordsPublisher = Flux + .just(new ReceiverRecord<>(new ConsumerRecord<>("topic", 0, 0, key, value), offset)); + + Mockito.when(sourceReceiver.receive()).thenReturn(recordsPublisher); + Mockito.when(reactiveKafkaTracingPropagator.propagateSpanContextToReactiveContext(Mockito.any())) + .thenAnswer(invocation -> invocation.getArguments()[0]); + + StepVerifier.create(tracingReceiverTest.receive()).expectNextCount(1).expectComplete().verify(); + } + +} diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/TracingReactorKafkaAutoConfigurationTest.java b/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/TracingReactorKafkaAutoConfigurationTest.java index 4992bc281..7243d45c8 100644 --- a/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/TracingReactorKafkaAutoConfigurationTest.java +++ b/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/TracingReactorKafkaAutoConfigurationTest.java @@ -21,7 +21,6 @@ import org.junit.jupiter.api.Test; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.SpringBootTest; -import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaConsumerFactory; import org.springframework.cloud.sleuth.instrument.kafka.TracingKafkaProducerFactory; import org.springframework.context.annotation.Configuration; @@ -31,17 +30,9 @@ import static org.assertj.core.api.BDDAssertions.then; webEnvironment = SpringBootTest.WebEnvironment.NONE) public class TracingReactorKafkaAutoConfigurationTest { - @Autowired - TracingKafkaConsumerFactory kafkaConsumerFactory; - @Autowired TracingKafkaProducerFactory kafkaProducerFactory; - @Test - public void should_register_consumer_factory() { - then(this.kafkaConsumerFactory).isNotNull(); - } - @Test public void should_register_producer_factory() { then(this.kafkaProducerFactory).isNotNull(); diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/KafkaReceiverTest.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/KafkaReceiverTest.java index 897aafd0f..549a9c818 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/KafkaReceiverTest.java +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/KafkaReceiverTest.java @@ -37,29 +37,38 @@ import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Answers; -import org.mockito.BDDMockito; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.reactivestreams.Publisher; import org.testcontainers.containers.KafkaContainer; import org.testcontainers.containers.wait.strategy.Wait; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; import org.testcontainers.utility.DockerImageName; import reactor.core.Disposable; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Hooks; +import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; import reactor.kafka.receiver.KafkaReceiver; import reactor.kafka.receiver.ReceiverOptions; +import reactor.kafka.receiver.ReceiverRecord; +import reactor.test.StepVerifier; +import reactor.util.context.ContextView; import org.springframework.beans.factory.BeanFactory; +import org.springframework.cloud.sleuth.CurrentTraceContext; import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.exporter.FinishedSpan; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; import org.springframework.cloud.sleuth.propagation.Propagator; import org.springframework.cloud.sleuth.test.TestSpanHandler; import org.springframework.cloud.sleuth.test.TestTracingAwareSupplier; -import org.springframework.core.ParameterizedTypeReference; -import org.springframework.core.ResolvableType; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; +import static org.assertj.core.api.Assertions.assertThat; import static org.awaitility.Awaitility.await; @Testcontainers @@ -67,6 +76,8 @@ import static org.awaitility.Awaitility.await; @Tag("DockerRequired") public abstract class KafkaReceiverTest implements TestTracingAwareSupplier { + static final String HOOK_KEY = "org.springframework.cloud.sleuth.autoconfig.instrument.reactor.TraceReactorAutoConfiguration.TraceReactorConfiguration"; + protected String testTopic; protected Tracer tracer = tracerTest().tracing().tracer(); @@ -75,10 +86,20 @@ public abstract class KafkaReceiverTest implements TestTracingAwareSupplier { protected TestSpanHandler spans = tracerTest().handler(); + protected CurrentTraceContext currentTraceContext = tracerTest().tracing().currentTraceContext(); + + protected Propagator.Getter> extractor = new TracingKafkaPropagatorGetter(); + private Disposable consumerSubscription; + private Disposable shareableReceiverDisposable; + + protected Flux> shareableReceiver; + protected final AtomicInteger receivedCounter = new AtomicInteger(0); + AnnotationConfigApplicationContext springContext = new AnnotationConfigApplicationContext(); + @Mock(answer = Answers.RETURNS_DEEP_STUBS) BeanFactory beanFactory; @@ -99,10 +120,11 @@ public abstract class KafkaReceiverTest implements TestTracingAwareSupplier { @BeforeEach void setup() { - BDDMockito.given(this.beanFactory.getBean(Propagator.class)).willReturn(this.propagator); - BDDMockito.given(this.beanFactory.getBeanProvider(ResolvableType.forClassWithGenerics(Propagator.Getter.class, - ResolvableType.forType(new ParameterizedTypeReference>() { - }))).getIfAvailable()).willReturn(new TracingKafkaPropagatorGetter()); + // We need to enable scope passing to actually see the context downstream + Hooks.resetOnEachOperator(); + Hooks.resetOnLastOperator(); + Schedulers.resetOnScheduleHooks(); + testTopic = UUID.randomUUID().toString(); Map consumerProperties = new HashMap<>(); consumerProperties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaContainer.getBootstrapServers()); @@ -111,18 +133,35 @@ public abstract class KafkaReceiverTest implements TestTracingAwareSupplier { consumerProperties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); consumerProperties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); ReceiverOptions options = ReceiverOptions.create(consumerProperties); + options = options.withKeyDeserializer(new StringDeserializer()).withValueDeserializer(new StringDeserializer()) .subscription(Collections.singletonList(testTopic)); - KafkaReceiver kafkaReceiver = KafkaReceiver.create(new TracingKafkaConsumerFactory(beanFactory), - options); - this.consumerSubscription = kafkaReceiver.receive().subscribeOn(Schedulers.single()) + + KafkaReceiver kafkaReceiver = new TracingKafkaReceiver<>( + new ReactiveKafkaTracingPropagator(tracer, propagator, extractor), KafkaReceiver.create(options)); + + // Create shareable receiver + this.shareableReceiver = kafkaReceiver.receive().publish().autoConnect(0, + disposable -> this.shareableReceiverDisposable = disposable); + + // Create subscription for previous tests compatibility + this.consumerSubscription = shareableReceiver.subscribeOn(Schedulers.single()) .subscribe(record -> receivedCounter.incrementAndGet()); + this.receivedCounter.set(0); } @AfterEach void destroy() { + springContext.close(); + Hooks.resetOnEachOperator(); + Hooks.resetOnLastOperator(); + Schedulers.resetOnScheduleHooks(); + this.consumerSubscription.dispose(); + if (this.shareableReceiverDisposable != null) { + this.shareableReceiverDisposable.dispose(); + } } @Test @@ -144,6 +183,31 @@ public abstract class KafkaReceiverTest implements TestTracingAwareSupplier { BDDAssertions.then(span.getTags().get("kafka.partition")).isEqualTo("0"); } + @Test + public void should_pass_tracing_context_for_consumers() { + springContext.registerBean(Tracer.class, () -> this.tracer); + springContext.registerBean(CurrentTraceContext.class, () -> this.currentTraceContext); + springContext.refresh(); + + Hooks.onEachOperator(HOOK_KEY, ReactorSleuth.onEachOperatorForOnEachInstrumentation(springContext)); + Hooks.onLastOperator(HOOK_KEY, ReactorSleuth.onLastOperatorForOnEachInstrumentation(springContext)); + + KafkaProducer kafkaProducer = KafkaTestUtils + .buildTestKafkaProducer(kafkaContainer.getBootstrapServers()); + ProducerRecord producerRecord = new ProducerRecord<>(testTopic, "test", "test-with-trace"); + producerRecord.headers().add("b3", "80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1".getBytes()); + kafkaProducer.send(producerRecord); + + Publisher testFlux = shareableReceiver.flatMap(record -> Mono.deferContextual(Mono::just)); + + StepVerifier.create(testFlux).assertNext(contextView -> { + TraceContext traceContext = contextView.get(TraceContext.class); + + assertThat(traceContext).returns("80f198ee56343ba864fe8b2a57d3eff7", TraceContext::traceId) + .returns("e457b5a2e4d86bd1", TraceContext::parentId); + }).thenCancel().verify(Duration.ofSeconds(15)); + } + @Override public void cleanUpTracing() { this.spans.clear();