Added a wrapper for KafkaPlatformTransactionManager; fixes gh-2063

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceKafkaPlatformTransactionManagerBeanPostProcessor.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.tx;
+
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.cloud.sleuth.instrument.tx.TraceKafkaAwareTransactionManager;
+import org.springframework.kafka.transaction.KafkaAwareTransactionManager;
+
+/**
+ * Post processor that wraps a {@link KafkaAwareTransactionManager}.
+ *
+ * @author Marcin Grzejszczak
+ * @since 3.1.0
+ */
+public class TraceKafkaPlatformTransactionManagerBeanPostProcessor implements BeanPostProcessor {
+
+	private final BeanFactory beanFactory;
+
+	private final TracePlatformTransactionManagerBeanPostProcessor delegate;
+
+	public TraceKafkaPlatformTransactionManagerBeanPostProcessor(BeanFactory beanFactory) {
+		this.beanFactory = beanFactory;
+		this.delegate = new TracePlatformTransactionManagerBeanPostProcessor(beanFactory);
+	}
+
+	@Override
+	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
+		if (bean instanceof KafkaAwareTransactionManager && !(bean instanceof TraceKafkaAwareTransactionManager)) {
+			return new TraceKafkaAwareTransactionManager((KafkaAwareTransactionManager) bean, this.beanFactory);
+		}
+		return this.delegate.postProcessAfterInitialization(bean, beanName);
+	}
+
+}

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfiguration.java

