diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/PartitionNameProvider.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/PartitionNameProvider.java new file mode 100644 index 000000000..fa1457302 --- /dev/null +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/PartitionNameProvider.java @@ -0,0 +1,38 @@ +/* + * Copyright 2006-2009 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.core.partition.support; + +import java.util.Collection; + +/** + *

+ * Optional interface for {@link Partitioner} implementations that need to use a + * custom naming scheme for partitions. It is not necessary to implement this + * interface if a partitioner extends {@link SimplePartitioner} and re-uses the + * default partition names. + *

+ * + * @author Dave Syer + * + * @since 2.1.3 + * + */ +public interface PartitionNameProvider { + + Collection getPartitionNames(int gridSize); + +} diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java index 1642d1974..9a81c766b 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java @@ -16,6 +16,8 @@ package org.springframework.batch.core.partition.support; +import java.util.Collection; +import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Set; @@ -189,7 +191,23 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi result = partitioner.partition(splitSize); } else { - result = new SimplePartitioner().partition(splitSize); + /* + * We need to return the same keys as the original (failed) + * execution, but the execution contexts will be discarded so they + * can be empty. + */ + if (partitioner instanceof PartitionNameProvider) { + result = new HashMap(); + Collection names = ((PartitionNameProvider) partitioner).getPartitionNames(splitSize); + for (String name : names) { + result.put(name, new ExecutionContext()); + } + } + else { + // If no names are provided, assume they follow the default + // pattern. + result = new SimplePartitioner().partition(splitSize); + } } return result; @@ -245,9 +263,13 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi return true; } - if (stepStatus == BatchStatus.STARTED || stepStatus == BatchStatus.STARTING || stepStatus == BatchStatus.STOPPING) { - throw new JobExecutionException("Cannot restart step from " + stepStatus + " status. " - + "The old execution may still be executing, so you may need to verify manually that this is the case."); + if (stepStatus == BatchStatus.STARTED || stepStatus == BatchStatus.STARTING + || stepStatus == BatchStatus.STOPPING) { + throw new JobExecutionException( + "Cannot restart step from " + + stepStatus + + " status. " + + "The old execution may still be executing, so you may need to verify manually that this is the case."); } throw new JobExecutionException("Cannot restart step from " + stepStatus + " status. " diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/PartitionStepTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/PartitionStepTests.java index e9e8477c3..c5d71ee30 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/PartitionStepTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/PartitionStepTests.java @@ -19,7 +19,9 @@ import static org.junit.Assert.assertEquals; import java.util.Arrays; import java.util.Collection; +import java.util.Date; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.Before; import org.junit.Test; @@ -102,6 +104,53 @@ public class PartitionStepTests { assertEquals(BatchStatus.FAILED, stepExecution.getStatus()); } + @Test + public void testRestartStepExecution() throws Exception { + final AtomicBoolean started = new AtomicBoolean(false); + step.setStepExecutionSplitter(new SimpleStepExecutionSplitter(jobRepository, remote, new SimplePartitioner())); + step.setPartitionHandler(new PartitionHandler() { + public Collection handle(StepExecutionSplitter stepSplitter, StepExecution stepExecution) + throws Exception { + Set executions = stepSplitter.split(stepExecution, 2); + if (!started.get()) { + started.set(true); + for (StepExecution execution : executions) { + execution.setStatus(BatchStatus.FAILED); + execution.setExitStatus(ExitStatus.FAILED); + execution.getExecutionContext().putString("foo", execution.getStepName()); + } + } + else { + for (StepExecution execution : executions) { + // On restart the execution context should have been restored + assertEquals(execution.getStepName(), execution.getExecutionContext().getString("foo")); + } + } + for (StepExecution execution : executions) { + jobRepository.update(execution); + jobRepository.updateExecutionContext(execution); + } + return executions; + } + }); + step.afterPropertiesSet(); + JobExecution jobExecution = jobRepository.createJobExecution("vanillaJob", new JobParameters()); + StepExecution stepExecution = jobExecution.createStepExecution("foo"); + jobRepository.add(stepExecution); + step.execute(stepExecution); + jobExecution.setStatus(BatchStatus.FAILED); + jobExecution.setEndTime(new Date()); + jobRepository.update(jobExecution); + // Now restart... + jobExecution = jobRepository.createJobExecution("vanillaJob", new JobParameters()); + stepExecution = jobExecution.createStepExecution("foo"); + jobRepository.add(stepExecution); + step.execute(stepExecution); + // one master and two workers + assertEquals(3, stepExecution.getJobExecution().getStepExecutions().size()); + assertEquals(BatchStatus.COMPLETED, stepExecution.getStatus()); + } + @Test public void testStoppedStepExecution() throws Exception { step.setStepExecutionSplitter(new SimpleStepExecutionSplitter(jobRepository, remote, new SimplePartitioner())); diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitterTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitterTests.java index 265ab5ebc..247dafb88 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitterTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitterTests.java @@ -4,6 +4,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; +import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.Date; import java.util.Map; @@ -71,6 +73,26 @@ public class SimpleStepExecutionSplitterTests { assertEquals(2, provider.split(stepExecution, 3).size()); } + @Test + public void testRememberPartitionNames() throws Exception { + class CustomPartitioner implements Partitioner, PartitionNameProvider { + public Map partition(int gridSize) { + return Collections.singletonMap("foo", new ExecutionContext()); + } + public Collection getPartitionNames(int gridSize) { + return Arrays.asList("foo"); + } + } + SimpleStepExecutionSplitter provider = new SimpleStepExecutionSplitter(jobRepository, step, + new CustomPartitioner()); + Set split = provider.split(stepExecution, 2); + assertEquals(1, split.size()); + assertEquals("step:foo", split.iterator().next().getStepName()); + stepExecution = update(split, stepExecution, BatchStatus.FAILED); + split = provider.split(stepExecution, 2); + assertEquals("step:foo", split.iterator().next().getStepName()); + } + @Test public void testGetStepName() { SimpleStepExecutionSplitter provider = new SimpleStepExecutionSplitter(jobRepository, step,