Get KafkaStreamsTracing via BeanFactory to prevent eager initialization (#1623)

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/SleuthKafkaStreamsConfiguration.java

@@ -23,6 +23,7 @@ import org.apache.commons.logging.LogFactory;
 import org.apache.kafka.streams.KafkaStreams;
 
 import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
@@ -44,8 +45,7 @@ import org.springframework.kafka.config.StreamsBuilderFactoryBean;
 @ConditionalOnBean(Tracing.class)
 @AutoConfigureAfter({ TraceAutoConfiguration.class })
 @OnMessagingEnabled
-@ConditionalOnProperty(value = "spring.sleuth.messaging.kafka.streams.enabled",
-		matchIfMissing = true)
+@ConditionalOnProperty(value = "spring.sleuth.messaging.kafka.streams.enabled", matchIfMissing = true)
 @ConditionalOnClass(KafkaStreams.class)
 public class SleuthKafkaStreamsConfiguration {
 
@@ -55,6 +55,7 @@ public class SleuthKafkaStreamsConfiguration {
 	/**
 	 * Expose {@link KafkaStreamsTracing} as bean to allow for filter/map/peek/transform
 	 * operations.
+	 *
 	 * @param tracing Brave Tracing instance from TraceAutoConfiguration
 	 * @return instance for use in further manual instrumentation
 	 */
@@ -64,10 +65,17 @@ public class SleuthKafkaStreamsConfiguration {
 		return KafkaStreamsTracing.create(tracing);
 	}
 
+	/**
+	 * Call {@link StreamsBuilderFactoryBean#setClientSupplier(org.apache.kafka.streams.KafkaClientSupplier)} with
+	 * Brave's TracingKafkaClientSupplier.
+	 *
+	 * @param objectProvider provides KafkaStreamsTracing; prevents eager initialization
+	 * @return
+	 */
 	@Bean
 	static KafkaStreamsBuilderFactoryBeanPostProcessor kafkaStreamsBuilderFactoryBeanPostProcessor(
-			KafkaStreamsTracing kafkaStreamsTracing) {
-		return new KafkaStreamsBuilderFactoryBeanPostProcessor(kafkaStreamsTracing);
+			ObjectProvider<KafkaStreamsTracing> objectProvider) {
+		return new KafkaStreamsBuilderFactoryBeanPostProcessor(objectProvider);
 	}
 
 }
@@ -83,25 +91,23 @@ public class SleuthKafkaStreamsConfiguration {
  */
 class KafkaStreamsBuilderFactoryBeanPostProcessor implements BeanPostProcessor {
 
-	private static final Log log = LogFactory
-			.getLog(KafkaStreamsBuilderFactoryBeanPostProcessor.class);
+	private static final Log log = LogFactory.getLog(KafkaStreamsBuilderFactoryBeanPostProcessor.class);
 
-	private final KafkaStreamsTracing kafkaStreamsTracing;
+	private final ObjectProvider<KafkaStreamsTracing> objectProvider;
 
-	KafkaStreamsBuilderFactoryBeanPostProcessor(KafkaStreamsTracing kafkaStreamsTracing) {
-		this.kafkaStreamsTracing = kafkaStreamsTracing;
+	KafkaStreamsBuilderFactoryBeanPostProcessor(ObjectProvider<KafkaStreamsTracing> objectProvider) {
+		this.objectProvider = objectProvider;
 	}
 
 	@Override
-	public Object postProcessAfterInitialization(Object bean, String beanName)
-			throws BeansException {
+	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
 		if (bean instanceof StreamsBuilderFactoryBean) {
-			StreamsBuilderFactoryBean sbfb = (StreamsBuilderFactoryBean) bean;
+			// KafkaStreamsTracing is created in SleuthKafkaStreamsConfiguration above, so should not be null here
+			KafkaStreamsTracing kafkaStreamsTracing = this.objectProvider.getIfAvailable();
+			((StreamsBuilderFactoryBean) bean).setClientSupplier(kafkaStreamsTracing.kafkaClientSupplier());
 			if (log.isDebugEnabled()) {
-				log.debug(
-						"StreamsBuilderFactoryBean bean is auto-configured to enable tracing.");
+				log.debug("StreamsBuilderFactoryBean bean is auto-configured to enable tracing.");
 			}
-			sbfb.setClientSupplier(kafkaStreamsTracing.kafkaClientSupplier());
 		}
 		return bean;
 	}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/SleuthKafkaStreamsConfigurationIntegrationTests.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright 2013-2019 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.messaging;
+
+import javax.annotation.PostConstruct;
+
+import brave.Tracing;
+import brave.kafka.streams.KafkaStreamsTracing;
+import org.apache.kafka.streams.KafkaClientSupplier;
+import org.apache.kafka.streams.KafkaStreams;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.FilteredClassLoader;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.boot.test.system.CapturedOutput;
+import org.springframework.boot.test.system.OutputCaptureExtension;
+import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.config.StreamsBuilderFactoryBean;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+@ExtendWith(OutputCaptureExtension.class)
+class SleuthKafkaStreamsConfigurationIntegrationTests {
+
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+			.withConfiguration(AutoConfigurations.of(
+					TraceAutoConfiguration.class,
+					SleuthKafkaStreamsConfiguration.class))
+			.withUserConfiguration(UserConfig.class);
+
+	@Test
+	void should_create_KafkaStreamsTracing() {
+		this.contextRunner
+				.run(context -> assertThat(context).hasSingleBean(KafkaStreamsTracing.class));
+	}
+
+	@Test
+	void should_not_create_KafkaStreamsTracing_when_KafkaStreams_not_present() {
+		this.contextRunner
+				.withClassLoader(new FilteredClassLoader(KafkaStreams.class))
+				.run(context -> assertThat(context).doesNotHaveBean(KafkaStreamsTracing.class));
+	}
+
+	@Test
+	void should_not_create_KafkaStreamsTracing_when_kafkastreams_disabled() {
+		this.contextRunner
+				.withPropertyValues("spring.sleuth.messaging.kafka.streams.enabled=false")
+				.run(context -> assertThat(context).doesNotHaveBean(KafkaStreamsTracing.class));
+	}
+
+	@Test
+	void should_not_create_KafkaStreamsTracing_when_messaging_disabled() {
+		this.contextRunner
+				.withPropertyValues("spring.sleuth.messaging.enabled=false")
+				.run(context -> assertThat(context).doesNotHaveBean(KafkaStreamsTracing.class));
+	}
+
+	@Test
+	void should_set_KafkaClientSupplier_on_StreamsBuilderFactoryBean() {
+		this.contextRunner
+				.run(context -> verify(UserConfig.streamsBuilderFactoryBean)
+						.setClientSupplier(any(KafkaClientSupplier.class)));
+	}
+
+	@Test
+	void should_not_complain_about_eager_initialization() {
+		this.contextRunner
+				.withUserConfiguration(EagerInitializationConfig.class)
+				.run(context -> verify(UserConfig.streamsBuilderFactoryBean)
+						.setClientSupplier(any(KafkaClientSupplier.class)));
+	}
+
+	@AfterEach
+	void afterEach(CapturedOutput output) {
+		assertThat(output).doesNotContain("is not eligible for getting processed by all BeanPostProcessors");
+	}
+
+	@Configuration
+	static class UserConfig {
+		static StreamsBuilderFactoryBean streamsBuilderFactoryBean;
+
+		@Bean
+		StreamsBuilderFactoryBean streamsBuilderFactoryBean() {
+			streamsBuilderFactoryBean = mock(StreamsBuilderFactoryBean.class);
+			return UserConfig.streamsBuilderFactoryBean;
+		}
+	}
+
+	@Configuration
+	static class EagerInitializationConfig {
+		@Bean
+		EagerInitializationComponent eagerInitializationComponent() {
+			return new EagerInitializationComponent();
+		}
+	}
+
+	static class EagerInitializationComponent {
+
+		@Autowired
+		private Tracing tracing;
+		private KafkaStreamsTracing kafkaStreamsTracing;
+
+		@PostConstruct
+		void init() {
+			kafkaStreamsTracing = KafkaStreamsTracing.create(tracing);
+		}
+	}
+}