GH-1236: Add Streams infrastructure customizer

Resolves #1236

Allow modification of the builder and/or topolgy before the stream is
created.

* Fix typo, docs, javadocs

Author: garyrussell

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2020 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.kafka.config;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.Topology;
+
+/**
+ * Composite {@link KafkaStreamsInfrastructureCustomizer} customizes {@link KafkaStreams}
+ * by delegating to a list of provided {@link KafkaStreamsInfrastructureCustomizer}.
+ *
+ * @author Gary Russell
+ *
+ * @since 2.4.1
+ */
+public class CompositeKafkaStreamsInfrastructureCustomizer implements KafkaStreamsInfrastructureCustomizer {
+
+	private final List<KafkaStreamsInfrastructureCustomizer> infrastructureCustomizers = new ArrayList<>();
+
+	/**
+	 * Construct an instance with the provided customizers.
+	 * @param customizers the customizers;
+	 */
+	public CompositeKafkaStreamsInfrastructureCustomizer(KafkaStreamsInfrastructureCustomizer... customizers) {
+		this.infrastructureCustomizers.addAll(Arrays.asList(customizers));
+	}
+
+	/**
+	 * Add customizers.
+	 * @param customizers the customizers.
+	 */
+	public void addKafkaStreamsCustomizers(KafkaStreamsInfrastructureCustomizer... customizers) {
+		this.infrastructureCustomizers.addAll(Arrays.asList(customizers));
+	}
+
+	@Override
+	public void configureBuilder(StreamsBuilder builder) {
+		this.infrastructureCustomizers.forEach(cust -> cust.configureBuilder(builder));
+	}
+
+	@Override
+	public void configureTopology(Topology topology) {
+		this.infrastructureCustomizers.forEach(cust -> cust.configureTopology(topology));
+	}
+
+}

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2020 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.kafka.config;
+
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.Topology;
+
+/**
+ * A customizer for infrastructure components such as the {@code StreamsBuilder} and
+ * {@code Topology}. It can be provided to the {@link StreamsBuilderFactoryBean} which
+ * will apply the changes before creating the stream.
+ *
+ * @author Gary Russell
+ * @since 2.4.1
+ *
+ */
+public interface KafkaStreamsInfrastructureCustomizer {
+
+	/**
+	 * Configure the builder.
+	 * @param builder the builder.
+	 */
+	default void configureBuilder(StreamsBuilder builder) {
+		// no-op
+	}
+
+	/**
+	 * Configure the topology.
+	 * @param topology the topology
+	 */
+	default void configureTopology(Topology topology) {
+		// no-op
+	}
+
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2019 the original author or authors.
+ * Copyright 2017-2020 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.
@@ -72,6 +72,9 @@ public class StreamsBuilderFactoryBean extends AbstractFactoryBean<StreamsBuilde
 
 	private final CleanupConfig cleanupConfig;
 
+	private KafkaStreamsInfrastructureCustomizer infrastructureCustomizer = new KafkaStreamsInfrastructureCustomizer() {
+	};
+
 	private KafkaStreamsCustomizer kafkaStreamsCustomizer;
 
 	private KafkaStreams.StateListener stateListener;
@@ -144,6 +147,16 @@ public void setClientSupplier(KafkaClientSupplier clientSupplier) {
 		this.clientSupplier = clientSupplier; // NOSONAR (sync)
 	}
 
+	/**
+	 * Set a customizer to configure the builder and/or topology before creating the stream.
+	 * @param infrastructureCustomizer the customizer
+	 * @since 2.4.1
+	 */
+	public void setInfrastructureCustomizer(KafkaStreamsInfrastructureCustomizer infrastructureCustomizer) {
+		Assert.notNull(infrastructureCustomizer, "'infrastructureCustomizer' must not be null");
+		this.infrastructureCustomizer = infrastructureCustomizer;
+	}
+
 	/**
 	 * Specify a {@link KafkaStreamsCustomizer} to customize a {@link KafkaStreams}
 	 * instance during {@link #start()}.
@@ -233,7 +246,10 @@ public synchronized void start() {
 			try {
 				Assert.state(this.properties != null,
 						"streams configuration properties must not be null");
-				Topology topology = getObject().build(this.properties); // NOSONAR: getObject() cannot return null
+				StreamsBuilder builder = getObject();
+				this.infrastructureCustomizer.configureBuilder(builder);
+				Topology topology = builder.build(this.properties); // NOSONAR: getObject() cannot return null
+				this.infrastructureCustomizer.configureTopology(topology);
 				LOGGER.debug(() -> topology.describe().toString());
 				this.kafkaStreams = new KafkaStreams(topology, this.properties, this.clientSupplier);
 				this.kafkaStreams.setStateListener(this.stateListener);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 the original author or authors.
+ * Copyright 2020 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.

spring-kafka/src/test/java/org/springframework/kafka/config/KafkaStreamsCustomizerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2019 the original author or authors.
+ * Copyright 2018-2020 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.
@@ -22,10 +22,13 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.StreamsBuilder;
 import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.Topology;
 import org.apache.kafka.streams.errors.DeserializationExceptionHandler;
 import org.apache.kafka.streams.processor.ProcessorContext;
 import org.junit.jupiter.api.Test;
@@ -61,6 +64,9 @@ public class KafkaStreamsCustomizerTests {
 	@Autowired
 	private StreamsBuilderFactoryBean streamsBuilderFactoryBean;
 
+	@Autowired
+	private KafkaStreamsConfig config;
+
 	@Test
 	public void testKafkaStreamsCustomizer(@Autowired KafkaStreamsConfiguration configuration,
 			@Autowired KafkaStreamsConfig config) {
@@ -74,13 +80,19 @@ public void testKafkaStreamsCustomizer(@Autowired KafkaStreamsConfiguration conf
 				.isEqualTo(Foo.class);
 		assertThat(properties.get(StreamsConfig.CACHE_MAX_BYTES_BUFFERING_CONFIG))
 			.isEqualTo(1000);
+		assertThat(this.config.builderConfigured.get()).isTrue();
+		assertThat(this.config.topologyConfigured.get()).isTrue();
 	}
 
 	@Configuration
 	@EnableKafka
 	@EnableKafkaStreams
 	public static class KafkaStreamsConfig {
 
+		final AtomicBoolean builderConfigured = new AtomicBoolean();
+
+		final AtomicBoolean topologyConfigured = new AtomicBoolean();
+
 		@Value("${" + EmbeddedKafkaBroker.SPRING_EMBEDDED_KAFKA_BROKERS + "}")
 		private String brokerAddresses;
 
@@ -89,6 +101,20 @@ public StreamsBuilderFactoryBean defaultKafkaStreamsBuilder() {
 			StreamsBuilderFactoryBean streamsBuilderFactoryBean =
 					new StreamsBuilderFactoryBean(kStreamsConfigs());
 			streamsBuilderFactoryBean.setKafkaStreamsCustomizer(customizer());
+			streamsBuilderFactoryBean.setInfrastructureCustomizer(new KafkaStreamsInfrastructureCustomizer() {
+
+				@SuppressWarnings("unchecked")
+				@Override
+				public void configureBuilder(StreamsBuilder builder) {
+					KafkaStreamsConfig.this.builderConfigured.set(true);
+				}
+
+				@Override
+				public void configureTopology(Topology topology) {
+					KafkaStreamsConfig.this.topologyConfigured.set(true);
+				}
+
+			});
 			return streamsBuilderFactoryBean;
 		}
 

src/reference/asciidoc/streams.adoc

@@ -43,6 +43,7 @@ In other words, all streams defined by a `StreamsBuilder` are tied with a single
 Once a `KafkaStreams` instance has been closed by `streams.close()`, it cannot be restarted.
 Instead, a new `KafkaStreams` instance to restart stream processing must be created.
 
+[[streams-spring]]
 ==== Spring Management
 
 To simplify using Kafka Streams from the Spring application context perspective and use the lifecycle management through a container, the Spring for Apache Kafka introduces `StreamsBuilderFactoryBean`.
@@ -121,6 +122,25 @@ private StreamsBuilderFactoryBean myKStreamBuilderFactoryBean;
 ----
 ====
 
+Starting with version 2.4.1, the factory bean has a new property `infrastructureCustomizer` with type `KafkaStreamsInfrastructureCustomizer`; this allows customization of the `StreamsBuilder` (e.g. to add a state store) and/or the `Topology` before the stream is created.
+
+====
+[source, java]
+----
+public interface KafkaStreamsInfrastructureCustomizer {
+
+	void configureBuilder(StreamsBuilder builder);
+
+	void configureTopology(Topology topology);
+
+}
+----
+====
+
+Default no-op implementations are provided to avoid having to implement both methods if one is not required.
+
+A `CompositeKafkaStreamsInfrastructureCustomizer` is provided, for when you need to apply multiple customizers.
+
 [[serde]]
 ==== Streams JSON Serialization and Deserialization
 

src/reference/asciidoc/whats-new.adoc

@@ -44,6 +44,12 @@ The `@KafkaListener` annotation has a new property `splitIterables`; default tru
 When a replying listener returns an `Iterable` this property controls whether the return result is sent as a single record or a record for each element is sent.
 See <<annotation-send-to>> for more information.
 
+==== Kafka Streams
+
+The `StreamsBuilderFactoryBean` accepts a new property `KafkaStreamsInfrastructureCustomizer`.
+This allows configuration of the builder and/or topology before the stream is created.
+See <<streams-spring>> for more information.
+
 === Migration Guide
 
 * This release is essentially the same as the 2.3.x line, except it has been compiled against the 2.4 `kafka-clients` jar, due to a binary incompatibility.