GH-725: Fix return Collection<Message<?>>

Fixes #725
Fixes #726

Handle return type `Collection<Message<?>>`.
Allow empty `@SendTo` when the return type is `Message<?>` (or a collection thereof).

* Polishing according PR comments

Author: garyrussell

Diff for: spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/DelegatingInvocableHandler.java

@@ -35,6 +35,7 @@
 import org.springframework.expression.common.TemplateParserContext;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.kafka.KafkaException;
+import org.springframework.kafka.support.KafkaUtils;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.handler.annotation.Header;
@@ -66,6 +67,8 @@ public class DelegatingInvocableHandler {
 
 	private final Map<InvocableHandlerMethod, Expression> handlerSendTo = new HashMap<>();
 
+	private final Map<InvocableHandlerMethod, Boolean> handlerReturnsMessage = new HashMap<>();
+
 	private final Object bean;
 
 	private final BeanExpressionResolver resolver;
@@ -124,10 +127,7 @@ public Object invoke(Message<?> message, Object... providedArgs) throws Exceptio
 		InvocableHandlerMethod handler = getHandlerForPayload(payloadClass);
 		Object result = handler.invoke(message, providedArgs);
 		Expression replyTo = this.handlerSendTo.get(handler);
-		if (replyTo != null) {
-			result = new InvocationResult(result, replyTo);
-		}
-		return result;
+		return new InvocationResult(result, replyTo, this.handlerReturnsMessage.get(handler));
 	}
 
 	/**
@@ -162,6 +162,7 @@ private void setupReplyTo(InvocableHandlerMethod handler) {
 		if (replyTo != null) {
 			this.handlerSendTo.put(handler, PARSER.parseExpression(replyTo, PARSER_CONTEXT));
 		}
+		this.handlerReturnsMessage.put(handler, KafkaUtils.returnTypeMessageOrCollectionOf(method));
 	}
 
 	private String extractSendTo(String element, SendTo ann) {

Diff for: spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/InvocationResult.java

@@ -30,9 +30,12 @@ public final class InvocationResult {
 
 	private final Expression sendTo;
 
-	public InvocationResult(Object result, Expression sendTo) {
+	private final boolean messageReturnType;
+
+	public InvocationResult(Object result, Expression sendTo, boolean messageReturnType) {
 		this.result = result;
 		this.sendTo = sendTo;
+		this.messageReturnType = messageReturnType;
 	}
 
 	public Object getResult() {
@@ -43,6 +46,10 @@ public Expression getSendTo() {
 		return this.sendTo;
 	}
 
+	public boolean isMessageReturnType() {
+		return this.messageReturnType;
+	}
+
 	@Override
 	public String toString() {
 		return this.result.toString();

Diff for: spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/MessagingMessageListenerAdapter.java

@@ -49,6 +49,7 @@
 import org.springframework.kafka.listener.MessageListener;
 import org.springframework.kafka.support.Acknowledgment;
 import org.springframework.kafka.support.KafkaHeaders;
+import org.springframework.kafka.support.KafkaUtils;
 import org.springframework.kafka.support.converter.MessagingMessageConverter;
 import org.springframework.kafka.support.converter.RecordMessageConverter;
 import org.springframework.lang.Nullable;
@@ -104,6 +105,8 @@ public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerS
 
 	private boolean hasAckParameter;
 
+	private boolean messageReturnType;
+
 	public MessagingMessageListenerAdapter(Object bean, Method method) {
 		this.bean = bean;
 		this.inferredType = determineInferredType(method);
@@ -288,11 +291,13 @@ protected void handleResult(Object resultArg, Object request, Object source) {
 			this.logger.debug("Listener method returned result [" + resultArg
 					+ "] - generating response message for it");
 		}
-		Object result = resultArg instanceof InvocationResult ? ((InvocationResult) resultArg).getResult() : resultArg;
+		boolean isInvocationResult = resultArg instanceof InvocationResult;
+		Object result = isInvocationResult ? ((InvocationResult) resultArg).getResult() : resultArg;
 		String replyTopic = evaluateReplyTopic(request, source, resultArg);
 		Assert.state(replyTopic == null || this.replyTemplate != null,
 				"a KafkaTemplate is required to support replies");
-		sendResponse(result, replyTopic, source);
+		sendResponse(result, replyTopic, source, isInvocationResult
+				? ((InvocationResult) resultArg).isMessageReturnType() : this.messageReturnType);
 	}
 
 	private String evaluateReplyTopic(Object request, Object source, Object result) {
@@ -311,12 +316,13 @@ private String evaluateTopic(Object request, Object source, Object result, Expre
 			return sendTo.getValue(String.class);
 		}
 		else {
-			Object value = sendTo.getValue(this.evaluationContext, new ReplyExpressionRoot(request, source, result));
+			Object value = sendTo == null ? null
+					: sendTo.getValue(this.evaluationContext, new ReplyExpressionRoot(request, source, result));
 			boolean isByteArray = value instanceof byte[];
-			if (!(value instanceof String || isByteArray)) {
+			if (!(value == null || value instanceof String || isByteArray)) {
 				throw new IllegalStateException(
 					"replyTopic expression must evaluate to a String or byte[], it is: "
-					+ (value == null ? null : value.getClass().getName()));
+					+ value.getClass().getName());
 			}
 			if (isByteArray) {
 				return new String((byte[]) value, StandardCharsets.UTF_8);
@@ -334,7 +340,7 @@ private String evaluateTopic(Object request, Object source, Object result, Expre
 	 */
 	@Deprecated
 	protected void sendResponse(Object result, String topic) {
-		sendResponse(result, topic, null);
+		sendResponse(result, topic, null, false);
 	}
 
 	/**
@@ -343,11 +349,12 @@ protected void sendResponse(Object result, String topic) {
 	 * @param result the result.
 	 * @param topic the topic.
 	 * @param source the source (input).
+	 * @param messageReturnType true if we are returning message(s).
 	 * @since 2.1.3
 	 */
 	@SuppressWarnings("unchecked")
