GH-981: Add StreamsBuilderFactoryBeanCustomizer

Resolves #981

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/annotation/KafkaStreamsDefaultConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2018 the original author or authors.
+ * Copyright 2017-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.
@@ -23,6 +23,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.kafka.config.KafkaStreamsConfiguration;
 import org.springframework.kafka.config.StreamsBuilderFactoryBean;
+import org.springframework.kafka.config.StreamsBuilderFactoryBeanCustomizer;
 
 /**
  * {@code @Configuration} class that registers a {@link StreamsBuilderFactoryBean}
@@ -55,11 +56,17 @@ public class KafkaStreamsDefaultConfiguration {
 	@Bean(name = DEFAULT_STREAMS_BUILDER_BEAN_NAME)
 	public StreamsBuilderFactoryBean defaultKafkaStreamsBuilder(
 			@Qualifier(DEFAULT_STREAMS_CONFIG_BEAN_NAME)
-				ObjectProvider<KafkaStreamsConfiguration> streamsConfigProvider) {
+				ObjectProvider<KafkaStreamsConfiguration> streamsConfigProvider,
+				ObjectProvider<StreamsBuilderFactoryBeanCustomizer> customizerProvider) {
 
 		KafkaStreamsConfiguration streamsConfig = streamsConfigProvider.getIfAvailable();
 		if (streamsConfig != null) {
-			return new StreamsBuilderFactoryBean(streamsConfig);
+			StreamsBuilderFactoryBean fb = new StreamsBuilderFactoryBean(streamsConfig);
+			StreamsBuilderFactoryBeanCustomizer customizer = customizerProvider.getIfUnique();
+			if (customizer != null) {
+				customizer.configure(fb);
+			}
+			return fb;
 		}
 		else {
 			throw new UnsatisfiedDependencyException(KafkaStreamsDefaultConfiguration.class.getName(),

spring-kafka/src/main/java/org/springframework/kafka/config/StreamsBuilderFactoryBeanCustomizer.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 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
+ *
+ *      http://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.kafka.config;
+
+/**
+ * A customizer for the {@link StreamsBuilderFactoryBean} that is implicitly created by
+ * {@link org.springframework.kafka.annotation.EnableKafkaStreams}. If exactly one
+ * implementation of this interface is found in the application context (or one is marked
+ * as {@link org.springframework.context.annotation.Primary}, it will be invoked after the
+ * factory bean has been created and before it is started.
+ *
+ * @author Gary Russell
+ * @since 2.3
+ *
+ */
+@FunctionalInterface
+public interface StreamsBuilderFactoryBeanCustomizer {
+
+	/**
+	 * Configure the factory bean.
+	 * @param factoryBean the factory bean.
+	 */
+	void configure(StreamsBuilderFactoryBean factoryBean);
+
+}

spring-kafka/src/test/java/org/springframework/kafka/kstream/KafkaStreamsTests.java

@@ -25,6 +25,7 @@
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -56,6 +57,7 @@
 import org.springframework.kafka.config.KafkaListenerContainerFactory;
 import org.springframework.kafka.config.KafkaStreamsConfiguration;
 import org.springframework.kafka.config.StreamsBuilderFactoryBean;
+import org.springframework.kafka.config.StreamsBuilderFactoryBeanCustomizer;
 import org.springframework.kafka.core.ConsumerFactory;
 import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
 import org.springframework.kafka.core.DefaultKafkaProducerFactory;
