Add MessageChannelPartitionHandler

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

@@ -0,0 +1,47 @@
+package org.springframework.batch.integration.partition;
+
+import java.util.Arrays;
+import java.util.Collection;
+
+import org.springframework.batch.core.Step;
+import org.springframework.batch.core.step.StepLocator;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanFactoryAware;
+import org.springframework.beans.factory.ListableBeanFactory;
+import org.springframework.util.Assert;
+
+/**
+ * A {@link StepLocator} implementation that just looks in its enclosing bean
+ * factory for components of type {@link Step}.
+ * 
+ * @author Dave Syer
+ * 
+ */
+public class BeanFactoryStepLocator implements StepLocator, BeanFactoryAware {
+
+	private BeanFactory beanFactory;
+
+	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
+		this.beanFactory = beanFactory;
+	}
+
+	/**
+	 * Look up a bean with the provided name of type {@link Step}.
+	 * @see StepLocator#getStep(String)
+	 */
+	public Step getStep(String stepName) {
+		return (Step) beanFactory.getBean(stepName, Step.class);
+	}
+
+	/**
+	 * Look in the bean factory for all beans of type {@link Step}.
+	 * @throws IllegalStateException if the {@link BeanFactory} is not listable
+	 * @see StepLocator#getStepNames()
+	 */
+	public Collection<String> getStepNames() {
+		Assert.state(beanFactory instanceof ListableBeanFactory, "BeanFactory is not listable.");
+		return Arrays.asList(((ListableBeanFactory) beanFactory).getBeanNamesForType(Step.class));
+	}
+
+}

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

@@ -0,0 +1,138 @@
+package org.springframework.batch.integration.partition;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Set;
+
+import org.springframework.batch.core.Step;
+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.annotation.Aggregator;
+import org.springframework.integration.annotation.MessageEndpoint;
+import org.springframework.integration.core.Message;
+import org.springframework.integration.core.MessageChannel;
+import org.springframework.integration.gateway.MessagingGateway;
+import org.springframework.integration.message.MessageBuilder;
+import org.springframework.util.Assert;
+
+/**
+ * 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.
+ * 
+ * @author Dave Syer
+ * 
+ */
+@MessageEndpoint
+public class MessageChannelPartitionHandler implements PartitionHandler {
+
+	private int gridSize = 1;
+
+	private MessagingGateway messagingGateway;
+
+	private String stepName;
+
+	public void afterPropertiesSet() throws Exception {
+		Assert.notNull(stepName, "A step name must be provided for the remote workers.");
+		Assert.state(messagingGateway != null, "The MessagingGateway must be set");
+	}
+
+	/**
+	 * A pre-configured gateway for sending and receiving messages to the remote
+	 * workers. Using this property allows a large degree of control over the
+	 * timeouts and other properties of the send. It should have channels set up
+	 * internally:
+	 * <ul>
+	 * <li>request channel capable of accepting {@link StepExecutionRequest}
+	 * payloads</li>
+	 * <li>reply channel that returns a list of {@link StepExecution} results</li>
+	 * </ul>
+	 * The timeout for the repoy should be set sufficiently long that the remote
+	 * steps have time to complete.
+	 * 
+	 * @param messagingGateway the {@link MessagingGateway} to set
+	 */
+	public void setMessagingGateway(MessagingGateway messagingGateway) {
+		this.messagingGateway = messagingGateway;
+	}
+
+	/**
+	 * Passed to the {@link StepExecutionSplitter} in the
+	 * {@link #handle(StepExecutionSplitter, StepExecution)} method, instructing
+	 * it how many {@link StepExecution} instances are required, ideally. The
+	 * {@link StepExecutionSplitter} is allowed to ignore the grid size in the
+	 * case of a restart, since the input data partitions must be preserved.
+	 * 
+	 * @param gridSize the number of step executions that will be created
+	 */
+	public void setGridSize(int gridSize) {
+		this.gridSize = gridSize;
+	}
+
+	/**
+	 * The name of the {@link Step} that will be used to execute the partitioned
+	 * {@link StepExecution}. This is a regular Spring Batch step, with all the
+	 * business logic required to complete an execution based on the input
+	 * parameters in its {@link StepExecution} context. The name will be
+	 * translated into a {@link Step} instance by the remote worker.
+	 * 
+	 * @param stepName the name of the {@link Step} instance to execute business
+	 * logic
+	 */
+	public void setStepName(String stepName) {
+		this.stepName = stepName;
+	}
+
+	/**
+	 * @param messages the messages to be aggregated
+	 * @return the list as it was passed in
+	 */
+	@Aggregator(sendPartialResultsOnTimeout = true)
+	public List<?> aggregate(List<?> messages) {
+		return messages;
+	}
+
+	/**
+	 * Sends {@link StepExecutionRequest} objects to the request channel of the
+	 * {@link MessagingGateway}, 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 MessagingGateway} <b>and</b> the aggregator, so that there is a
+	 * good chance of all work being done.
+	 * 
+	 * @see PartitionHandler#handle(StepExecutionSplitter, StepExecution)
+	 */
+	public Collection<StepExecution> handle(StepExecutionSplitter stepExecutionSplitter,
+			StepExecution masterStepExecution) throws Exception {
+
+		Set<StepExecution> split = stepExecutionSplitter.split(masterStepExecution, gridSize);
+		int count = 0;
+		for (StepExecution stepExecution : split) {
+			messagingGateway.send(createMessage(count++, split.size(), new StepExecutionRequest(stepName, stepExecution
+					.getJobExecutionId(), stepExecution.getId())));
+		}
+
+		@SuppressWarnings("unchecked")
+		Collection<StepExecution> result = (Collection<StepExecution>) messagingGateway.receive();
+		return result;
+
+	}
+
+	private Message<StepExecutionRequest> createMessage(int sequenceNumber, int sequenceSize,
+			StepExecutionRequest stepExecutionRequest) {
+		return MessageBuilder.withPayload(stepExecutionRequest).setSequenceNumber(sequenceNumber).setSequenceSize(
+				sequenceSize).setCorrelationId(
+				stepExecutionRequest.getJobExecutionId() + ":" + stepExecutionRequest.getStepName()).build();
+	}
+}

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

