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

@@ -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<K,V>` with `TracingKafkaReceiver<K,V>` 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<String, String> reactiveKafkaReceiver(TracingKafkaConsumerFactory tracingKafkaConsumerFactory, KafkaReceiverOptions kafkaReceiverOptions) {
-    return KafkaReceiver.create(tracingKafkaConsumerFactory, kafkaReceiverOptions);
+KafkaReceiver<K, V> reactiveKafkaReceiver(ReceiverOptions<K,V> 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<String, String> 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

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);
+	}
+
+}

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<ConsumerRecord<?, ?>> 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);
-	}
-
 }

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 <K, V> Consumer<K, V> createConsumer(ReceiverOptions<K, V> config) {
+			return new MockConsumer<>(OffsetResetStrategy.NONE);
+		}
+
 	}
 
 }

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<ConsumerRecord<?, ?>> extractor;
+
+	public ReactiveKafkaTracingPropagator(Tracer tracer, Propagator propagator,
+			Propagator.Getter<ConsumerRecord<?, ?>> extractor) {
+		this.tracer = tracer;
+		this.propagator = propagator;
+		this.extractor = extractor;
+	}
+
+	public <K, V, T extends ConsumerRecord<K, V>> Flux<T> propagateSpanContextToReactiveContext(Flux<T> 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));
+		}));
+	}
+
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/kafka/TracingKafkaConsumerFactory.java

@@ -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 <K, V> Consumer<K, V> createConsumer(ReceiverOptions<K, V> receiverOptions) {
-		return new TracingKafkaConsumer<>(super.createConsumer(receiverOptions), beanFactory);
-	}
-
-}

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<K, V> implements KafkaReceiver<K, V> {
+
+	private final ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator;
+
+	private final KafkaReceiver<K, V> delegate;
+
+	public TracingKafkaReceiver(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator,
+			KafkaReceiver<K, V> 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 <K> Key of the record
+	 * @param <V> Value of the record
+	 */
+	public static <K, V> KafkaReceiver<K, V> create(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator,
+			ReceiverOptions<K, V> 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 <K> Key of the record
+	 * @param <V> Value of the record
+	 */
+	public static <K, V> KafkaReceiver<K, V> create(ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator,
+			ConsumerFactory factory, ReceiverOptions<K, V> options) {
+		return new TracingKafkaReceiver<>(reactiveKafkaTracingPropagator, KafkaReceiver.create(factory, options));
+	}
+
+	@Override
+	public Flux<ReceiverRecord<K, V>> receive(Integer prefetch) {
+		return delegate.receive(prefetch)
+				.transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext);
+	}
+
+	@Override
+	public Flux<ReceiverRecord<K, V>> receive() {
+		return delegate.receive()
+				.transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext);
+	}
+
+	@Override
+	public Flux<Flux<ConsumerRecord<K, V>>> receiveAutoAck(Integer prefetch) {
+		return delegate.receiveAutoAck(prefetch)
+				.map(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext);
+	}
+
+	@Override
+	public Flux<Flux<ConsumerRecord<K, V>>> receiveAutoAck() {
+		return delegate.receiveAutoAck().map(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext);
+	}
+
+	@Override
+	public Flux<ConsumerRecord<K, V>> receiveAtmostOnce(Integer prefetch) {
+		return delegate.receiveAtmostOnce(prefetch)
+				.transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext);
+	}
+
+	@Override
+	public Flux<ConsumerRecord<K, V>> receiveAtmostOnce() {
+		return delegate.receiveAtmostOnce()
+				.transformDeferred(reactiveKafkaTracingPropagator::propagateSpanContextToReactiveContext);
+	}
+
+	@Override
+	public Flux<Flux<ConsumerRecord<K, V>>> receiveExactlyOnce(TransactionManager transactionManager,
+			Integer prefetch) {
+		return delegate.receiveExactlyOnce(transactionManager, prefetch);
+	}
+
+	@Override
+	public <T> Mono<T> doOnConsumer(Function<Consumer<K, V>, ? extends T> function) {
+		return delegate.doOnConsumer(function);
+	}
+
+}

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<String, String> sourceReceiver;
+
+	@Mock
+	ReactiveKafkaTracingPropagator reactiveKafkaTracingPropagator;
+
+	@Test
+	void should_wrap_delegate_kafka_receiver() {
+
+		ReceiverOffset offset = Mockito.mock(ReceiverOffset.class);
+		TracingKafkaReceiver<String, String> tracingReceiverTest = new TracingKafkaReceiver<>(
+				reactiveKafkaTracingPropagator, sourceReceiver);
+
+		String key = "foo";
+		String value = "bar";
+
+		Flux<ReceiverRecord<String, String>> 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();
+	}
+
+}

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();

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<ConsumerRecord<?, ?>> extractor = new TracingKafkaPropagatorGetter();
+
 	private Disposable consumerSubscription;
 
+	private Disposable shareableReceiverDisposable;
+
+	protected Flux<ReceiverRecord<String, String>> 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<ConsumerRecord<?, ?>>() {
-				}))).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<String, Object> 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<String, String> options = ReceiverOptions.create(consumerProperties);
+
 		options = options.withKeyDeserializer(new StringDeserializer()).withValueDeserializer(new StringDeserializer())
 				.subscription(Collections.singletonList(testTopic));
-		KafkaReceiver<String, String> kafkaReceiver = KafkaReceiver.create(new TracingKafkaConsumerFactory(beanFactory),
-				options);
-		this.consumerSubscription = kafkaReceiver.receive().subscribeOn(Schedulers.single())
+
+		KafkaReceiver<String, String> 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<String, String> kafkaProducer = KafkaTestUtils
+				.buildTestKafkaProducer(kafkaContainer.getBootstrapServers());
+		ProducerRecord<String, String> producerRecord = new ProducerRecord<>(testTopic, "test", "test-with-trace");
+		producerRecord.headers().add("b3", "80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1".getBytes());
+		kafkaProducer.send(producerRecord);
+
+		Publisher<ContextView> 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();