Add job repository polling option to MessageChannelPartitionHandler

When using remote partitioning, each slave worker persists it's current
status in the same job repsository that the master uses.  Because of
this, there is no hard need for the master to wait for each worker to
send a formal response once it's work is complete.  Instead, the master
(at the cost of polling a db periodically) can determine if the workers
are done by looking up each partition's status in the job repository.
This commit removes the requirement for a reply channel and implements
the polling of the job repository to determine if the workers are done.

BATCH-2332

spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java

@@ -1,12 +1,29 @@
 package org.springframework.batch.integration.partition;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import javax.sql.DataSource;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
 import org.springframework.batch.core.Step;
 import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.explore.JobExplorer;
+import org.springframework.batch.core.explore.support.JobExplorerFactoryBean;
 import org.springframework.batch.core.partition.PartitionHandler;
 import org.springframework.batch.core.partition.StepExecutionSplitter;
 import org.springframework.batch.core.repository.JobRepository;
+import org.springframework.batch.poller.DirectPoller;
+import org.springframework.batch.poller.Poller;
+import org.springframework.beans.factory.InitializingBean;
 import org.springframework.integration.MessageTimeoutException;
 import org.springframework.integration.annotation.Aggregator;
 import org.springframework.integration.annotation.MessageEndpoint;
@@ -20,18 +37,22 @@ import org.springframework.messaging.PollableChannel;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 
-import java.util.Collection;
-import java.util.List;
-import java.util.Set;
-
 /**
  * A {@link PartitionHandler} that uses {@link MessageChannel} instances to send instructions to remote workers and
  * receive their responses. The {@link MessageChannel} provides a nice abstraction so that the location of the workers
  * and the transport used to communicate with them can be changed at run time. The communication with the remote workers
  * does not need to be transactional or have guaranteed delivery, so a local thread pool based implementation works as
- * well as a remote web service or JMS implementation. If a remote worker fails or doesn't send a reply message, the job
- * will fail and can be restarted to pick up missing messages and processing. The remote workers need access to the
- * Spring Batch {@link JobRepository} so that the shared state across those restarts can be managed centrally.
+ * well as a remote web service or JMS implementation. If a remote worker fails, the job will fail and can be restarted
+ * to pick up missing messages and processing. The remote workers need access to the Spring Batch {@link JobRepository}
+ * so that the shared state across those restarts can be managed centrally.
+ *
+ * While a {@link org.springframework.messaging.MessageChannel} is used for sending the requests to the workers, the
+ * worker's responses can be obtained in one of two ways:
+ * <ul>
+ *     <li>A reply channel - Slaves will respond with messages that will be aggregated via this component.</li>
+ *     <li>Polling the job repository - Since the state of each slave is maintained independently within the job
+ *     repository, we can poll the store to determine the state without the need of the slaves to formally respond.</li>
+ * </ul>
  *
  * @author Dave Syer
  * @author Will Schipp
@@ -39,7 +60,7 @@ import java.util.Set;
  *
  */
 @MessageEndpoint
