Add SleuthKafkaStreamsConfiguration to enable starting and reporting…

… spans (#1365)

docs/src/main/asciidoc/spring-cloud-sleuth.adoc

@@ -1359,6 +1359,15 @@ so that tracing headers get injected into the created Spring Kafka's
 
 To block this feature, set `spring.sleuth.messaging.kafka.enabled` to `false`.
 
+==== Spring Kafka Streams
+
+We instrument the `KafkaStreams` `KafkaClientSupplier` so that tracing headers
+get injected into the `Producer` and `Consumer`s. A `KafkaStreamsTracing` bean
+allows for further instrumentation through additional `TransformerSupplier` and
+`ProcessorSupplier` methods.
+
+To block this feature, set `spring.sleuth.messaging.kafka.streams.enabled` to `false`.
+
 ==== Spring JMS
 
 We instrument the `JmsTemplate` so that tracing headers get injected

spring-cloud-sleuth-core/pom.xml

@@ -112,6 +112,11 @@
 			<artifactId>spring-kafka</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>org.apache.kafka</groupId>
+			<artifactId>kafka-streams</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>org.springframework.security.oauth</groupId>
 			<artifactId>spring-security-oauth2</artifactId>
@@ -197,6 +202,10 @@
 			<groupId>io.zipkin.brave</groupId>
 			<artifactId>brave-instrumentation-kafka-clients</artifactId>
 		</dependency>
+		<dependency>
+			<groupId>io.zipkin.brave</groupId>
+			<artifactId>brave-instrumentation-kafka-streams</artifactId>
+		</dependency>
 		<dependency>
 			<groupId>io.zipkin.brave</groupId>
 			<artifactId>brave-instrumentation-httpclient</artifactId>

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

@@ -0,0 +1,109 @@
+/*
+ * 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 brave.Tracing;
+import brave.kafka.streams.KafkaStreamsTracing;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.apache.kafka.streams.KafkaStreams;
+
+import org.springframework.beans.BeansException;
+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;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+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;
+
+/**
+ * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration
+ * Auto-configuration} enables Kafka Streams span creation and reporting.
+ *
+ * @author Tim te Beek
+ */
+@Configuration
+@ConditionalOnBean(Tracing.class)
+@AutoConfigureAfter({ TraceAutoConfiguration.class })
+@OnMessagingEnabled
+@ConditionalOnProperty(value = "spring.sleuth.messaging.kafka.streams.enabled",
+		matchIfMissing = true)
+@ConditionalOnClass(KafkaStreams.class)
+public class SleuthKafkaStreamsConfiguration {
+
+	protected 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
+	 */
+	@Bean
+	@ConditionalOnMissingBean
+	KafkaStreamsTracing kafkaStreamsTracing(Tracing tracing) {
+		return KafkaStreamsTracing.create(tracing);
+	}
+
+	@Bean
+	KafkaStreamsBuilderFactoryBeanPostProcessor kafkaStreamsBuilderFactoryBeanPostProcessor(
+			KafkaStreamsTracing kafkaStreamsTracing) {
+		return new KafkaStreamsBuilderFactoryBeanPostProcessor(kafkaStreamsTracing);
+	}
+
+}
+
+/**
+ * Invoke
+ * {@link StreamsBuilderFactoryBean#setClientSupplier(org.apache.kafka.streams.KafkaClientSupplier)}
+ * with {@link KafkaStreamsTracing#kafkaClientSupplier()} to enable producer/consumer
+ * header injection.<br/>
+ * Explicitly not using
+ * {@link org.springframework.kafka.config.StreamsBuilderFactoryBeanCustomizer} as that
+ * only allows for a single instance, which could conflict with a user supplied instance.
+ */
+class KafkaStreamsBuilderFactoryBeanPostProcessor implements BeanPostProcessor {
+
+	private static final Log log = LogFactory
+			.getLog(KafkaStreamsBuilderFactoryBeanPostProcessor.class);
+
+	private final KafkaStreamsTracing kafkaStreamsTracing;
+
+	KafkaStreamsBuilderFactoryBeanPostProcessor(KafkaStreamsTracing kafkaStreamsTracing) {
+		this.kafkaStreamsTracing = kafkaStreamsTracing;
+	}
+
+	@Override
+	public Object postProcessAfterInitialization(Object bean, String beanName)
+			throws BeansException {
+		if (bean instanceof StreamsBuilderFactoryBean) {
+			StreamsBuilderFactoryBean sbfb = (StreamsBuilderFactoryBean) bean;
+			if (log.isDebugEnabled()) {
+				log.debug(
+						"StreamsBuilderFactoryBean bean is auto-configured to enable tracing.");
+			}
+			sbfb.setClientSupplier(kafkaStreamsTracing.kafkaClientSupplier());
+		}
+		return bean;
+	}
+
+}

spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories

@@ -21,6 +21,7 @@ org.springframework.cloud.sleuth.instrument.reactor.TraceReactorAutoConfiguratio
 org.springframework.cloud.sleuth.instrument.web.TraceWebFluxAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.zuul.TraceZuulAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.grpc.TraceGrpcAutoConfiguration,\
+org.springframework.cloud.sleuth.instrument.messaging.SleuthKafkaStreamsConfiguration,\
 org.springframework.cloud.sleuth.instrument.messaging.TraceMessagingAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.messaging.TraceSpringIntegrationAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.messaging.websocket.TraceWebSocketAutoConfiguration,\

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

@@ -0,0 +1,81 @@
+/*
+ * 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 brave.Tracing;
+import org.apache.kafka.streams.KafkaClientSupplier;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.context.SpringBootTest.WebEnvironment;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.config.StreamsBuilderFactoryBean;
+import org.springframework.test.context.junit4.SpringRunner;
+
+import static org.assertj.core.api.BDDAssertions.then;
+
+/**
+ * @author Tim te Beek
+ */
+@RunWith(SpringRunner.class)
+@SpringBootTest(classes = SleuthKafkaStreamsConfigurationTest.Config.class,
+		webEnvironment = WebEnvironment.NONE)
+public class SleuthKafkaStreamsConfigurationTest {
+
+	@Autowired
+	TestTraceStreamsBuilderFactoryBean streamsBuilderFactoryBean;
+
+	@Test
+	public void clientSupplierInvokedOnStreamsBuilderFactoryBean() {
+		then(streamsBuilderFactoryBean.clientSupplierInvoked).isTrue();
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	protected static class Config {
+
+		@Bean
+		Tracing tracing() {
+			return Tracing.newBuilder().build();
+		}
+
+		@Bean
+		StreamsBuilderFactoryBean streamsBuilderFactoryBean() {
+			TestTraceStreamsBuilderFactoryBean factoryBean = new TestTraceStreamsBuilderFactoryBean();
+			factoryBean.setAutoStartup(false);
+			return factoryBean;
+		}
+
+	}
+
+}
+
+class TestTraceStreamsBuilderFactoryBean extends StreamsBuilderFactoryBean {
+
+	boolean clientSupplierInvoked;
+
+	@Override
+	public void setClientSupplier(KafkaClientSupplier clientSupplier) {
+		this.clientSupplierInvoked = true;
+		super.setClientSupplier(clientSupplier);
+	}
+
+}