From 1d38c149354cfb1f863e022bdd879fdb080c0e1e Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@pivotal.io>
Date: Tue, 13 Mar 2018 16:06:20 -0400
Subject: [PATCH] GH-719: Container Factory: Add replyPostProcessor

Resolves https://github.com/spring-projects/spring-amqp/issues/719

Add `replyPostProcessor` to the container factory.

* Polishing - support an array of MPPs for consistency with other places.

* Polishing - PR Comments
---
 ...bstractRabbitListenerContainerFactory.java | 20 +++++++++-
 .../AbstractAdaptableMessageListener.java     | 37 +++++++++++++------
 .../EnableRabbitIntegrationTests.java         |  5 +++
 .../core/RabbitTemplateIntegrationTests.java  |  2 +-
 src/reference/asciidoc/amqp.adoc              |  2 +
 src/reference/asciidoc/whats-new.adoc         |  2 +
 6 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/config/AbstractRabbitListenerContainerFactory.java b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/config/AbstractRabbitListenerContainerFactory.java
index 81f3d71f..5c5f576f 100644
--- a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/config/AbstractRabbitListenerContainerFactory.java
+++ b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/config/AbstractRabbitListenerContainerFactory.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2017 the original author or authors.
+ * Copyright 2014-2018 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.
@@ -31,6 +31,7 @@ import org.springframework.amqp.rabbit.listener.AbstractMessageListenerContainer
 import org.springframework.amqp.rabbit.listener.RabbitListenerContainerFactory;
 import org.springframework.amqp.rabbit.listener.RabbitListenerEndpoint;
 import org.springframework.amqp.rabbit.listener.SimpleMessageListenerContainer;
+import org.springframework.amqp.rabbit.listener.adapter.AbstractAdaptableMessageListener;
 import org.springframework.amqp.support.ConsumerTagStrategy;
 import org.springframework.amqp.support.converter.MessageConverter;
 import org.springframework.beans.BeansException;
@@ -103,6 +104,8 @@ public abstract class AbstractRabbitListenerContainerFactory<C extends AbstractM
 
 	private MessagePostProcessor[] afterReceivePostProcessors;
 
+	private MessagePostProcessor[] beforeSendReplyPostProcessors;
+
 	protected final AtomicInteger counter = new AtomicInteger();
 
 	/**
@@ -274,13 +277,23 @@ public abstract class AbstractRabbitListenerContainerFactory<C extends AbstractM
 	}
 
 	/**
+	 * Set post processors which will be applied after the Message is received.
 	 * @param afterReceivePostProcessors the post processors.
+	 * @since 2.0
 	 * @see AbstractMessageListenerContainer#setAfterReceivePostProcessors(MessagePostProcessor...)
 	 */
 	public void setAfterReceivePostProcessors(MessagePostProcessor... afterReceivePostProcessors) {
 		this.afterReceivePostProcessors = afterReceivePostProcessors;
 	}
 