-public class MessageChannelPartitionHandler implements PartitionHandler {
+public class MessageChannelPartitionHandler implements PartitionHandler, InitializingBean {
 
 	private static Log logger = LogFactory.getLog(MessageChannelPartitionHandler.class);
 
@@ -49,14 +70,75 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
 
 	private String stepName;
 
+	private long pollInterval = 10000;
+
+	private JobExplorer jobExplorer;
+
+	private boolean pollRepositoryForResults = false;
+
+	private long timeout = -1;
+
+	private DataSource dataSource;
+
 	/**
 	 * pollable channel for the replies
 	 */
 	private PollableChannel replyChannel;
 
+	@Override
 	public void afterPropertiesSet() throws Exception {
 		Assert.notNull(stepName, "A step name must be provided for the remote workers.");
 		Assert.state(messagingGateway != null, "The MessagingOperations must be set");
+
+		pollRepositoryForResults = !(dataSource == null && jobExplorer == null);
+
+		if(pollRepositoryForResults) {
+			logger.debug("MessageChannelPartitionHandler is configured to poll the job repository for slave results");
+		}
+
+		if(dataSource != null && jobExplorer == null) {
+			JobExplorerFactoryBean jobExplorerFactoryBean = new JobExplorerFactoryBean();
+			jobExplorerFactoryBean.setDataSource(dataSource);
+			jobExplorerFactoryBean.afterPropertiesSet();
+			jobExplorer = jobExplorerFactoryBean.getObject();
+		}
+	}
+
+	/**
+	 * When using job repository polling, the time limit to wait.
+	 *
+	 * @param timeout millisconds to wait, defaults to -1 (no timeout).
+	 */
+	public void setTimeout(long timeout) {
+		this.timeout = timeout;
+	}
+
+	/**
+	 * {@link org.springframework.batch.core.explore.JobExplorer} to use to query the job repository.  Either this or
+	 * a {@link javax.sql.DataSource} is required when using job repository polling.
+	 *
+	 * @param jobExplorer {@link org.springframework.batch.core.explore.JobExplorer} to use for lookups
+	 */
+	public void setJobExplorer(JobExplorer jobExplorer) {
+		this.jobExplorer = jobExplorer;
+	}
+
+	/**
+	 * How often to poll the job repository for the status of the slaves.
+	 *
+	 * @param pollInterval milliseconds between polls, defaults to 10000 (10 seconds).
+	 */
+	public void setPollInterval(long pollInterval) {
+		this.pollInterval = pollInterval;
+	}
+
+	/**
+	 * {@link javax.sql.DataSource} pointing to the job repository
+	 *
+	 * @param dataSource {@link javax.sql.DataSource} that points to the job repository's store
+	 */
+	public void setDataSource(DataSource dataSource) {
+		this.dataSource = dataSource;
 	}
 
 	/**
@@ -117,9 +199,9 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
 	 * @see PartitionHandler#handle(StepExecutionSplitter, StepExecution)
 	 */
 	public Collection<StepExecution> handle(StepExecutionSplitter stepExecutionSplitter,
-			StepExecution masterStepExecution) throws Exception {
+			final StepExecution masterStepExecution) throws Exception {
 
-		Set<StepExecution> split = stepExecutionSplitter.split(masterStepExecution, gridSize);
+		final Set<StepExecution> split = stepExecutionSplitter.split(masterStepExecution, gridSize);
 
 		if(CollectionUtils.isEmpty(split)) {
 			return null;
@@ -127,21 +209,76 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
 
 		int count = 0;
 
-		if (replyChannel == null) {
-			replyChannel = new QueueChannel();
+		PollableChannel currentReplyChannel = replyChannel;
+
+		if (!pollRepositoryForResults && currentReplyChannel == null) {
+			currentReplyChannel = new QueueChannel();
 		}//end if
 
 		for (StepExecution stepExecution : split) {
 			Message<StepExecutionRequest> request = createMessage(count++, split.size(), new StepExecutionRequest(
-					stepName, stepExecution.getJobExecutionId(), stepExecution.getId()), replyChannel);
+					stepName, stepExecution.getJobExecutionId(), stepExecution.getId()), currentReplyChannel);
 			if (logger.isDebugEnabled()) {
 				logger.debug("Sending request: " + request);
 			}
 			messagingGateway.send(request);
 		}
 
+		if(!pollRepositoryForResults) {
+			return receiveReplies(currentReplyChannel);
+		}
+		else {
+			return pollReplies(masterStepExecution, split);
+		}
+	}
+
+	private Collection<StepExecution> pollReplies(final StepExecution masterStepExecution, final Set<StepExecution> split) throws Exception {
+		final Collection<StepExecution> result = new ArrayList<StepExecution>(split.size());
+
+		Callable<Collection<StepExecution>> callback = new Callable<Collection<StepExecution>>() {
+			@Override
+			public Collection<StepExecution> call() throws Exception {
+
+				for(Iterator<StepExecution> stepExecutionIterator = split.iterator(); stepExecutionIterator.hasNext(); ) {
+					StepExecution curStepExecution = stepExecutionIterator.next();
+
+					if(!result.contains(curStepExecution)) {
+						StepExecution partitionStepExecution =
+								jobExplorer.getStepExecution(masterStepExecution.getJobExecutionId(), curStepExecution.getId());
+
+						if(!partitionStepExecution.getStatus().isRunning()) {
+							result.add(partitionStepExecution);
+						}
+					}
+				}
+
+				if(logger.isDebugEnabled()) {
+					logger.debug(String.format("Currently waiting on %s partitions to finish", split.size()));
+				}
+
+				if(result.size() == split.size()) {
+					return result;
+				}
+				else {
+					return null;
+				}
+			}
+		};
+
+		Poller<Collection<StepExecution>> poller = new DirectPoller<Collection<StepExecution>>(pollInterval);
+		Future<Collection<StepExecution>> resultsFuture = poller.poll(callback);
+
+		if(timeout >= 0) {
+			return resultsFuture.get(timeout, TimeUnit.MILLISECONDS);
+		}
+		else {
+			return resultsFuture.get();
+		}
+	}
+
+	private Collection<StepExecution> receiveReplies(PollableChannel currentReplyChannel) {
 		@SuppressWarnings("unchecked")
-		Message<Collection<StepExecution>> message = (Message<Collection<StepExecution>>) messagingGateway.receive(replyChannel);
+		Message<Collection<StepExecution>> message = (Message<Collection<StepExecution>>) messagingGateway.receive(currentReplyChannel);
 
 		if(message == null) {
 			throw new MessageTimeoutException("Timeout occurred before all partitions returned");
@@ -149,9 +286,7 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
 			logger.debug("Received replies: " + message);
 		}
 
-		Collection<StepExecution> result = message.getPayload();
-		return result;
-
+		return message.getPayload();
 	}
 
 	private Message<StepExecutionRequest> createMessage(int sequenceNumber, int sequenceSize,

spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java

@@ -1,26 +1,35 @@
 package org.springframework.batch.integration.partition;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.anyLong;
+import static org.mockito.Matchers.anyObject;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.concurrent.TimeoutException;
+
 import org.junit.Test;
+
+import org.springframework.batch.core.BatchStatus;
 import org.springframework.batch.core.JobExecution;
+import org.springframework.batch.core.JobParameters;
 import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.explore.JobExplorer;
 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 java.util.HashSet;
-
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Matchers.anyObject;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 /**
  *
  * @author Will Schipp
@@ -119,4 +128,85 @@ public class MessageChannelPartitionHandlerTests {
 		//execute
 		Collection<StepExecution> executions = messageChannelPartitionHandler.handle(stepExecutionSplitter, masterStepExecution);
 	}
+
+	@Test
+	public void testHandleWithJobRepositoryPolling() throws Exception {
+		//execute with no default set
+		messageChannelPartitionHandler = new MessageChannelPartitionHandler();
+		//mock
+		JobExecution jobExecution = new JobExecution(5l, new JobParameters());
+		StepExecution masterStepExecution = new StepExecution("step1", jobExecution, 1l);
+		StepExecutionSplitter stepExecutionSplitter = mock(StepExecutionSplitter.class);
+		MessagingTemplate operations = mock(MessagingTemplate.class);
+		JobExplorer jobExplorer = mock(JobExplorer.class);
+		//when
+		HashSet<StepExecution> stepExecutions = new HashSet<StepExecution>();
+		StepExecution partition1 = new StepExecution("step1:partition1", jobExecution, 2l);
+		StepExecution partition2 = new StepExecution("step1:partition2", jobExecution, 3l);
+		StepExecution partition3 = new StepExecution("step1:partition3", jobExecution, 4l);
+		StepExecution partition4 = new StepExecution("step1:partition3", jobExecution, 4l);
+		partition1.setStatus(BatchStatus.COMPLETED);
+		partition2.setStatus(BatchStatus.COMPLETED);
+		partition3.setStatus(BatchStatus.STARTED);
+		partition4.setStatus(BatchStatus.COMPLETED);
+		stepExecutions.add(partition1);
+		stepExecutions.add(partition2);
+		stepExecutions.add(partition3);
+		when(stepExecutionSplitter.split((StepExecution) anyObject(), eq(1))).thenReturn(stepExecutions);
+		when(jobExplorer.getStepExecution(eq(5l), anyLong())).thenReturn(partition2, partition1, partition3, partition3, partition3, partition3, partition4);
+
+		//set
+		messageChannelPartitionHandler.setMessagingOperations(operations);
+		messageChannelPartitionHandler.setJobExplorer(jobExplorer);
+		messageChannelPartitionHandler.setStepName("step1");
+		messageChannelPartitionHandler.setPollInterval(500l);
+		messageChannelPartitionHandler.afterPropertiesSet();
+
+		//execute
+		Collection<StepExecution> executions = messageChannelPartitionHandler.handle(stepExecutionSplitter, masterStepExecution);
+		//verify
+		assertNotNull(executions);
+		assertEquals(3, executions.size());
+		assertTrue(executions.contains(partition1));
+		assertTrue(executions.contains(partition2));
+		assertTrue(executions.contains(partition4));
+
+		//verify
+		verify(operations, times(3)).send((Message) anyObject());
+	}
+
+	@Test(expected = TimeoutException.class)
+	public void testHandleWithJobRepositoryPollingTimeout() throws Exception {
+		//execute with no default set
+		messageChannelPartitionHandler = new MessageChannelPartitionHandler();
+		//mock
+		JobExecution jobExecution = new JobExecution(5l, new JobParameters());
+		StepExecution masterStepExecution = new StepExecution("step1", jobExecution, 1l);
+		StepExecutionSplitter stepExecutionSplitter = mock(StepExecutionSplitter.class);
+		MessagingTemplate operations = mock(MessagingTemplate.class);
+		JobExplorer jobExplorer = mock(JobExplorer.class);
+		//when
+		HashSet<StepExecution> stepExecutions = new HashSet<StepExecution>();
+		StepExecution partition1 = new StepExecution("step1:partition1", jobExecution, 2l);
+		StepExecution partition2 = new StepExecution("step1:partition2", jobExecution, 3l);
+		StepExecution partition3 = new StepExecution("step1:partition3", jobExecution, 4l);
+		partition1.setStatus(BatchStatus.COMPLETED);
+		partition2.setStatus(BatchStatus.COMPLETED);
+		partition3.setStatus(BatchStatus.STARTED);
+		stepExecutions.add(partition1);
+		stepExecutions.add(partition2);
+		stepExecutions.add(partition3);
+		when(stepExecutionSplitter.split((StepExecution) anyObject(), eq(1))).thenReturn(stepExecutions);
+		when(jobExplorer.getStepExecution(eq(5l), anyLong())).thenReturn(partition2, partition1, partition3);
+
+		//set
+		messageChannelPartitionHandler.setMessagingOperations(operations);
+		messageChannelPartitionHandler.setJobExplorer(jobExplorer);
+		messageChannelPartitionHandler.setStepName("step1");
+		messageChannelPartitionHandler.setTimeout(1000l);
+		messageChannelPartitionHandler.afterPropertiesSet();
+
+		//execute
+		Collection<StepExecution> executions = messageChannelPartitionHandler.handle(stepExecutionSplitter, masterStepExecution);
+	}
 }

spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/PollingIntegrationTests.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2006-2007 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.batch.integration.partition;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+import java.util.List;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.springframework.batch.core.BatchStatus;
+import org.springframework.batch.core.Job;
+import org.springframework.batch.core.JobExecution;
+import org.springframework.batch.core.JobInstance;
+import org.springframework.batch.core.JobParameters;
+import org.springframework.batch.core.explore.JobExplorer;
+import org.springframework.batch.core.launch.JobLauncher;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+/**
+ * @author Dave Syer
+ * 
+ */
+@ContextConfiguration
+@RunWith(SpringJUnit4ClassRunner.class)
+public class PollingIntegrationTests {
+
+	@Autowired
+	private JobLauncher jobLauncher;
+
+	@Autowired
+	private Job job;
+	
+	@Autowired
+	private JobExplorer jobExplorer;
+
+	@Test
+	public void testSimpleProperties() throws Exception {
+		assertNotNull(jobLauncher);
+	}
+
+	@Test
+	public void testLaunchJob() throws Exception {
+		int before = jobExplorer.getJobInstances(job.getName(), 0, 100).size();
+		assertNotNull(jobLauncher.run(job, new JobParameters()));
+		List<JobInstance> jobInstances = jobExplorer.getJobInstances(job.getName(), 0, 100);
+		int after = jobInstances.size();
+		assertEquals(1, after-before);
+		JobExecution jobExecution = jobExplorer.getJobExecutions(jobInstances.get(jobInstances.size()-1)).get(0);
+		assertEquals(BatchStatus.COMPLETED, jobExecution.getStatus());
+		assertEquals(3, jobExecution.getStepExecutions().size());
+	}
+
+}

spring-batch-integration/src/test/resources/org/springframework/batch/integration/partition/PollingIntegrationTests-context.xml

@@ -0,0 +1,63 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans" xmlns:p="http://www.springframework.org/schema/p"
+	   xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	   xsi:schemaLocation="
+		http://www.springframework.org/schema/integration http://www.springframework.org/schema/integration/spring-integration.xsd
+		http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd
+		http://www.springframework.org/schema/batch http://www.springframework.org/schema/batch/spring-batch.xsd">
+
+	<import resource="classpath:/simple-job-launcher-context.xml" />
+
+	<channel id="requests" xmlns="http://www.springframework.org/schema/integration">
+		<queue />
+	</channel>
+
+	<service-activator ref="stepExecutionRequestHandler" input-channel="requests" output-channel="nullChannel"
+		xmlns="http://www.springframework.org/schema/integration">
+		<poller fixed-delay="10"/>
+	</service-activator>
+
+	<!-- This is the "remote" worker (which in this case is local) -->
+	<bean id="stepExecutionRequestHandler" class="org.springframework.batch.integration.partition.StepExecutionRequestHandler"
+		p:jobExplorer-ref="jobExplorer" p:stepLocator-ref="stepLocator" />
+
+	<bean id="stepLocator" class="org.springframework.batch.integration.partition.BeanFactoryStepLocator" />
+
+	<bean id="jobExplorer" class="org.springframework.batch.core.explore.support.MapJobExplorerFactoryBean">
+		<property name="repositoryFactory" ref="&amp;jobRepository" />
+	</bean>
+
+	<bean id="partitionHandler" class="org.springframework.batch.integration.partition.MessageChannelPartitionHandler">
+		<property name="messagingOperations">
+			<bean class="org.springframework.integration.core.MessagingTemplate">
+				<property name="defaultChannel" ref="requests" />
+			</bean>
+		</property>
+		<property name="jobExplorer" ref="jobExplorer"/>
+		<property name="stepName" value="step1" />
+		<property name="gridSize" value="2" />
+	</bean>
+
+	<job id="job1" xmlns="http://www.springframework.org/schema/batch">
+		<step id="step1-master">
+			<partition handler="partitionHandler" partitioner="partitioner" />
+		</step>
+	</job>
+
+	<bean id="partitioner" class="org.springframework.batch.core.partition.support.SimplePartitioner" />
+
+	<step id="step1" xmlns="http://www.springframework.org/schema/batch">
+		<tasklet>
+			<chunk commit-interval="10">
+				<reader>
+					<bean class="org.springframework.batch.integration.partition.ExampleItemReader" scope="step"
+						xmlns="http://www.springframework.org/schema/beans" />
+				</reader>
+				<writer>
+					<bean class="org.springframework.batch.integration.partition.ExampleItemWriter" xmlns="http://www.springframework.org/schema/beans" />
+				</writer>
+			</chunk>
+		</tasklet>
+	</step>
+
+</beans>