@@ -113,6 +115,9 @@ public class KafkaStreamsTests {
 	@Value("${streaming.topic.two}")
 	private String streamingTopic2;
 
+	@Autowired
+	private AtomicBoolean stateChangeCalled;
+
 	@Test
 	public void testKStreams() throws Exception {
 		assertThat(this.embeddedKafka.getKafkaServer(0).config().autoCreateTopicsEnable()).isFalse();
@@ -150,6 +155,7 @@ public void testKStreams() throws Exception {
 		StreamThread[] threads = KafkaTestUtils.getPropertyValue(kafkaStreams, "threads", StreamThread[].class);
 		assertThat(threads).isNotEmpty();
 		assertThat(threads[0].getUncaughtExceptionHandler()).isSameAs(exceptionHandler);
+		assertThat(this.stateChangeCalled.get()).isTrue();
 	}
 
 	@Configuration
@@ -193,6 +199,18 @@ public KafkaStreamsConfiguration kStreamsConfigs() {
 			return new KafkaStreamsConfiguration(props);
 		}
 
+		@Bean
+		public AtomicBoolean stateChangeCalled() {
+			return new AtomicBoolean();
+		}
+
+		@Bean
+		public StreamsBuilderFactoryBeanCustomizer customizer() {
+			return fb -> fb.setStateListener((newState, oldState) -> {
+				stateChangeCalled().set(true);
+			});
+		}
+
 		@Bean
 		public KStream<Integer, String> kStream(StreamsBuilder kStreamBuilder) {
 			KStream<Integer, String> stream = kStreamBuilder.stream(STREAMING_TOPIC1);

src/reference/asciidoc/streams.adoc

@@ -179,10 +179,12 @@ See the Apache Kafka https://kafka.apache.org/0102/documentation/#streamsconfigs
 
 IMPORTANT: Starting with version 2.2, the stream configuration is now provided as a `KafkaStreamsConfiguration` object, rather than as a `StreamsConfig`.
 
-To avoid boilerplate code for most cases, especially when you develop microservices, Spring for Apache Kafka provides the `@EnableKafkaStreams` annotation, which you should placed on a `@Configuration` class.
+To avoid boilerplate code for most cases, especially when you develop microservices, Spring for Apache Kafka provides the `@EnableKafkaStreams` annotation, which you should place on a `@Configuration` class.
 All you need is to declare a `KafkaStreamsConfiguration` bean named `defaultKafkaStreamsConfig`.
-A `StreamsBuilder` bean, named `defaultKafkaStreamsBuilder`, is automatically declared in the application context.
+A `StreamsBuilderFactoryBean` bean, named `defaultKafkaStreamsBuilder`, is automatically declared in the application context.
 You can declare and use any additional `StreamsBuilderFactoryBean` beans as well.
+Starting with version 2.3, you can perform additional customization of that bean, by providing a bean that implements `StreamsBuilderFactoryBeanCustomizer`.
+There must only be one such bean, or one must be marked `@Primary`.
 
 By default, when the factory bean is stopped, the `KafkaStreams.cleanUp()` method is called.
 Starting with version 2.1.2, the factory bean has additional constructors, taking a `CleanupConfig` object that has properties to let you control whether the `cleanUp()` method is called during `start()` or `stop()` or neither.
@@ -209,6 +211,13 @@ public static class KafkaStreamsConfig {
         return new KafkaStreamsConfiguration(props);
     }
 
+    @Bean
+    public StreamsBuilderFactoryBeanCustomizer customizer() {
+        return fb -> fb.setStateListener((newState, oldState) -> {
+            System.out.println("State transition from " + oldState + " to " + newState);
+        });
+    }
+
     @Bean
     public KStream<Integer, String> kStream(StreamsBuilder kStreamBuilder) {
         KStream<Integer, String> stream = kStreamBuilder.stream("streamingTopic1");

src/reference/asciidoc/whats-new.adoc

@@ -20,3 +20,8 @@ It now sets it to false automatically unless specifically set in the consumer fa
 
 A new class `TopicBuilder` is provided for more convenient creation of `NewTopic` `@Bean` s for automatic topic provisioning.
 See <<configuring-topics>> for more information.
+
+==== Kafka Streams Changes
+
+You can now perform additional configuration of the `StreamsBuilderFactoryBean` created by `@EnableKafkaStreams`.
+See <<streams-config, Streams Configuration>> for more information.