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 269fcdf46..e25314213 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 @@ -150,11 +150,7 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi JobExecution jobExecution = stepExecution.getJobExecution(); - // If this is a restart we must retain the same grid size, ignoring the - // one passed in... - int splitSize = getSplitSize(stepExecution, gridSize); - - Map contexts = partitioner.partition(splitSize); + Map contexts = getContexts(stepExecution, gridSize); Set set = new HashSet(contexts.size()); for (Entry context : contexts.entrySet()) { @@ -177,14 +173,25 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi } - private int getSplitSize(StepExecution stepExecution, int gridSize) { + private Map getContexts(StepExecution stepExecution, int gridSize) { + ExecutionContext context = stepExecution.getExecutionContext(); String key = SimpleStepExecutionSplitter.class.getSimpleName() + ".GRID_SIZE"; - int result = (int) context.getLong(key, gridSize); - context.putLong(key, result); + + // If this is a restart we must retain the same grid size, ignoring the + // one passed in... + int splitSize = (int) context.getLong(key, gridSize); + context.putLong(key, splitSize); + + Map result; if (context.isDirty()) { + // The context changed so we didn't already know the partitions jobRepository.updateExecutionContext(stepExecution); + result = partitioner.partition(splitSize); + } else { + result = new SimplePartitioner().partition(splitSize); } + return result; } diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemReader.java b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemReader.java index ca8ab61ae..bf11cb9da 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemReader.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemReader.java @@ -18,13 +18,31 @@ public class ExampleItemReader implements ItemReader, ItemStream { private int index = 0; + private int min = 0; + + private int max = Integer.MAX_VALUE; + public static volatile boolean fail = false; + /** + * @param min the min to set + */ + public void setMin(int min) { + this.min = min; + } + + /** + * @param max the max to set + */ + public void setMax(int max) { + this.max = max; + } + /** * Reads next record from input */ public String read() throws Exception { - if (index >= input.length) { + if (index >= input.length || index >= max) { return null; } logger.info(String.format("Processing input index=%s, item=%s, in (%s)", index, input[index], this)); @@ -47,7 +65,7 @@ public class ExampleItemReader implements ItemReader, ItemStream { } public void open(ExecutionContext executionContext) throws ItemStreamException { - index = (int) executionContext.getLong("POSITION", 0); + index = (int) executionContext.getLong("POSITION", min); } public void update(ExecutionContext executionContext) throws ItemStreamException { diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemWriter.java b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemWriter.java index 2984e21fe..7b3e626e7 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemWriter.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/ExampleItemWriter.java @@ -1,5 +1,6 @@ package org.springframework.batch.core.partition; +import java.util.ArrayList; import java.util.List; import org.apache.commons.logging.Log; @@ -9,15 +10,26 @@ import org.springframework.batch.item.ItemWriter; /** * Dummy {@link ItemWriter} which only logs data it receives. */ -public class ExampleItemWriter implements ItemWriter { +public class ExampleItemWriter implements ItemWriter { private static final Log log = LogFactory.getLog(ExampleItemWriter.class); + + private static List items = new ArrayList(); + + public static void clear() { + items.clear(); + } + + public static List getItems() { + return items; + } /** * @see ItemWriter#write(List) */ - public void write(List data) throws Exception { + public void write(List data) throws Exception { log.info(data); + items.addAll(data); } } diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/MinMaxPartitioner.java b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/MinMaxPartitioner.java new file mode 100644 index 000000000..a0d93de6b --- /dev/null +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/MinMaxPartitioner.java @@ -0,0 +1,44 @@ +/* + * 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; + +import java.util.Map; + +import org.springframework.batch.core.partition.support.SimplePartitioner; +import org.springframework.batch.item.ExecutionContext; + +/** + * @author Dave Syer + * + */ +public class MinMaxPartitioner extends SimplePartitioner { + + public Map partition(int gridSize) { + Map partition = super.partition(gridSize); + int total = 8; // The number of items in the ExampleItemReader + int range = total/gridSize; + int i = 0; + for (ExecutionContext context : partition.values()) { + int min = (i++)*range; + int max = Math.min(total, (min+1)*range); + context.putInt("min", min); + context.putInt("max", max); + } + return partition; + } + +} diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/RestartIntegrationTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/RestartIntegrationTests.java index 4d3a617c7..506c42796 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/partition/RestartIntegrationTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/partition/RestartIntegrationTests.java @@ -77,9 +77,16 @@ public class RestartIntegrationTests { int beforeMaster = jdbcTemplate.queryForInt("SELECT COUNT(*) from BATCH_STEP_EXECUTION where STEP_NAME='step1:master'"); int beforePartition = jdbcTemplate.queryForInt("SELECT COUNT(*) from BATCH_STEP_EXECUTION where STEP_NAME like 'step1:partition%'"); + ExampleItemWriter.clear(); JobExecution execution = jobLauncher.run(job, jobParameters); assertEquals(BatchStatus.FAILED,execution.getStatus()); + // Only 4 because the others were in the failed step execution + assertEquals(4, ExampleItemWriter.getItems().size()); + + ExampleItemWriter.clear(); assertNotNull(jobLauncher.run(job, jobParameters)); + // Only 4 because the others were processed in the first attempt + assertEquals(4, ExampleItemWriter.getItems().size()); int afterMaster = jdbcTemplate.queryForInt("SELECT COUNT(*) from BATCH_STEP_EXECUTION where STEP_NAME='step1:master'"); int afterPartition = jdbcTemplate.queryForInt("SELECT COUNT(*) from BATCH_STEP_EXECUTION where STEP_NAME like 'step1:partition%'"); diff --git a/spring-batch-core/src/test/resources/org/springframework/batch/core/partition/launch-context.xml b/spring-batch-core/src/test/resources/org/springframework/batch/core/partition/launch-context.xml index ca6e10659..585d51331 100644 --- a/spring-batch-core/src/test/resources/org/springframework/batch/core/partition/launch-context.xml +++ b/spring-batch-core/src/test/resources/org/springframework/batch/core/partition/launch-context.xml @@ -30,7 +30,7 @@ + class="org.springframework.batch.core.partition.MinMaxPartitioner" /> @@ -42,10 +42,8 @@ - + +