Merge pull request #16755 from garyrussell

* pr/16755:
  Polish "Auto-configure Kafka listener container with rebalance listener"
  Auto-configure Kafka listener container with rebalance listener

Author: snicoll

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/kafka/ConcurrentKafkaListenerContainerFactoryConfigurer.java

@@ -25,6 +25,7 @@
 import org.springframework.kafka.core.KafkaTemplate;
 import org.springframework.kafka.listener.AfterRollbackProcessor;
 import org.springframework.kafka.listener.BatchErrorHandler;
+import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
 import org.springframework.kafka.listener.ContainerProperties;
 import org.springframework.kafka.listener.ErrorHandler;
 import org.springframework.kafka.support.converter.MessageConverter;
@@ -47,6 +48,8 @@ public class ConcurrentKafkaListenerContainerFactoryConfigurer {
 
 	private KafkaAwareTransactionManager<Object, Object> transactionManager;
 
+	private ConsumerAwareRebalanceListener rebalanceListener;
+
 	private ErrorHandler errorHandler;
 
 	private BatchErrorHandler batchErrorHandler;
@@ -86,6 +89,15 @@ void setTransactionManager(
 		this.transactionManager = transactionManager;
 	}
 
+	/**
+	 * Set the {@link ConsumerAwareRebalanceListener} to use.
+	 * @param rebalanceListener the rebalance listener.
+	 * @since 2.2
+	 */
+	void setRebalanceListener(ConsumerAwareRebalanceListener rebalanceListener) {
+		this.rebalanceListener = rebalanceListener;
+	}
+
 	/**
 	 * Set the {@link ErrorHandler} to use.
 	 * @param errorHandler the error handler
@@ -160,6 +172,7 @@ private void configureContainer(ContainerProperties container) {
 		map.from(properties::getLogContainerConfig).to(container::setLogContainerConfig);
 		map.from(properties::isMissingTopicsFatal).to(container::setMissingTopicsFatal);
 		map.from(this.transactionManager).to(container::setTransactionManager);
+		map.from(this.rebalanceListener).to(container::setConsumerRebalanceListener);
 	}
 
 }

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/kafka/KafkaAnnotationDrivenConfiguration.java

@@ -29,6 +29,7 @@
 import org.springframework.kafka.core.KafkaTemplate;
 import org.springframework.kafka.listener.AfterRollbackProcessor;
 import org.springframework.kafka.listener.BatchErrorHandler;
+import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
 import org.springframework.kafka.listener.ErrorHandler;
 import org.springframework.kafka.support.converter.BatchMessageConverter;
 import org.springframework.kafka.support.converter.BatchMessagingMessageConverter;
@@ -57,6 +58,8 @@ class KafkaAnnotationDrivenConfiguration {
 
 	private final KafkaAwareTransactionManager<Object, Object> transactionManager;
 
+	private final ConsumerAwareRebalanceListener rebalanceListener;
+
 	private final ErrorHandler errorHandler;
 
 	private final BatchErrorHandler batchErrorHandler;
@@ -68,6 +71,7 @@ class KafkaAnnotationDrivenConfiguration {
 			ObjectProvider<BatchMessageConverter> batchMessageConverter,
 			ObjectProvider<KafkaTemplate<Object, Object>> kafkaTemplate,
 			ObjectProvider<KafkaAwareTransactionManager<Object, Object>> kafkaTransactionManager,
+			ObjectProvider<ConsumerAwareRebalanceListener> rebalanceListener,
 			ObjectProvider<ErrorHandler> errorHandler,
 			ObjectProvider<BatchErrorHandler> batchErrorHandler,
 			ObjectProvider<AfterRollbackProcessor<Object, Object>> afterRollbackProcessor) {
@@ -77,6 +81,7 @@ class KafkaAnnotationDrivenConfiguration {
 				() -> new BatchMessagingMessageConverter(this.messageConverter));
 		this.kafkaTemplate = kafkaTemplate.getIfUnique();
 		this.transactionManager = kafkaTransactionManager.getIfUnique();
+		this.rebalanceListener = rebalanceListener.getIfUnique();
 		this.errorHandler = errorHandler.getIfUnique();
 		this.batchErrorHandler = batchErrorHandler.getIfUnique();
 		this.afterRollbackProcessor = afterRollbackProcessor.getIfUnique();
@@ -92,6 +97,7 @@ public ConcurrentKafkaListenerContainerFactoryConfigurer kafkaListenerContainerF
 		configurer.setMessageConverter(messageConverterToUse);
 		configurer.setReplyTemplate(this.kafkaTemplate);
 		configurer.setTransactionManager(this.transactionManager);
+		configurer.setRebalanceListener(this.rebalanceListener);
 		configurer.setErrorHandler(this.errorHandler);
 		configurer.setBatchErrorHandler(this.batchErrorHandler);
 		configurer.setAfterRollbackProcessor(this.afterRollbackProcessor);

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/kafka/KafkaAutoConfigurationTests.java

@@ -55,6 +55,7 @@
 import org.springframework.kafka.core.KafkaAdmin;
 import org.springframework.kafka.core.KafkaTemplate;
 import org.springframework.kafka.listener.AfterRollbackProcessor;
+import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
 import org.springframework.kafka.listener.ContainerProperties;
 import org.springframework.kafka.listener.ContainerProperties.AckMode;
 import org.springframework.kafka.listener.SeekToCurrentBatchErrorHandler;
@@ -674,6 +675,18 @@ public void testConcurrentKafkaListenerContainerFactoryWithCustomAfterRollbackPr
 				});
 	}
 
+	@Test
+	public void testConcurrentKafkaListenerContainerFactoryWithCustomRebalanceListener() {
+		this.contextRunner.withUserConfiguration(RebalanceListenerConfiguration.class)
+				.run((context) -> {
+					ConcurrentKafkaListenerContainerFactory<?, ?> factory = context
+							.getBean(ConcurrentKafkaListenerContainerFactory.class);
+					assertThat(factory.getContainerProperties())
+							.hasFieldOrPropertyWithValue("consumerRebalanceListener",
+									context.getBean("rebalanceListener"));
+				});
+	}
+
 	@Test
 	public void testConcurrentKafkaListenerContainerFactoryWithKafkaTemplate() {
 		this.contextRunner.run((context) -> {
@@ -749,6 +762,16 @@ public AfterRollbackProcessor<Object, Object> afterRollbackProcessor() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	protected static class RebalanceListenerConfiguration {
+
+		@Bean
+		public ConsumerAwareRebalanceListener rebalanceListener() {
+			return mock(ConsumerAwareRebalanceListener.class);
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	@EnableKafkaStreams
 	protected static class EnableKafkaStreamsConfiguration {

spring-boot-project/spring-boot-docs/src/main/asciidoc/spring-boot-features.adoc

@@ -6158,8 +6158,9 @@ The following component creates a listener endpoint on the `someTopic` topic:
 ----
 
 If a `KafkaTransactionManager` bean is defined, it is automatically associated to the
-container factory. Similarly, if a `ErrorHandler` or `AfterRollbackProcessor` bean is
-defined, it is automatically associated to the default factory.
+container factory. Similarly, if a `ErrorHandler`, `AfterRollbackProcessor` or
+`ConsumerAwareRebalanceListener` bean is defined, it is automatically associated to the
+default factory.
 
 Depending on the listener type, a `RecordMessageConverter` or `BatchMessageConverter` bean
 is associated to the default factory. If only a `RecordMessageConverter` bean is present