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 f85e59285..3aefe17e7 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 @@ -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: + * * * @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 handle(StepExecutionSplitter stepExecutionSplitter, - StepExecution masterStepExecution) throws Exception { + final StepExecution masterStepExecution) throws Exception { - Set split = stepExecutionSplitter.split(masterStepExecution, gridSize); + final Set 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 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 pollReplies(final StepExecution masterStepExecution, final Set split) throws Exception { + final Collection result = new ArrayList(split.size()); + + Callable> callback = new Callable>() { + @Override + public Collection call() throws Exception { + + for(Iterator 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> poller = new DirectPoller>(pollInterval); + Future> resultsFuture = poller.poll(callback); + + if(timeout >= 0) { + return resultsFuture.get(timeout, TimeUnit.MILLISECONDS); + } + else { + return resultsFuture.get(); + } + } + + private Collection receiveReplies(PollableChannel currentReplyChannel) { @SuppressWarnings("unchecked") - Message> message = (Message>) messagingGateway.receive(replyChannel); + Message> message = (Message>) 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 result = message.getPayload(); - return result; - + return message.getPayload(); } private Message createMessage(int sequenceNumber, int sequenceSize, diff --git a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java index 2491e1a39..32227c75b 100644 --- a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java +++ b/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 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 stepExecutions = new HashSet(); + 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 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 stepExecutions = new HashSet(); + 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 executions = messageChannelPartitionHandler.handle(stepExecutionSplitter, masterStepExecution); + } } diff --git a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/PollingIntegrationTests.java b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/PollingIntegrationTests.java new file mode 100644 index 000000000..9404ee44c --- /dev/null +++ b/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 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()); + } + +} diff --git a/spring-batch-integration/src/test/resources/org/springframework/batch/integration/partition/PollingIntegrationTests-context.xml b/spring-batch-integration/src/test/resources/org/springframework/batch/integration/partition/PollingIntegrationTests-context.xml new file mode 100644 index 000000000..4ce3a247c --- /dev/null +++ b/spring-batch-integration/src/test/resources/org/springframework/batch/integration/partition/PollingIntegrationTests-context.xml @@ -0,0 +1,63 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +