Apply container customizer bean to AMQP MessageListenerContainer

See gh-27625

Author: garyrussell

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitAnnotationDrivenConfiguration.java

@@ -19,10 +19,13 @@
 import java.util.stream.Collectors;
 
 import org.springframework.amqp.rabbit.annotation.EnableRabbit;
+import org.springframework.amqp.rabbit.config.ContainerCustomizer;
 import org.springframework.amqp.rabbit.config.DirectRabbitListenerContainerFactory;
 import org.springframework.amqp.rabbit.config.RabbitListenerConfigUtils;
 import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
 import org.springframework.amqp.rabbit.connection.ConnectionFactory;
+import org.springframework.amqp.rabbit.listener.DirectMessageListenerContainer;
+import org.springframework.amqp.rabbit.listener.SimpleMessageListenerContainer;
 import org.springframework.amqp.rabbit.retry.MessageRecoverer;
 import org.springframework.amqp.support.converter.MessageConverter;
 import org.springframework.beans.factory.ObjectProvider;
@@ -48,14 +51,23 @@ class RabbitAnnotationDrivenConfiguration {
 
 	private final ObjectProvider<RabbitRetryTemplateCustomizer> retryTemplateCustomizers;
 
+	private final ObjectProvider<ContainerCustomizer<SimpleMessageListenerContainer>> simpleContainerCustomizer;
+
+	private final ObjectProvider<ContainerCustomizer<DirectMessageListenerContainer>> directContainerCustomizer;
+
 	private final RabbitProperties properties;
 
 	RabbitAnnotationDrivenConfiguration(ObjectProvider<MessageConverter> messageConverter,
 			ObjectProvider<MessageRecoverer> messageRecoverer,
-			ObjectProvider<RabbitRetryTemplateCustomizer> retryTemplateCustomizers, RabbitProperties properties) {
+			ObjectProvider<RabbitRetryTemplateCustomizer> retryTemplateCustomizers,
+			ObjectProvider<ContainerCustomizer<SimpleMessageListenerContainer>> simpleContainerCustomizer,
+			ObjectProvider<ContainerCustomizer<DirectMessageListenerContainer>> directContainerCustomizer,
+			RabbitProperties properties) {
 		this.messageConverter = messageConverter;
 		this.messageRecoverer = messageRecoverer;
 		this.retryTemplateCustomizers = retryTemplateCustomizers;
+		this.simpleContainerCustomizer = simpleContainerCustomizer;
+		this.directContainerCustomizer = directContainerCustomizer;
 		this.properties = properties;
 	}
 
@@ -79,6 +91,7 @@ SimpleRabbitListenerContainerFactory simpleRabbitListenerContainerFactory(
 			SimpleRabbitListenerContainerFactoryConfigurer configurer, ConnectionFactory connectionFactory) {
 		SimpleRabbitListenerContainerFactory factory = new SimpleRabbitListenerContainerFactory();
 		configurer.configure(factory, connectionFactory);
+		this.simpleContainerCustomizer.ifUnique(factory::setContainerCustomizer);
 		return factory;
 	}
 
@@ -101,6 +114,7 @@ DirectRabbitListenerContainerFactory directRabbitListenerContainerFactory(
 			DirectRabbitListenerContainerFactoryConfigurer configurer, ConnectionFactory connectionFactory) {
 		DirectRabbitListenerContainerFactory factory = new DirectRabbitListenerContainerFactory();
 		configurer.configure(factory, connectionFactory);
+		this.directContainerCustomizer.ifUnique(factory::setContainerCustomizer);
 		return factory;
 	}
 

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/amqp/RabbitAutoConfigurationTests.java

@@ -37,7 +37,9 @@
 import org.springframework.amqp.core.AmqpAdmin;
 import org.springframework.amqp.core.Message;
 import org.springframework.amqp.rabbit.annotation.EnableRabbit;
+import org.springframework.amqp.rabbit.annotation.RabbitListener;
 import org.springframework.amqp.rabbit.config.AbstractRabbitListenerContainerFactory;
+import org.springframework.amqp.rabbit.config.ContainerCustomizer;
 import org.springframework.amqp.rabbit.config.DirectRabbitListenerContainerFactory;
 import org.springframework.amqp.rabbit.config.RabbitListenerConfigUtils;
 import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
@@ -49,7 +51,9 @@
 import org.springframework.amqp.rabbit.core.RabbitAdmin;
 import org.springframework.amqp.rabbit.core.RabbitMessagingTemplate;
 import org.springframework.amqp.rabbit.core.RabbitTemplate;
+import org.springframework.amqp.rabbit.listener.DirectMessageListenerContainer;
 import org.springframework.amqp.rabbit.listener.RabbitListenerContainerFactory;
+import org.springframework.amqp.rabbit.listener.SimpleMessageListenerContainer;
 import org.springframework.amqp.rabbit.retry.MessageRecoverer;
 import org.springframework.amqp.support.converter.MessageConverter;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
@@ -849,6 +853,20 @@ void whenMultipleConnectionFactoryCustomizersAreDefinedThenTheyAreCalledInOrder(
 				});
 	}
 
+	@Test
+	void simpleContainerCustomizer() {
+		this.contextRunner.withUserConfiguration(SimpleContainerCustomizerConfiguration.class).run(
+				(context) -> assertThat(context.getBean(SimpleContainerCustomizerConfiguration.class).customizerCalled)
+						.isTrue());
+	}
+
+	@Test
+	void directContainerCustomizer() {
+		this.contextRunner.withUserConfiguration(DirectContainerCustomizerConfiguration.class)
+				.withPropertyValues("spring.rabbitmq.listener.type:direct").run((context) -> assertThat(
+						context.getBean(DirectContainerCustomizerConfiguration.class).customizerCalled).isTrue());
+	}
+
 	private com.rabbitmq.client.ConnectionFactory getTargetConnectionFactory(AssertableApplicationContext context) {
 		CachingConnectionFactory connectionFactory = context.getBean(CachingConnectionFactory.class);
 		return connectionFactory.getRabbitConnectionFactory();
@@ -1113,4 +1131,36 @@ ConnectionFactoryCustomizer firstCustomizer() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class SimpleContainerCustomizerConfiguration {
+
+		boolean customizerCalled;
+
+		@RabbitListener(queues = "test", autoStartup = "false")
+		void listen(String in) {
+		}
+
+		@Bean
+		ContainerCustomizer<SimpleMessageListenerContainer> customizer() {
+			return (container) -> this.customizerCalled = true;
+		}
+
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class DirectContainerCustomizerConfiguration {
+
+		boolean customizerCalled;
+
+		@RabbitListener(queues = "test", autoStartup = "false")
+		void listen(String in) {
+		}
+
+		@Bean
+		ContainerCustomizer<DirectMessageListenerContainer> customizer() {
+			return (container) -> this.customizerCalled = true;
+		}
+
+	}
+
 }