+	/**
+	 * Set post processors that will be applied before sending replies.
+	 * @param beforeSendReplyPostProcessors the post processors.
+	 * @since 2.0.3
+	 */
+	public void setBeforeSendReplyPostProcessors(MessagePostProcessor... beforeSendReplyPostProcessors) {
+		this.beforeSendReplyPostProcessors = beforeSendReplyPostProcessors;
+	}
 
 	@Override
 	public C createListenerContainer(RabbitListenerEndpoint endpoint) {
@@ -355,6 +368,11 @@ public abstract class AbstractRabbitListenerContainerFactory<C extends AbstractM
 		instance.setListenerId(endpoint.getId());
 
 		endpoint.setupListenerContainer(instance);
+		if (this.beforeSendReplyPostProcessors != null
+				&& instance.getMessageListener() instanceof AbstractAdaptableMessageListener) {
+			((AbstractAdaptableMessageListener) instance.getMessageListener())
+					.setBeforeSendReplyPostProcessors(this.beforeSendReplyPostProcessors);
+		}
 		initializeContainer(instance, endpoint);
 
 		return instance;
diff --git a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/adapter/AbstractAdaptableMessageListener.java b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/adapter/AbstractAdaptableMessageListener.java
index f35ab7ad..aab6fea7 100644
--- a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/adapter/AbstractAdaptableMessageListener.java
+++ b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/adapter/AbstractAdaptableMessageListener.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2017 the original author or authors.
+ * Copyright 2014-2018 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,6 +16,8 @@
 
 package org.springframework.amqp.rabbit.listener.adapter;
 
+import java.util.Arrays;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
@@ -88,7 +90,7 @@ public abstract class AbstractAdaptableMessageListener implements MessageListene
 
 	private String encoding = DEFAULT_ENCODING;
 
-	private MessagePostProcessor replyPostProcessor;
+	private MessagePostProcessor[] beforeSendReplyPostProcessors;
 
 	/**
 	 * Set the routing key to use when sending response messages.
@@ -173,12 +175,26 @@ public abstract class AbstractAdaptableMessageListener implements MessageListene
 	}
 
 	/**
-	 * Set a post processor to process the reply immediately before {@code Channel#basicPublish()}.
-	 * Often used to compress the data.
+	 * Set a post processor to process the reply immediately before
+	 * {@code Channel#basicPublish()}. Often used to compress the data.
 	 * @param replyPostProcessor the reply post processor.
+	 * @deprecated in favor of
+	 * {@link #setBeforeSendReplyPostProcessors(MessagePostProcessor...)}.
 	 */
+	@Deprecated
 	public void setReplyPostProcessor(MessagePostProcessor replyPostProcessor) {
-		this.replyPostProcessor = replyPostProcessor;
+		setBeforeSendReplyPostProcessors(replyPostProcessor);
+	}
+
+	/**
+	 * Set post processors that will be applied before sending replies.
+	 * @param beforeSendReplyPostProcessors the post processors.
+	 * @since 2.0.3
+	 */
+	public void setBeforeSendReplyPostProcessors(MessagePostProcessor... beforeSendReplyPostProcessors) {
+		Assert.noNullElements(beforeSendReplyPostProcessors, "'replyPostProcessors' must not have any null elements");
+		this.beforeSendReplyPostProcessors = Arrays.copyOf(beforeSendReplyPostProcessors,
+				beforeSendReplyPostProcessors.length);
 	}
 
 	/**
@@ -414,12 +430,11 @@ public abstract class AbstractAdaptableMessageListener implements MessageListene
 	 * @see #postProcessResponse(Message, Message)
 	 */
 	protected void sendResponse(Channel channel, Address replyTo, Message messageIn) throws Exception {
-		Message message;
-		if (this.replyPostProcessor == null) {
-			message = messageIn;
-		}
-		else {
-			message = this.replyPostProcessor.postProcessMessage(messageIn);
+		Message message = messageIn;
+		if (this.beforeSendReplyPostProcessors != null) {
+			for (MessagePostProcessor postProcessor : this.beforeSendReplyPostProcessors) {
+				message = postProcessor.postProcessMessage(message);
+			}
 		}
 		postProcessChannel(channel, message);
 
diff --git a/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/annotation/EnableRabbitIntegrationTests.java b/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/annotation/EnableRabbitIntegrationTests.java
index 33afc0a8..b1b1cabb 100644
--- a/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/annotation/EnableRabbitIntegrationTests.java
+++ b/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/annotation/EnableRabbitIntegrationTests.java
@@ -370,6 +370,7 @@ public class EnableRabbitIntegrationTests {
 		Message request = MessageTestUtils.createTextMessage("foo", properties);
 		Message reply = rabbitTemplate.sendAndReceive("test.header", request);
 		assertEquals("prefix-FOO", MessageTestUtils.extractText(reply));
+		assertEquals(reply.getMessageProperties().getHeaders().get("replyMPPApplied"), Boolean.TRUE);
 	}
 
 	@Test
@@ -1105,6 +1106,10 @@ public class EnableRabbitIntegrationTests {
 			factory.setErrorHandler(errorHandler());
 			factory.setConsumerTagStrategy(consumerTagStrategy());
 			factory.setReceiveTimeout(10L);
+			factory.setBeforeSendReplyPostProcessors(m -> {
+				m.getMessageProperties().getHeaders().put("replyMPPApplied", true);
+				return m;
+			});
 			return factory;
 		}
 
diff --git a/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/core/RabbitTemplateIntegrationTests.java b/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/core/RabbitTemplateIntegrationTests.java
index 260fe567..0709908b 100644
--- a/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/core/RabbitTemplateIntegrationTests.java
+++ b/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/core/RabbitTemplateIntegrationTests.java
@@ -1382,7 +1382,7 @@ public class RabbitTemplateIntegrationTests {
 				return message.toUpperCase();
 			}
 		});
-		messageListener.setReplyPostProcessor(new GZipPostProcessor());
+		messageListener.setBeforeSendReplyPostProcessors(new GZipPostProcessor());
 		container.setMessageListener(messageListener);
 		container.setReceiveTimeout(100);
 		container.afterPropertiesSet();
diff --git a/src/reference/asciidoc/amqp.adoc b/src/reference/asciidoc/amqp.adoc
index 5b52a1ea..cf736bc3 100644
--- a/src/reference/asciidoc/amqp.adoc
+++ b/src/reference/asciidoc/amqp.adoc
@@ -1850,6 +1850,8 @@ It is possible to customize the listener container factory to use per annotation
 The default is only required if at least one endpoint is registered without a specific container factory.
 See the javadoc for full details and examples.
 
+The container factories provide methods for adding `MessagePostProcessor` s that will be applied after receiving messages (before invoking the listener) and before sending replies.
+
 If you prefer XML configuration, use the `<rabbit:annotation-driven>` element; any beans annotated with `@RabbitListener` will be detected.
 
 For `SimpleRabbitListenerContainer` s:
diff --git a/src/reference/asciidoc/whats-new.adoc b/src/reference/asciidoc/whats-new.adoc
index 74cff063..c34b564b 100644
--- a/src/reference/asciidoc/whats-new.adoc
+++ b/src/reference/asciidoc/whats-new.adoc
@@ -141,6 +141,8 @@ You can now set the `concurrency` of the listener container at the annotation le
 
 You can now set the `autoStartup` property of the listener container at the annotation level, overriding the default setting in the container factory.
 
+You can now set after receive and before send (reply) `MessagePostProcessor` s in the `RabbitListener` container factories.
+
 See <<async-annotation-driven>> for more information.
 
 Starting with _version 2.0.3_, one of the `@RabbitHandler` s on a class-level `@RabbitListener` can be designated as the default.