Support custom KafkaStreams implementations

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2019 the original author or authors.
+ * Copyright 2018-2024 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.
@@ -16,7 +16,11 @@
 
 package org.springframework.kafka.config;
 
+import java.util.Properties;
+
+import org.apache.kafka.streams.KafkaClientSupplier;
 import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.Topology;
 
 /**
  * Callback interface that can be used to configure {@link KafkaStreams} directly.
@@ -30,6 +34,14 @@
 @FunctionalInterface
 public interface KafkaStreamsCustomizer {
 
+	default KafkaStreams initKafkaStreams(
+			Topology topology,
+			Properties properties,
+			KafkaClientSupplier clientSupplier
+	) {
+		return new KafkaStreams(topology, properties, clientSupplier);
+	}
+
 	void customize(KafkaStreams kafkaStreams);
 
 }

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

@@ -92,7 +92,7 @@ public class StreamsBuilderFactoryBean extends AbstractFactoryBean<StreamsBuilde
 	private KafkaStreamsInfrastructureCustomizer infrastructureCustomizer = new KafkaStreamsInfrastructureCustomizer() {
 	};
 
-	private KafkaStreamsCustomizer kafkaStreamsCustomizer;
+	private KafkaStreamsCustomizer kafkaStreamsCustomizer = kafkaStreams -> { };
 
 	private KafkaStreams.StateListener stateListener;
 
@@ -361,15 +361,15 @@ public void start() {
 				try {
 					Assert.state(this.properties != null,
 							"streams configuration properties must not be null");
-					this.kafkaStreams = new KafkaStreams(this.topology, this.properties, this.clientSupplier);
+					this.kafkaStreams = this.kafkaStreamsCustomizer.initKafkaStreams(
+							this.topology, this.properties, this.clientSupplier
+					);
 					this.kafkaStreams.setStateListener(this.stateListener);
 					this.kafkaStreams.setGlobalStateRestoreListener(this.stateRestoreListener);
 					if (this.streamsUncaughtExceptionHandler != null) {
 						this.kafkaStreams.setUncaughtExceptionHandler(this.streamsUncaughtExceptionHandler);
 					}
-					if (this.kafkaStreamsCustomizer != null) {
-						this.kafkaStreamsCustomizer.customize(this.kafkaStreams);
-					}
+					this.kafkaStreamsCustomizer.customize(this.kafkaStreams);
 					if (this.cleanupConfig.cleanupOnStart()) {
 						this.kafkaStreams.cleanUp();
 					}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2023 the original author or authors.
+ * Copyright 2018-2024 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.
@@ -26,6 +26,7 @@
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.streams.KafkaClientSupplier;
 import org.apache.kafka.streams.KafkaStreams;
 import org.apache.kafka.streams.StreamsBuilder;
 import org.apache.kafka.streams.StreamsConfig;
@@ -95,6 +96,7 @@ public void testKafkaStreamsCustomizer(@Autowired KafkaStreamsConfiguration conf
 				.isEqualTo(1000);
 		assertThat(this.config.builderConfigured.get()).isTrue();
 		assertThat(this.config.topologyConfigured.get()).isTrue();
+		assertThat(this.config.ksInitialized.get()).isTrue();
 		assertThat(this.meterRegistry.get("kafka.consumer.coordinator.join.total")
 				.tag("customTag", "stream")
 				.tag("spring.id", KafkaStreamsDefaultConfiguration.DEFAULT_STREAMS_BUILDER_BEAN_NAME)
@@ -118,6 +120,8 @@ public static class KafkaStreamsConfig {
 
 		final AtomicBoolean topologyConfigured = new AtomicBoolean();
 
+		final AtomicBoolean ksInitialized = new AtomicBoolean();
+
 		@Autowired
 		EmbeddedKafkaBroker broker;
 
@@ -168,7 +172,26 @@ public KafkaStreamsConfiguration kStreamsConfigs() {
 		}
 
 		private KafkaStreamsCustomizer customizer() {
-			return kafkaStreams -> kafkaStreams.setStateListener(STATE_LISTENER);
+			return new KafkaStreamsCustomizer() {
+				@Override
+				public KafkaStreams initKafkaStreams(
+						final Topology topology,
+						final Properties properties,
+						final KafkaClientSupplier clientSupplier
+				) {
+					ksInitialized.set(true);
+					return KafkaStreamsCustomizer.super.initKafkaStreams(
+							topology,
+							properties,
+							clientSupplier
+					);
+				}
+
+				@Override
+				public void customize(final KafkaStreams kafkaStreams) {
+					kafkaStreams.setStateListener(STATE_LISTENER);
+				}
+			};
 		}
 
 		@Bean