@@ -20,6 +20,7 @@ import org.springframework.beans.factory.BeanFactory;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.autoconfig.brave.BraveAutoConfiguration;
@@ -41,11 +42,20 @@ public class TraceTxAutoConfiguration {
 
 	@Bean
 	@ConditionalOnClass(name = "org.springframework.transaction.PlatformTransactionManager")
+	@ConditionalOnMissingClass("org.springframework.kafka.transaction.KafkaAwareTransactionManager")
 	static TracePlatformTransactionManagerBeanPostProcessor tracePlatformTransactionManagerBeanPostProcessor(
 			BeanFactory beanFactory) {
 		return new TracePlatformTransactionManagerBeanPostProcessor(beanFactory);
 	}
 
+	@Bean
+	@ConditionalOnClass(name = { "org.springframework.transaction.PlatformTransactionManager",
+			"org.springframework.kafka.transaction.KafkaAwareTransactionManager" })
+	static TraceKafkaPlatformTransactionManagerBeanPostProcessor traceKafkaPlatformTransactionManagerBeanPostProcessor(
+			BeanFactory beanFactory) {
+		return new TraceKafkaPlatformTransactionManagerBeanPostProcessor(beanFactory);
+	}
+
 	@Bean
 	@ConditionalOnClass(
 			name = { "org.springframework.transaction.ReactiveTransactionManager", "reactor.core.publisher.Mono" })

spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfigurationTests.java

@@ -36,10 +36,22 @@ class TraceTxAutoConfigurationTests {
 	@Test
 	void should_register_bean_post_processors() {
 		this.contextRunner.run(context -> Assertions.assertThat(context)
-				.hasSingleBean(TracePlatformTransactionManagerBeanPostProcessor.class)
+				.hasSingleBean(TraceKafkaPlatformTransactionManagerBeanPostProcessor.class)
+				.doesNotHaveBean(TracePlatformTransactionManagerBeanPostProcessor.class)
 				.hasSingleBean(TraceReactiveTransactionManagerBeanPostProcessor.class));
 	}
 
+	@Test
+	void should_register_non_kafka_bean_post_processors_when_kafka_not_on_classpath() {
+		this.contextRunner
+				.withClassLoader(
+						new FilteredClassLoader("org.springframework.kafka.transaction.KafkaAwareTransactionManager"))
+				.run(context -> Assertions.assertThat(context)
+						.doesNotHaveBean(TraceKafkaPlatformTransactionManagerBeanPostProcessor.class)
+						.hasSingleBean(TracePlatformTransactionManagerBeanPostProcessor.class)
+						.hasSingleBean(TraceReactiveTransactionManagerBeanPostProcessor.class));
+	}
+
 	@Test
 	void should_not_register_bean_post_processor_when_tx_not_on_classpath() {
 		this.contextRunner.withClassLoader(new FilteredClassLoader(PlatformTransactionManager.class))

spring-cloud-sleuth-instrumentation/pom.xml

@@ -116,6 +116,11 @@
 			<artifactId>spring-boot-starter-websocket</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework.kafka</groupId>
+			<artifactId>spring-kafka</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
 			<artifactId>spring-boot-starter-batch</artifactId>

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TraceKafkaAwareTransactionManager.java

@@ -0,0 +1,41 @@
+/*
+ * 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.tx;
+
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.kafka.core.ProducerFactory;
+import org.springframework.kafka.transaction.KafkaAwareTransactionManager;
+
+/**
+ * A trace representation of a {@link KafkaAwareTransactionManager}.
+ *
+ * @author Marcin Grzejszczak
+ * @since 3.1.0
+ */
+public class TraceKafkaAwareTransactionManager extends TracePlatformTransactionManager
+		implements KafkaAwareTransactionManager {
+
+	public TraceKafkaAwareTransactionManager(KafkaAwareTransactionManager delegate, BeanFactory beanFactory) {
+		super(delegate, beanFactory);
+	}
+
+	@Override
+	public ProducerFactory getProducerFactory() {
+		return ((KafkaAwareTransactionManager) this.delegate).getProducerFactory();
+	}
+
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManager.java

@@ -41,7 +41,7 @@ public class TracePlatformTransactionManager implements PlatformTransactionManag
 
 	private static final Log log = LogFactory.getLog(TracePlatformTransactionManager.class);
 
-	private final PlatformTransactionManager delegate;
+	protected final PlatformTransactionManager delegate;
 
 	private final BeanFactory beanFactory;
 

spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TraceKafkaPlatformTransactionManagerTests.java

@@ -0,0 +1,40 @@
+/*
+ * 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.tx;
+
+import org.mockito.Mockito;
+
+import org.springframework.kafka.transaction.KafkaAwareTransactionManager;
+
+class TraceKafkaPlatformTransactionManagerTests extends TracePlatformTransactionManagerTests {
+
+	KafkaAwareTransactionManager delegate = Mockito.mock(KafkaAwareTransactionManager.class);
+
+	@Override
+	TraceKafkaAwareTransactionManager manager() {
+		final TraceKafkaAwareTransactionManager manager = new TraceKafkaAwareTransactionManager(getDelegate(),
+				beanFactory());
+		manager.initialize();
+		return manager;
+	}
+
+	@Override
+	KafkaAwareTransactionManager getDelegate() {
+		return this.delegate;
+	}
+
+}

spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManagerTests.java

@@ -108,7 +108,7 @@ class TracePlatformTransactionManagerTests {
 	@Test
 	void should_report_a_fallback_span_when_exception_occurred_while_getting_transaction() {
 		TracePlatformTransactionManager manager = manager();
-		BDDMockito.given(this.delegate.getTransaction(BDDMockito.any()))
+		BDDMockito.given(getDelegate().getTransaction(BDDMockito.any()))
 				.willThrow(new TransactionTimedOutException("boom"));
 		thenThreadLocalIsClear(manager);
 
@@ -121,7 +121,7 @@ class TracePlatformTransactionManagerTests {
 	@Test
 	void should_report_a_span_when_exception_occurred_while_committing_transaction() {
 		TracePlatformTransactionManager manager = manager();
-		BDDMockito.willThrow(new TransactionTimedOutException("boom")).given(this.delegate).commit(BDDMockito.any());
+		BDDMockito.willThrow(new TransactionTimedOutException("boom")).given(getDelegate()).commit(BDDMockito.any());
 		threadLocalSpan(manager);
 
 		BDDAssertions.thenThrownBy(() -> manager.commit(null)).isInstanceOf(TransactionTimedOutException.class);
@@ -135,7 +135,7 @@ class TracePlatformTransactionManagerTests {
 	@Test
 	void should_report_a_span_when_exception_occurred_while_rolling_back_transaction() {
 		TracePlatformTransactionManager manager = manager();
-		BDDMockito.willThrow(new TransactionTimedOutException("boom")).given(this.delegate).rollback(BDDMockito.any());
+		BDDMockito.willThrow(new TransactionTimedOutException("boom")).given(getDelegate()).rollback(BDDMockito.any());
 		threadLocalSpan(manager);
 
 		BDDAssertions.thenThrownBy(() -> manager.rollback(null)).isInstanceOf(TransactionTimedOutException.class);
@@ -147,7 +147,7 @@ class TracePlatformTransactionManagerTests {
 	}
 
 	private void setupTransactionStatusWithNewTransactionStatusEqualTo(boolean transactionStatus) {
-		BDDMockito.given(this.delegate.getTransaction(BDDMockito.any()))
+		BDDMockito.given(getDelegate().getTransaction(BDDMockito.any()))
 				.willReturn(new SimpleTransactionStatus(transactionStatus));
 	}
 
@@ -165,15 +165,15 @@ class TracePlatformTransactionManagerTests {
 		then(tracer.getOnlySpan()).as("The previously created span was reported").isSameAs(firstSpan);
 	}
 
-	private TracePlatformTransactionManager manager() {
-		final TracePlatformTransactionManager manager = new TracePlatformTransactionManager(this.delegate,
+	TracePlatformTransactionManager manager() {
+		final TracePlatformTransactionManager manager = new TracePlatformTransactionManager(getDelegate(),
 				beanFactory());
 		manager.initialize();
 		return manager;
 	}
 
 	private TracePlatformTransactionManager managerWithManualFallback() {
-		final TracePlatformTransactionManager manager = new TracePlatformTransactionManager(this.delegate,
+		final TracePlatformTransactionManager manager = new TracePlatformTransactionManager(getDelegate(),
 				beanFactory()) {
 			@Override
 			Span fallbackSpan() {
@@ -193,10 +193,14 @@ class TracePlatformTransactionManagerTests {
 		return span;
 	}
 
-	private BeanFactory beanFactory() {
+	BeanFactory beanFactory() {
 		StaticListableBeanFactory beanFactory = new StaticListableBeanFactory();
 		beanFactory.addBean("tracer", tracer);
 		return beanFactory;
 	}
 
+	PlatformTransactionManager getDelegate() {
+		return this.delegate;
+	}
+
 }

tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/KafkaSenderTest.java

@@ -23,6 +23,7 @@ import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.header.Header;
 import org.assertj.core.api.BDDAssertions;
+import org.awaitility.Awaitility;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -50,15 +51,18 @@ public class KafkaSenderTest extends org.springframework.cloud.sleuth.instrument
 		ProducerRecord<String, String> producerRecord = new ProducerRecord<>(testTopic, "test", "test");
 		startKafkaConsumer();
 
-		Flux<SenderResult<Object>> senderResultFlux = this.kafkaSender
-				.send(Mono.just(SenderRecord.create(producerRecord, null)));
-		StepVerifier.create(senderResultFlux).expectNextCount(1).verifyComplete();
-		ConsumerRecord<String, String> consumerRecord = consumerRecords.poll(15, TimeUnit.SECONDS);
+			Flux<SenderResult<Object>> senderResultFlux = this.kafkaSender
+					.send(Mono.just(SenderRecord.create(producerRecord, null)));
+			StepVerifier.create(senderResultFlux).expectNextCount(1).verifyComplete();
 
-		BDDAssertions.then(consumerRecord).isNotNull();
-		BDDAssertions.then(getHeaderValueOrNull(consumerRecord, "X-B3-TraceId")).isNotNull();
-		BDDAssertions.then(getHeaderValueOrNull(consumerRecord, "X-B3-SpanId")).isNotNull();
-		BDDAssertions.then(getHeaderValueOrNull(consumerRecord, "X-B3-Sampled")).isNotNull();
+		Awaitility.await().untilAsserted(() -> {
+			ConsumerRecord<String, String> consumerRecord = consumerRecords.poll(1, TimeUnit.SECONDS);
+
+			BDDAssertions.then(consumerRecord).isNotNull();
+			BDDAssertions.then(getHeaderValueOrNull(consumerRecord, "X-B3-TraceId")).isNotNull();
+			BDDAssertions.then(getHeaderValueOrNull(consumerRecord, "X-B3-SpanId")).isNotNull();
+			BDDAssertions.then(getHeaderValueOrNull(consumerRecord, "X-B3-Sampled")).isNotNull();
+		});
 	}
 
 	private static String getHeaderValueOrNull(ConsumerRecord<?, ?> consumerRecord, String header) {