From f6b83167ac577b1bdca7cd27ded3fff3ae7420e8 Mon Sep 17 00:00:00 2001
From: Michael Minella <mminella@gopivotal.com>
Date: Thu, 1 May 2014 09:48:02 -0500
Subject: [PATCH] XD-1389: Added null check when receiving partitioned messages
 back to handle timeouts

---
 .../MessageChannelPartitionHandler.java       | 12 ++++--
 .../MessageChannelPartitionHandlerTest.java   | 39 ++++++++++++++-----
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java b/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java
index e84e2507b..e57cbb076 100644
--- a/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java
+++ b/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java
@@ -11,6 +11,7 @@ import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.partition.PartitionHandler;
 import org.springframework.batch.core.partition.StepExecutionSplitter;
 import org.springframework.batch.core.repository.JobRepository;
+import org.springframework.integration.MessageTimeoutException;
 import org.springframework.integration.annotation.Aggregator;
 import org.springframework.integration.annotation.MessageEndpoint;
 import org.springframework.integration.annotation.Payloads;
@@ -33,6 +34,7 @@ import org.springframework.util.Assert;
  *
  * @author Dave Syer
  * @author Will Schipp
+ * @author Michael Minella
  *
  */
 @MessageEndpoint
@@ -106,10 +108,10 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
 	}
 
 	/**
-	 * Sends {@link StepExecutionRequest} objects to the request channel of the {@link MessagingOperations}, and then
+	 * Sends {@link StepExecutionRequest} objects to the request channel of the {@link MessagingTemplate}, and then
 	 * receives the result back as a list of {@link StepExecution} on a reply channel. Use the {@link #aggregate(List)}
 	 * method as an aggregator of the individual remote replies. The receive timeout needs to be set realistically in
-	 * the {@link MessagingOperations} <b>and</b> the aggregator, so that there is a good chance of all work being done.
+	 * the {@link MessagingTemplate} <b>and</b> the aggregator, so that there is a good chance of all work being done.
 	 *
 	 * @see PartitionHandler#handle(StepExecutionSplitter, StepExecution)
 	 */
@@ -134,9 +136,13 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
 
 		@SuppressWarnings("unchecked")
 		Message<Collection<StepExecution>> message = (Message<Collection<StepExecution>>) messagingGateway.receive(replyChannel);
-		if (logger.isDebugEnabled()) {
+
+		if(message == null) {
+			throw new MessageTimeoutException("Timeout occurred before all partitions returned");
+		} else if (logger.isDebugEnabled()) {
 			logger.debug("Received replies: " + message);
 		}
+
 		Collection<StepExecution> result = message.getPayload();
 		return result;
 
diff --git a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java
index 3b5d812c4..d382a1df2 100644
--- a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java
+++ b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java
@@ -1,24 +1,26 @@
 package org.springframework.batch.integration.partition;
 
+import org.junit.Test;
+import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.partition.StepExecutionSplitter;
+import org.springframework.integration.MessageTimeoutException;
+import org.springframework.integration.core.MessagingTemplate;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.PollableChannel;
+
+import java.util.Collection;
+import java.util.Collections;
+
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.anyObject;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import java.util.Collection;
-import java.util.Collections;
-
-import org.junit.Test;
-import org.springframework.batch.core.StepExecution;
-import org.springframework.batch.core.partition.StepExecutionSplitter;
-import org.springframework.integration.core.MessagingTemplate;
-import org.springframework.messaging.Message;
-import org.springframework.messaging.PollableChannel;
-
 /**
  *
  * @author Will Schipp
+ * @author Michael Minella
  *
  */
 public class MessageChannelPartitionHandlerTest {
@@ -72,4 +74,21 @@ public class MessageChannelPartitionHandlerTest {
 
 	}
 
+	@Test(expected = MessageTimeoutException.class)
+	public void messageReceiveTimeout() throws Exception {
+		//execute with no default set
+		messageChannelPartitionHandler = new MessageChannelPartitionHandler();
+		//mock
+		StepExecution masterStepExecution = mock(StepExecution.class);
+		StepExecutionSplitter stepExecutionSplitter = mock(StepExecutionSplitter.class);
+		MessagingTemplate operations = mock(MessagingTemplate.class);
+		Message message = mock(Message.class);
+		//when
+		when(message.getPayload()).thenReturn(Collections.emptyList());
+		//set
+		messageChannelPartitionHandler.setMessagingOperations(operations);
+
+		//execute
+		Collection<StepExecution> executions = messageChannelPartitionHandler.handle(stepExecutionSplitter, masterStepExecution);
+	}
 }