diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceKafkaPlatformTransactionManagerBeanPostProcessor.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceKafkaPlatformTransactionManagerBeanPostProcessor.java new file mode 100644 index 000000000..d590d9b79 --- /dev/null +++ b/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); + } + +} diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfiguration.java index f1fd9b189..bf24fd56f 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfiguration.java +++ b/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" }) diff --git a/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfigurationTests.java b/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfigurationTests.java index 04f3e5acd..ceddc8905 100644 --- a/spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/tx/TraceTxAutoConfigurationTests.java +++ b/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)) diff --git a/spring-cloud-sleuth-instrumentation/pom.xml b/spring-cloud-sleuth-instrumentation/pom.xml index 37405dbc4..40306e9ed 100644 --- a/spring-cloud-sleuth-instrumentation/pom.xml +++ b/spring-cloud-sleuth-instrumentation/pom.xml @@ -116,6 +116,11 @@ spring-boot-starter-websocket true + + org.springframework.kafka + spring-kafka + true + org.springframework.boot spring-boot-starter-batch diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TraceKafkaAwareTransactionManager.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TraceKafkaAwareTransactionManager.java new file mode 100644 index 000000000..c20099b61 --- /dev/null +++ b/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(); + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManager.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManager.java index c12f605c8..9cd9d6367 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManager.java +++ b/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; diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TraceKafkaPlatformTransactionManagerTests.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TraceKafkaPlatformTransactionManagerTests.java new file mode 100644 index 000000000..a5ea8b9d8 --- /dev/null +++ b/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; + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManagerTests.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManagerTests.java index c7b47457b..5f2e21732 100644 --- a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/tx/TracePlatformTransactionManagerTests.java +++ b/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; + } + } diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/KafkaSenderTest.java b/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/KafkaSenderTest.java index 52b763679..e64f2b057 100644 --- a/tests/brave/spring-cloud-sleuth-instrumentation-kafka-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/kafka/KafkaSenderTest.java +++ b/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 producerRecord = new ProducerRecord<>(testTopic, "test", "test"); startKafkaConsumer(); - Flux> senderResultFlux = this.kafkaSender - .send(Mono.just(SenderRecord.create(producerRecord, null))); - StepVerifier.create(senderResultFlux).expectNextCount(1).verifyComplete(); - ConsumerRecord consumerRecord = consumerRecords.poll(15, TimeUnit.SECONDS); + Flux> 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 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) {