@@ -0,0 +1,37 @@
+package org.springframework.batch.integration.partition;
+
+import java.io.Serializable;
+
+public class StepExecutionRequest implements Serializable {
+
+	private final Long stepExecutionId;
+
+	private final String stepName;
+
+	private final Long jobExecutionId;
+
+	public StepExecutionRequest(String stepName, Long jobExecutionId, Long stepExecutionId) {
+		this.stepName = stepName;
+		this.jobExecutionId = jobExecutionId;
+		this.stepExecutionId = stepExecutionId;
+	}
+
+	public Long getJobExecutionId() {
+		return jobExecutionId;
+	}
+
+	public Long getStepExecutionId() {
+		return stepExecutionId;
+	}
+
+	public String getStepName() {
+		return stepName;
+	}
+
+	@Override
+	public String toString() {
+		return String.format("StepExecutionRequest: [jobExecutionId=%d, stepExecutionId=%d, stepName=%s]",
+				jobExecutionId, stepExecutionId, stepName);
+	}
+
+}

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

@@ -0,0 +1,81 @@
+package org.springframework.batch.integration.partition;
+
+import org.springframework.batch.core.BatchStatus;
+import org.springframework.batch.core.JobInterruptedException;
+import org.springframework.batch.core.Step;
+import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.explore.JobExplorer;
+import org.springframework.batch.core.step.StepLocator;
+import org.springframework.integration.annotation.MessageEndpoint;
+import org.springframework.integration.annotation.ServiceActivator;
+
+/**
+ * A {@link MessageEndpoint} that can handle a {@link StepExecutionRequest} and
+ * return a {@link StepExecution} as the result. Typically these need to be
+ * aggregated into a response to a partition handler.
+ * 
+ * @author Dave Syer
+ * 
+ */
+@MessageEndpoint
+public class StepExecutionRequestHandler {
+
+	private JobExplorer jobExplorer;
+
+	private StepLocator stepLocator;
+
+	/**
+	 * Used to locate a {@link Step} to execute for each request.
+	 * @param stepLocator a {@link StepLocator}
+	 */
+	public void setStepLocator(StepLocator stepLocator) {
+		this.stepLocator = stepLocator;
+	}
+
+	/**
+	 * An explorer that should be used to check for {@link StepExecution}
+	 * completion.
+	 * 
+	 * @param jobExplorer a {@link JobExplorer} that is linked to the shared
+	 * repository used by all remote workers.
+	 */
+	public void setJobExplorer(JobExplorer jobExplorer) {
+		this.jobExplorer = jobExplorer;
+	}
+
+	@ServiceActivator
+	public StepExecution handle(StepExecutionRequest request) {
+
+		Long jobExecutionId = request.getJobExecutionId();
+		Long stepExecutionId = request.getStepExecutionId();
+		StepExecution stepExecution = jobExplorer.getStepExecution(jobExecutionId, stepExecutionId);
+		if (stepExecution == null) {
+			// TODO: create new Exception
+			throw new IllegalStateException("No StepExecution could be located for this request: " + request);
+		}
+
+		String stepName = request.getStepName();
+		Step step = stepLocator.getStep(stepName);
+		if (step == null) {
+			// TODO: create new Exception
+			throw new IllegalStateException(String.format("No Step with name [%s] could be located.", stepName));
+		}
+
+		try {
+			step.execute(stepExecution);
+		}
+		catch (JobInterruptedException e) {
+			stepExecution.setStatus(BatchStatus.STOPPED);
+			// TODO: maybe update stepExecution in repository
+		}
+		catch (Throwable e) {
+			stepExecution.addFailureException(e);
+			stepExecution.setStatus(BatchStatus.FAILED);
+			// TODO: maybe update stepExecution in repository
+		}
+
+		return stepExecution;
+
+	}
+
+}