-	protected void sendResponse(Object result, String topic, @Nullable Object source) {
-		if (topic == null) {
+	protected void sendResponse(Object result, String topic, @Nullable Object source, boolean messageReturnType) {
+		if (!messageReturnType && topic == null) {
 			if (this.logger.isDebugEnabled()) {
 				this.logger.debug("No replyTopic to handle the reply: " + result);
 			}
@@ -358,7 +365,12 @@ else if (result instanceof Message) {
 		else {
 			if (result instanceof Collection) {
 				((Collection<V>) result).forEach(v -> {
-					this.replyTemplate.send(topic, v);
+					if (v instanceof Message) {
+						this.replyTemplate.send((Message<?>) v);
+					}
+					else {
+						this.replyTemplate.send(topic, v);
+					}
 				});
 			}
 			else {
@@ -486,6 +498,7 @@ else if (methodParameter.getGenericParameterType().equals(Acknowledgment.class))
 				() -> String.format(stateMessage, "List<ConsumerRecord>"));
 		Assert.state(!this.isMessageList || validParametersForBatch,
 				() -> String.format(stateMessage, "List<Message<?>>"));
+		this.messageReturnType = KafkaUtils.returnTypeMessageOrCollectionOf(method);
 		return genericParameterType;
 	}
 

Diff for: spring-kafka/src/main/java/org/springframework/kafka/support/KafkaUtils.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://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.support;
+
+import java.lang.reflect.Method;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
+import java.util.Collection;
+
+import org.springframework.messaging.Message;
+
+/**
+ * Utility methods.
+ *
+ * @author Gary Russell
+ *
+ * @since 2.2
+ *
+ */
+public final class KafkaUtils {
+
+	/**
+	 * Return true if the method return type is {@link Message} or
+	 * {@code Collection<Message<?>>}.
+	 * @param method the method.
+	 * @return true if it returns message(s).
+	 */
+	public static boolean returnTypeMessageOrCollectionOf(Method method) {
+		Type returnType = method.getGenericReturnType();
+		if (returnType.equals(Message.class)) {
+			return true;
+		}
+		if (returnType instanceof ParameterizedType) {
+			ParameterizedType prt = (ParameterizedType) returnType;
+			Type rawType = prt.getRawType();
+			if (rawType.equals(Message.class)) {
+				return true;
+			}
+			if (rawType.equals(Collection.class)) {
+				Type collectionType = prt.getActualTypeArguments()[0];
+				if (collectionType.equals(Message.class)) {
+					return true;
+				}
+				return collectionType instanceof ParameterizedType
+						&& ((ParameterizedType) collectionType).getRawType().equals(Message.class);
+			}
+		}
+		return false;
+
+	}
+
+	private KafkaUtils() {
+		super();
+	}
+
+}

Diff for: spring-kafka/src/test/java/org/springframework/kafka/annotation/BatchListenerConversionTests.java

@@ -18,11 +18,13 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.producer.ProducerConfig;
@@ -48,7 +50,9 @@
 import org.springframework.kafka.test.utils.KafkaTestUtils;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.handler.annotation.Header;
+import org.springframework.messaging.handler.annotation.SendTo;
 import org.springframework.messaging.support.GenericMessage;
+import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -66,7 +70,8 @@ public class BatchListenerConversionTests {
 	private static final String DEFAULT_TEST_GROUP_ID = "blc";
 
 	@ClassRule // one topic to preserve order
-	public static KafkaEmbedded embeddedKafka = new KafkaEmbedded(1, true, 1, "blc1", "blc2", "blc3");
+	public static KafkaEmbedded embeddedKafka = new KafkaEmbedded(1, true, 1, "blc1", "blc2", "blc3",
+			"blc4", "blc5");
 
 	@Autowired
 	private Config config;
@@ -108,22 +113,40 @@ public void testBatchOfPojoMessages() throws Exception {
 		assertThat(listener.received.get(0).getPayload().getBar()).isEqualTo("bar");
 	}
 
+	@Test
+	public void testBatchReplies() throws Exception {
+		Listener4 listener = this.config.listener4();
+		String topic = "blc4";
+		this.template.send(new GenericMessage<>(
+				new Foo("bar"), Collections.singletonMap(KafkaHeaders.TOPIC, topic)));
+		this.template.send(new GenericMessage<>(
+				new Foo("baz"), Collections.singletonMap(KafkaHeaders.TOPIC, topic)));
+		assertThat(listener.latch1.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(listener.received.size()).isGreaterThan(0);
+		assertThat(listener.received.get(0)).isInstanceOf(Foo.class);
+		assertThat(listener.received.get(0).bar).isEqualTo("bar");
+		assertThat(listener.replies.size()).isGreaterThan(0);
+		assertThat(listener.replies.get(0)).isInstanceOf(Foo.class);
+		assertThat(listener.replies.get(0).bar).isEqualTo("BAR");
+	}
+
 	@Configuration
 	@EnableKafka
 	public static class Config {
 
 		@Bean
 		public KafkaListenerContainerFactory<?> kafkaListenerContainerFactory() {
-			ConcurrentKafkaListenerContainerFactory<Integer, String> factory =
+			ConcurrentKafkaListenerContainerFactory<Integer, Foo> factory =
 					new ConcurrentKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(consumerFactory());
 			factory.setBatchListener(true);
 			factory.setMessageConverter(new BatchMessagingMessageConverter(converter()));
+			factory.setReplyTemplate(template());
 			return factory;
 		}
 
 		@Bean
-		public DefaultKafkaConsumerFactory<Integer, String> consumerFactory() {
+		public DefaultKafkaConsumerFactory<Integer, Foo> consumerFactory() {
 			return new DefaultKafkaConsumerFactory<>(consumerConfigs());
 		}
 
@@ -175,6 +198,11 @@ public Listener3 listener3() {
 			return new Listener3();
 		}
 
+		@Bean
+		public Listener4 listener4() {
+			return new Listener4();
+		}
+
 	}
 
 	public static class Listener {
@@ -234,6 +262,35 @@ public void listen1(List<Message<Foo>> foos) {
 
 	}
 
+	public static class Listener4 {
+
+		private final CountDownLatch latch1 = new CountDownLatch(1);
+
+		private List<Foo> received;
+
+		private List<Foo> replies;
+
+		@KafkaListener(topics = "blc4", groupId = "blc4")
+		@SendTo
+		public Collection<Message<?>> listen1(List<Foo> foos) {
+			if (this.received == null) {
+				this.received = foos;
+			}
+			return foos.stream().map(f -> MessageBuilder.withPayload(new Foo(f.getBar().toUpperCase()))
+					.setHeader(KafkaHeaders.TOPIC, "blc5")
+					.setHeader(KafkaHeaders.MESSAGE_KEY, 42)
+					.build())
+				.collect(Collectors.toList());
+		}
+
+		@KafkaListener(topics = "blc5", groupId = "blc5")
+		public void listen2(List<Foo> foos) {
+			this.replies = foos;
+			this.latch1.countDown();
+		}
+
+	}
+
 	public static class Foo {
 
 		public String bar;