Implemented the restart functionality for a partitioned step

* Updated parsers to address the allowStartIfComplete flag on
 partitioned steps
* Updated factory beans to address the above flag and injecting a
 JobRepository reference where needed
* Updated how state is restored for a JSR-352 partitioned step
* Fixed the synchronization around the queue used for a
 PartitionCollecotr and PartitionAnalyzer works

spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java

@@ -24,6 +24,7 @@ import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.locks.ReentrantLock;
 
 import javax.batch.api.chunk.listener.RetryProcessListener;
 import javax.batch.api.chunk.listener.RetryReadListener;
@@ -163,6 +164,8 @@ public class StepParserStepFactoryBean<I, O> implements FactoryBean, BeanNameAwa
 
 	private Queue<Serializable> partitionQueue;
 
+	private ReentrantLock partitionLock;
+
 	//
 	// Tasklet Elements
 	//
@@ -248,6 +251,15 @@ public class StepParserStepFactoryBean<I, O> implements FactoryBean, BeanNameAwa
 		this.partitionQueue = queue;
 	}
 
+	/**
+	 * Used to coordinate access to the partition queue between the {@link PartitionCollector} and {@link PartitionAnalyzer}
+	 *
+	 * @param lock a lock that will be locked around accessing the partition queue
+	 */
+	public void setPartitionLock(ReentrantLock lock) {
+		this.partitionLock = lock;
+	}
+
 	/**
 	 * Create a {@link Step} from the configuration provided.
 	 *
@@ -474,6 +486,10 @@ public class StepParserStepFactoryBean<I, O> implements FactoryBean, BeanNameAwa
 
 		enhanceCommonStep(builder);
 		for (ChunkListener listener : chunkListeners) {
+			if(listener instanceof PartitionCollectorAdapter) {
+				((PartitionCollectorAdapter) listener).setPartitionLock(partitionLock);
+			}
+
 			builder.listener(listener);
 
 		}

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/PartitionParser.java

@@ -17,6 +17,7 @@ package org.springframework.batch.core.jsr.configuration.xml;
 
 import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.locks.ReentrantLock;
 
 import org.springframework.batch.core.jsr.configuration.support.BatchArtifact.BatchArtifactType;
 import org.springframework.batch.core.jsr.partition.JsrPartitionHandler;
@@ -55,14 +56,17 @@ public class PartitionParser {
 	private static final String LISTENERS_PROPERTY = "listeners";
 	private static final String THREADS_PROPERTY = "threads";
 	private static final String PARTITIONS_PROPERTY = "partitions";
+	private static final String PARTITION_LOCK_PROPERTY = "partitionLock";
 
 	private final String name;
+	private boolean allowStartIfComplete = false;
 
 	/**
 	 * @param stepName the name of the step that is being partitioned
 	 */
-	public PartitionParser(String stepName) {
+	public PartitionParser(String stepName, boolean allowStartIfComplete) {
 		this.name = stepName;
+		this.allowStartIfComplete = allowStartIfComplete;
 	}
 
 	public void parse(Element element, AbstractBeanDefinition bd, ParserContext parserContext, String stepName) {
@@ -74,6 +78,8 @@ public class PartitionParser {
 
 		MutablePropertyValues properties = partitionHandlerDefinition.getPropertyValues();
 		properties.addPropertyValue(PARTITION_CONTEXT_PROPERTY, new RuntimeBeanReference("batchPropertyContext"));
+		properties.addPropertyValue("jobRepository", new RuntimeBeanReference("jobRepository"));
+		properties.addPropertyValue("allowStartIfComplete", allowStartIfComplete);
 
 		paserMapperElement(element, parserContext, properties);
 		parsePartitionPlan(element, parserContext, stepName, properties);
@@ -85,7 +91,6 @@ public class PartitionParser {
 		String partitionHandlerBeanName = name + ".partitionHandler";
 		registry.registerBeanDefinition(partitionHandlerBeanName, partitionHandlerDefinition);
 		factoryBeanProperties.add("partitionHandler", new RuntimeBeanReference(partitionHandlerBeanName));
-
 	}
 
 	private void parseCollectorElement(Element element,
@@ -98,7 +103,9 @@ public class PartitionParser {
 			// Only needed if a collector is used
 			registerCollectorAnalyzerQueue(parserContext);
 			properties.add(PARTITION_QUEUE_PROPERTY, new RuntimeBeanReference(name + "PartitionQueue"));
+			properties.add(PARTITION_LOCK_PROPERTY, new RuntimeBeanReference(name + "PartitionLock"));
 			factoryBeanProperties.add("partitionQueue", new RuntimeBeanReference(name + "PartitionQueue"));
+			factoryBeanProperties.add("partitionLock", new RuntimeBeanReference(name + "PartitionLock"));
 			String collectorName = collectorElement.getAttribute(REF);
 			factoryBeanProperties.add(LISTENERS_PROPERTY, new RuntimeBeanReference(collectorName));
 			new PropertyParser(collectorName, parserContext, BatchArtifactType.STEP_ARTIFACT, name).parseProperties(collectorElement);
@@ -142,8 +149,11 @@ public class PartitionParser {
 	private void registerCollectorAnalyzerQueue(ParserContext parserContext) {
 		AbstractBeanDefinition partitionQueueDefinition = BeanDefinitionBuilder.genericBeanDefinition(ConcurrentLinkedQueue.class)
 				.getBeanDefinition();
+		AbstractBeanDefinition partitionLockDefinition = BeanDefinitionBuilder.genericBeanDefinition(ReentrantLock.class)
+				.getBeanDefinition();
 
 		parserContext.getRegistry().registerBeanDefinition(name + "PartitionQueue", partitionQueueDefinition);
+		parserContext.getRegistry().registerBeanDefinition(name + "PartitionLock", partitionLockDefinition);
 	}
 
 	protected void parsePartitionPlan(Element element,

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepFactoryBean.java

@@ -31,7 +31,6 @@ import org.springframework.batch.core.jsr.step.builder.JsrBatchletStepBuilder;
 import org.springframework.batch.core.jsr.step.builder.JsrFaultTolerantStepBuilder;
 import org.springframework.batch.core.jsr.step.builder.JsrPartitionStepBuilder;
 import org.springframework.batch.core.jsr.step.builder.JsrSimpleStepBuilder;
-import org.springframework.batch.core.partition.JsrStepExecutionSplitter;
 import org.springframework.batch.core.step.builder.FaultTolerantStepBuilder;
 import org.springframework.batch.core.step.builder.SimpleStepBuilder;
 import org.springframework.batch.core.step.builder.StepBuilder;
@@ -179,8 +178,6 @@ public class StepFactoryBean extends StepParserStepFactoryBean {
 			builder.reducer(reducer);
 		}
 
-		builder.splitter(new JsrStepExecutionSplitter(getName(), getJobRepository()));
-
 		builder.aggregator(getStepExecutionAggergator());
 
 		return builder.build();

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepParser.java

@@ -68,9 +68,11 @@ public class StepParser extends AbstractSingleBeanDefinitionParser {
 		}
 
 		String allowStartIfComplete = element.getAttribute(ALLOW_START_IF_COMPLETE_ATTRIBUTE);
+		boolean allowStartIfCompletValue = false;
 		if(StringUtils.hasText(allowStartIfComplete)) {
 			bd.getPropertyValues().addPropertyValue("allowStartIfComplete",
 					allowStartIfComplete);
+			allowStartIfCompletValue = Boolean.valueOf(allowStartIfComplete);
 		}
 
 		new ListenerParser(StepListenerFactoryBean.class, "listeners").parseListeners(element, parserContext, bd, stepName);
@@ -91,7 +93,7 @@ public class StepParser extends AbstractSingleBeanDefinitionParser {
 				} else if(name.equals(CHUNK_ELEMENT)) {
 					new ChunkParser().parse(nestedElement, bd, parserContext, stepName);
 				} else if(name.equals(PARTITION_ELEMENT)) {
-					new PartitionParser(stepName).parse(nestedElement, bd, parserContext, stepName);
+					new PartitionParser(stepName, allowStartIfCompletValue).parse(nestedElement, bd, parserContext, stepName);
 				}
 			}
 		}

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -28,6 +28,7 @@ import java.util.concurrent.Callable;
 import java.util.concurrent.Future;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.locks.ReentrantLock;
 
 import javax.batch.api.partition.PartitionAnalyzer;
 import javax.batch.api.partition.PartitionCollector;
@@ -36,13 +37,17 @@ import javax.batch.api.partition.PartitionPlan;
 
 import org.springframework.batch.core.BatchStatus;
 import org.springframework.batch.core.ExitStatus;
+import org.springframework.batch.core.JobExecutionException;
 import org.springframework.batch.core.Step;
 import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.jsr.configuration.support.BatchArtifact.BatchArtifactType;
 import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext;
 import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext.BatchPropertyContextEntry;
+import org.springframework.batch.core.partition.JsrStepExecutionSplitter;
 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.item.ExecutionContext;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.core.task.TaskRejectedException;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
@@ -60,12 +65,19 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 
 	// TODO: Replace with proper Channel and Messages once minimum support level for Spring is 4
 	private Queue<Serializable> partitionDataQueue;
+	private ReentrantLock lock;
 	private Step step;
 	private int partitions;
 	private PartitionAnalyzer analyzer;
 	private PartitionMapper mapper;
 	private int threads;
 	private BatchPropertyContext propertyContext;
+	private JobRepository jobRepository;
+	private boolean allowStartIfComplete = false;
+
+	public void setAllowStartIfComplete(boolean allowStartIfComplete) {
+		this.allowStartIfComplete = allowStartIfComplete;
+	}
 
 	/**
 	 * @param queue {@link Queue} to receive the output of the {@link PartitionCollector}
@@ -74,6 +86,10 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 		this.partitionDataQueue = queue;
 	}
 
+	public void setPartitionLock(ReentrantLock lock) {
+		this.lock = lock;
+	}
+
 	/**
 	 * @param context {@link BatchPropertyContext} to resolve partition level step properties
 	 */
@@ -117,6 +133,13 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 		this.partitions = partitions;
 	}
 
+	/**
+	 * @param jobRepository {@link JobRepository}
+	 */
+	public void setJobRepository(JobRepository jobRepository) {
+		this.jobRepository = jobRepository;
+	}
+
 	/* (non-Javadoc)
 	 * @see org.springframework.batch.core.partition.PartitionHandler#handle(org.springframework.batch.core.partition.StepExecutionSplitter, org.springframework.batch.core.StepExecution)
 	 */
@@ -127,24 +150,11 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 		final Set<StepExecution> result = new HashSet<StepExecution>();
 		final ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
 
-		Set<StepExecution> partitionStepExecutions;
+		int stepExecutionCount = jobRepository.getStepExecutionCount(stepExecution.getJobExecution().getJobInstance(), stepExecution.getStepName());
 
-		if(mapper != null) {
-			PartitionPlan plan = mapper.mapPartitions();
-			if(plan.getThreads() > 0) {
-				threads = plan.getThreads();
-			} else if(plan.getPartitions() > 0) {
-				threads = plan.getPartitions();
-			} else {
-				throw new IllegalArgumentException("Either a number of threads or partitions are required");
-			}
+		boolean isRestart = stepExecutionCount > 1;
 
-			partitionStepExecutions = stepSplitter.split(stepExecution, plan.getPartitions());
-			registerPartitionProperties(partitionStepExecutions, plan);
-
-		} else {
-			partitionStepExecutions = stepSplitter.split(stepExecution, partitions);
-		}
+		Set<StepExecution> partitionStepExecutions = splitStepExecution(stepExecution, isRestart);
 
 		taskExecutor.setCorePoolSize(threads);
 		taskExecutor.setMaxPoolSize(threads);
@@ -171,8 +181,17 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 			}
 		}
 
+		processPartitionResults(tasks, result);
+
+		return result;
+	}
+
+	private void processPartitionResults(
+			final List<Future<StepExecution>> tasks,
+			final Set<StepExecution> result) throws Exception {
 		while(true) {
-			synchronized (partitionDataQueue) {
+			try {
+				lock.lock();
 				while(!partitionDataQueue.isEmpty()) {
 					analyzer.analyzeCollectorData(partitionDataQueue.remove());
 				}
@@ -182,10 +201,70 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 				if(tasks.size() == 0) {
 					break;
 				}
+			} finally {
+				if(lock.isHeldByCurrentThread()) {
+					lock.unlock();
+				}
 			}
 		}
+	}
 
-		return result;
+	private Set<StepExecution> splitStepExecution(StepExecution stepExecution,
+			boolean isRestart) throws Exception, JobExecutionException {
+		Set<StepExecution> partitionStepExecutions = new HashSet<StepExecution>();
+		if(isRestart) {
+			if(mapper != null) {
+				PartitionPlan plan = mapper.mapPartitions();
+
+				if(plan.getPartitionsOverride()) {
+					partitionStepExecutions = applyPartitionPlan(stepExecution, plan, false);
+
+					for (StepExecution curStepExecution : partitionStepExecutions) {
+						curStepExecution.setExecutionContext(new ExecutionContext());
+					}
+				} else {
+					Properties[] partitionProps = plan.getPartitionProperties();
+
+					plan = (PartitionPlanState) stepExecution.getExecutionContext().get("partitionPlanState");
+					plan.setPartitionProperties(partitionProps);
+
+					partitionStepExecutions = applyPartitionPlan(stepExecution, plan, true);
+				}
+
+			} else {
+				StepExecutionSplitter stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), true);
+				partitionStepExecutions = stepSplitter.split(stepExecution, partitions);
+			}
+		} else {
+			if(mapper != null) {
+				PartitionPlan plan = mapper.mapPartitions();
+				partitionStepExecutions = applyPartitionPlan(stepExecution, plan, true);
+			} else {
+				StepExecutionSplitter stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), true);
+				partitionStepExecutions = stepSplitter.split(stepExecution, partitions);
+			}
+		}
+		return partitionStepExecutions;
+	}
+
+	private Set<StepExecution> applyPartitionPlan(StepExecution stepExecution,
+			PartitionPlan plan, boolean restoreState) throws JobExecutionException {
+		StepExecutionSplitter stepSplitter;
+		Set<StepExecution> partitionStepExecutions;
+		if(plan.getThreads() > 0) {
+			threads = plan.getThreads();
+		} else if(plan.getPartitions() > 0) {
+			threads = plan.getPartitions();
+		} else {
+			throw new IllegalArgumentException("Either a number of threads or partitions are required");
+		}
+
+		stepExecution.getExecutionContext().put("partitionPlanState", new PartitionPlanState(plan));
+
+		stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), restoreState);
+		partitionStepExecutions = stepSplitter.split(stepExecution, plan.getPartitions());
+		registerPartitionProperties(partitionStepExecutions, plan);
+		return partitionStepExecutions;
 	}
 
 	private void processFinishedPartitions(
@@ -255,13 +334,110 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean {
 		});
 	}
 
+	/* (non-Javadoc)
+	 * @see org.springframework.beans.factory.InitializingBean#afterPropertiesSet()
+	 */
 	@Override
 	public void afterPropertiesSet() throws Exception {
 		Assert.notNull(propertyContext, "A BatchPropertyContext is required");
 		Assert.isTrue(mapper != null || threads > 0, "Either a mapper implementation or the number of partitions/threads is required");
+		Assert.notNull(jobRepository, "A JobRepository is required");
 
 		if(partitionDataQueue == null) {
 			partitionDataQueue = new LinkedBlockingQueue<Serializable>();
 		}
+
+		if(lock == null) {
+			lock = new ReentrantLock();
+		}
+	}
+
+	/**
+	 * Since a {@link PartitionPlan} could provide dynamic data (different results from run to run),
+	 * the batch runtime needs to save off the results for restarts.  This class serves as a container
+	 * used to save off that state.
+	 *
+	 * @author Michael Minella
+	 * @since 3.0
+	 */
+	public static class PartitionPlanState implements PartitionPlan, Serializable {
+
+		private static final long serialVersionUID = 1L;
+		private Properties[] partitionProperties;
+		private int partitions;
+		private int threads;
+
+		/**
+		 * @param plan the {@link PartitionPlan} that is the source of the state
+		 */
+		public PartitionPlanState(PartitionPlan plan) {
+			partitionProperties = plan.getPartitionProperties();
+			partitions = plan.getPartitions();
+			threads = plan.getThreads();
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#getPartitionProperties()
+		 */
+		@Override
+		public Properties[] getPartitionProperties() {
+			return partitionProperties;
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#getPartitions()
+		 */
+		@Override
+		public int getPartitions() {
+			return partitions;
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#getThreads()
+		 */
+		@Override
+		public int getThreads() {
+			return threads;
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#setPartitions(int)
+		 */
+		@Override
+		public void setPartitions(int count) {
+			this.partitions = count;
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#setPartitionsOverride(boolean)
+		 */
+		@Override
+		public void setPartitionsOverride(boolean override) {
+			// Intentional No-op
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#getPartitionsOverride()
+		 */
+		@Override
+		public boolean getPartitionsOverride() {
+			return false;
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#setThreads(int)
+		 */
+		@Override
+		public void setThreads(int count) {
+			this.threads = count;
+		}
+
+		/* (non-Javadoc)
+		 * @see javax.batch.api.partition.PartitionPlan#setPartitionProperties(java.util.Properties[])
+		 */
+		@Override
+		public void setPartitionProperties(Properties[] props) {
+			this.partitionProperties = props;
+		}
 	}
 }

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -17,6 +17,7 @@ package org.springframework.batch.core.jsr.partition;
 
 import java.io.Serializable;
 import java.util.Queue;
+import java.util.concurrent.locks.ReentrantLock;
 
 import javax.batch.api.partition.PartitionCollector;
 import javax.batch.operations.BatchRuntimeException;
@@ -38,6 +39,7 @@ public class PartitionCollectorAdapter implements ChunkListener {
 
 	private PartitionCollector collector;
 	private Queue<Serializable> partitionQueue;
+	private ReentrantLock lock;
 
 	public PartitionCollectorAdapter(Queue<Serializable> queue, PartitionCollector collector) {
 		Assert.notNull(queue, "A thread safe Queue is required");
@@ -47,6 +49,10 @@ public class PartitionCollectorAdapter implements ChunkListener {
 		this.collector = collector;
 	}
 
+	public void setPartitionLock(ReentrantLock lock) {
+		this.lock = lock;
+	}
+
 	@Override
 	public void beforeChunk(ChunkContext context) {
 	}
@@ -54,22 +60,40 @@ public class PartitionCollectorAdapter implements ChunkListener {
 	@Override
 	public void afterChunk(ChunkContext context) {
 		try {
-			synchronized (partitionQueue) {
-				partitionQueue.add(collector.collectPartitionData());
+			if(context.isComplete()) {
+				lock.lock();
+				Serializable collectPartitionData = collector.collectPartitionData();
+
+				if(collectPartitionData != null) {
+					partitionQueue.add(collectPartitionData);
+				}
 			}
 		} catch (Throwable e) {
 			throw new BatchRuntimeException("An error occured while collecting data from the PartionCollector", e);
+		} finally {
+			if(lock.isHeldByCurrentThread()) {
+				lock.unlock();
+			}
 		}
 	}
 
 	@Override
 	public void afterChunkError(ChunkContext context) {
 		try {
-			synchronized (partitionQueue) {
-				partitionQueue.add(collector.collectPartitionData());
+			lock.lock();
+			if(context.isComplete()) {
+				Serializable collectPartitionData = collector.collectPartitionData();
+
+				if(collectPartitionData != null) {
+					partitionQueue.add(collectPartitionData);
+				}
 			}
 		} catch (Throwable e) {
 			throw new BatchRuntimeException("An error occured while collecting data from the PartionCollector", e);
+		} finally {
+			if(lock.isHeldByCurrentThread()) {
+				lock.unlock();
+			}
 		}
 	}
 }

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/PartitionStep.java

@@ -70,7 +70,7 @@ public class PartitionStep extends org.springframework.batch.core.partition.supp
 		}
 
 		// Wait for task completion and then aggregate the results
-		Collection<StepExecution> stepExecutions = getPartitionHandler().handle(getStepExecutionSplitter(), stepExecution);
+		Collection<StepExecution> stepExecutions = getPartitionHandler().handle(null, stepExecution);
 		stepExecution.upgradeStatus(BatchStatus.COMPLETED);
 		stepExecutionAggregator.aggregate(stepExecution, stepExecutions);
 

spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -43,12 +43,18 @@ public class BatchletAdapter implements StoppableTasklet {
 	@Override
 	public RepeatStatus execute(StepContribution contribution,
 			ChunkContext chunkContext) throws Exception {
-		String exitStatus = batchlet.process();
+		String exitStatus;
+		try {
+			exitStatus = batchlet.process();
+		} finally {
+			chunkContext.setComplete();
+		}
 
 		if(StringUtils.hasText(exitStatus)) {
 			contribution.setExitStatus(new ExitStatus(exitStatus));
 		}
 
+
 		return RepeatStatus.FINISHED;
 	}
 

spring-batch-core/src/main/java/org/springframework/batch/core/partition/JsrStepExecutionSplitter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -23,7 +23,9 @@ import org.springframework.batch.core.JobExecution;
 import org.springframework.batch.core.JobExecutionException;
 import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.jsr.launch.JsrJobOperator;
+import org.springframework.batch.core.partition.support.SimpleStepExecutionSplitter;
 import org.springframework.batch.core.repository.JobRepository;
+import org.springframework.batch.item.ExecutionContext;
 
 /**
  * Provides JSR-352 specific behavior for the splitting of {@link StepExecution}s.
@@ -31,14 +33,17 @@ import org.springframework.batch.core.repository.JobRepository;
  * @author Michael Minella
  * @since 3.0
  */
-public class JsrStepExecutionSplitter implements StepExecutionSplitter {
+public class JsrStepExecutionSplitter extends SimpleStepExecutionSplitter {
 
 	private String stepName;
 	private JobRepository jobRepository;
+	private boolean restoreState;
 
-	public JsrStepExecutionSplitter(String stepName, JobRepository jobRepository) {
+	public JsrStepExecutionSplitter(JobRepository jobRepository, boolean allowStartIfComplete, String stepName, boolean restoreState) {
+		super(jobRepository, allowStartIfComplete, stepName, null);
 		this.stepName = stepName;
 		this.jobRepository = jobRepository;
+		this.restoreState = restoreState;
 	}
 
 	@Override
@@ -78,7 +83,10 @@ public class JsrStepExecutionSplitter implements StepExecutionSplitter {
 			String stepName = this.stepName + ":partition" + i;
 			JobExecution curJobExecution = new JobExecution(jobExecution);
 			StepExecution curStepExecution = new StepExecution(stepName, curJobExecution);
-			executions.add(curStepExecution);
+
+			if(!restoreState || getStartable(curStepExecution, new ExecutionContext())) {
+				executions.add(curStepExecution);
+			}
 		}
 
 		jobRepository.addAll(executions);

spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java

@@ -190,7 +190,7 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi
 				set.add(currentStepExecution);
 			}
 		}
-		
+
 		jobRepository.addAll(set);
 
 		return set;
@@ -235,7 +235,7 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi
 		return result;
 	}
 
-	private boolean getStartable(StepExecution stepExecution, ExecutionContext context) throws JobExecutionException {
+	protected boolean getStartable(StepExecution stepExecution, ExecutionContext context) throws JobExecutionException {
 
 		JobInstance jobInstance = stepExecution.getJobExecution().getJobInstance();
 		String stepName = stepExecution.getStepName();

spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandlerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -24,19 +24,26 @@ import java.util.Properties;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
+import javax.batch.api.BatchProperty;
 import javax.batch.api.partition.PartitionAnalyzer;
+import javax.batch.api.partition.PartitionCollector;
 import javax.batch.api.partition.PartitionMapper;
 import javax.batch.api.partition.PartitionPlan;
 import javax.batch.api.partition.PartitionPlanImpl;
+import javax.batch.api.partition.PartitionReducer;
 import javax.batch.runtime.BatchStatus;
+import javax.inject.Inject;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.JobExecution;
+import org.springframework.batch.core.JobInstance;
 import org.springframework.batch.core.JobInterruptedException;
 import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.jsr.JsrTestUtils;
 import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext;
+import org.springframework.batch.core.jsr.step.batchlet.BatchletSupport;
 import org.springframework.batch.core.partition.JsrStepExecutionSplitter;
 import org.springframework.batch.core.repository.JobRepository;
 import org.springframework.batch.core.repository.support.MapJobRepositoryFactoryBean;
@@ -47,14 +54,17 @@ public class JsrPartitionHandlerTests {
 
 	private JsrPartitionHandler handler;
 	private JobRepository repository = new JobRepositorySupport();
-	private StepExecution stepExecution = new StepExecution("step", new JobExecution(1L));
+	private StepExecution stepExecution;
 	private int count;
 	private BatchPropertyContext propertyContext;
 	private JsrStepExecutionSplitter stepSplitter;
 
 	@Before
 	public void setUp() throws Exception {
-		stepSplitter = new JsrStepExecutionSplitter("step1", repository);
+		JobExecution jobExecution = new JobExecution(1L);
+		jobExecution.setJobInstance(new JobInstance(1l, "job"));
+		stepExecution = new StepExecution("step1", jobExecution);
+		stepSplitter = new JsrStepExecutionSplitter(repository, false, "step1", true);
 		Analyzer.collectorData = "";
 		Analyzer.status = "";
 		count = 0;
@@ -70,6 +80,9 @@ public class JsrPartitionHandlerTests {
 		propertyContext = new BatchPropertyContext();
 		handler.setPropertyContext(propertyContext);
 		repository = new MapJobRepositoryFactoryBean().getJobRepository();
+		handler.setJobRepository(repository);
+		MyPartitionReducer.reset();
+		CountingPartitionCollector.reset();
 	}
 
 	@Test
@@ -93,6 +106,15 @@ public class JsrPartitionHandlerTests {
 		}
 
 		handler.setThreads(3);
+
+		try {
+			handler.afterPropertiesSet();
+			fail("JobRepository was not checked for");
+		} catch(IllegalArgumentException iae) {
+			assertEquals("A JobRepository is required", iae.getMessage());
+		}
+
+		handler.setJobRepository(repository);
 		handler.afterPropertiesSet();
 	}
 
@@ -123,7 +145,7 @@ public class JsrPartitionHandlerTests {
 
 		handler.afterPropertiesSet();
 
-		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution);
+		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);
 
 		assertEquals(3, executions.size());
 		assertEquals(3, count);
@@ -144,7 +166,7 @@ public class JsrPartitionHandlerTests {
 
 		handler.afterPropertiesSet();
 
-		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution);
+		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);
 
 		assertEquals(3, executions.size());
 		assertEquals(3, count);
@@ -171,7 +193,7 @@ public class JsrPartitionHandlerTests {
 
 		handler.afterPropertiesSet();
 
-		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution);
+		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);
 
 		assertEquals(3, executions.size());
 		assertEquals(3, count);
@@ -191,7 +213,7 @@ public class JsrPartitionHandlerTests {
 		handler.setPartitionAnalyzer(new Analyzer());
 		handler.afterPropertiesSet();
 
-		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution);
+		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);
 
 		assertEquals(2, executions.size());
 		assertEquals(2, count);
@@ -199,6 +221,159 @@ public class JsrPartitionHandlerTests {
 		assertEquals("COMPLETEDdone", Analyzer.status);
 	}
 
+	@Test
+	public void testRestartNoOverride() throws Exception {
+		javax.batch.runtime.JobExecution execution1 = JsrTestUtils.runJob("jsrPartitionHandlerRestartWithOverrideJob", null, 1000000l);
+		assertEquals(BatchStatus.FAILED, execution1.getBatchStatus());
+		assertEquals(1, MyPartitionReducer.beginCount);
+		assertEquals(0, MyPartitionReducer.beforeCount);
+		assertEquals(1, MyPartitionReducer.rollbackCount);
+		assertEquals(1, MyPartitionReducer.afterCount);
+		assertEquals(3, CountingPartitionCollector.collected);
+
+		MyPartitionReducer.reset();
+		CountingPartitionCollector.reset();
+
+		javax.batch.runtime.JobExecution execution2 = JsrTestUtils.restartJob(execution1.getExecutionId(), null, 1000000l);
+		assertEquals(BatchStatus.COMPLETED, execution2.getBatchStatus());
+		assertEquals(1, MyPartitionReducer.beginCount);
+		assertEquals(1, MyPartitionReducer.beforeCount);
+		assertEquals(0, MyPartitionReducer.rollbackCount);
+		assertEquals(1, MyPartitionReducer.afterCount);
+		assertEquals(1, CountingPartitionCollector.collected);
+	}
+
+
+	@Test
+	public void testRestartOverride() throws Exception {
+		Properties jobParameters = new Properties();
+		jobParameters.put("mapper.override", "true");
+
+		javax.batch.runtime.JobExecution execution1 = JsrTestUtils.runJob("jsrPartitionHandlerRestartWithOverrideJob", jobParameters, 1000000l);
+		assertEquals(BatchStatus.FAILED, execution1.getBatchStatus());
+		assertEquals(1, MyPartitionReducer.beginCount);
+		assertEquals(0, MyPartitionReducer.beforeCount);
+		assertEquals(1, MyPartitionReducer.rollbackCount);
+		assertEquals(1, MyPartitionReducer.afterCount);
+		assertEquals(3, CountingPartitionCollector.collected);
+
+		MyPartitionReducer.reset();
+		CountingPartitionCollector.reset();
+
+		javax.batch.runtime.JobExecution execution2 = JsrTestUtils.restartJob(execution1.getExecutionId(), jobParameters, 1000000l);
+		assertEquals(BatchStatus.COMPLETED, execution2.getBatchStatus());
+		assertEquals(1, MyPartitionReducer.beginCount);
+		assertEquals(1, MyPartitionReducer.beforeCount);
+		assertEquals(0, MyPartitionReducer.rollbackCount);
+		assertEquals(1, MyPartitionReducer.afterCount);
+		assertEquals(5, CountingPartitionCollector.collected);
+	}
+
+	public static class CountingPartitionCollector implements PartitionCollector {
+
+		public static int collected = 0;
+
+		public static void reset() {
+			collected = 0;
+		}
+
+		@Override
+		public Serializable collectPartitionData() throws Exception {
+			collected++;
+
+			return null;
+		}
+	}
+
+	public static class MyPartitionReducer implements PartitionReducer {
+
+		public static int beginCount = 0;
+		public static int beforeCount = 0;
+		public static int rollbackCount = 0;
+		public static int afterCount = 0;
+
+		public static void reset() {
+			beginCount = 0;
+			beforeCount = 0;
+			rollbackCount = 0;
+			afterCount = 0;
+		}
+
+		@Override
+		public void beginPartitionedStep() throws Exception {
+			beginCount++;
+		}
+
+		@Override
+		public void beforePartitionedStepCompletion() throws Exception {
+			beforeCount++;
+		}
+
+		@Override
+		public void rollbackPartitionedStep() throws Exception {
+			rollbackCount++;
+		}
+
+		@Override
+		public void afterPartitionedStepCompletion(PartitionStatus status)
+				throws Exception {
+			afterCount++;
+		}
+	}
+
+	public static class MyPartitionMapper implements PartitionMapper {
+
+		private static int count = 0;
+
+		@Inject
+		@BatchProperty
+		String overrideString = "false";
+
+		@Override
+		public PartitionPlan mapPartitions() throws Exception {
+			count++;
+
+			PartitionPlan plan = new PartitionPlanImpl();
+
+			if(count % 2 == 1) {
+				plan.setPartitions(3);
+				plan.setThreads(3);
+			} else {
+				plan.setPartitions(5);
+				plan.setThreads(5);
+			}
+
+			plan.setPartitionsOverride(Boolean.valueOf(overrideString));
+
+			Properties[] props = new Properties[3];
+			props[0] = new Properties();
+			props[1] = new Properties();
+			props[2] = new Properties();
+
+			if(count % 2 == 1) {
+				props[1].put("fail", "true");
+			}
+
+			plan.setPartitionProperties(props);
+			return plan;
+		}
+	}
+
+	public static class MyBatchlet extends BatchletSupport {
+		@Inject
+		@BatchProperty
+		String fail;
+
+		@Override
+		public String process() {
+			if("true".equalsIgnoreCase(fail)) {
+				throw new RuntimeException("Expected");
+			}
+
+			return null;
+		}
+	}
+
 	public static class Analyzer implements PartitionAnalyzer {
 
 		public static String collectorData;

spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrStepExecutionSplitterTests.java

@@ -33,7 +33,7 @@ public class JsrStepExecutionSplitterTests {
 
 	@Before
 	public void setUp() throws Exception {
-		splitter = new JsrStepExecutionSplitter("step1", new JobRepositorySupport());
+		splitter = new JsrStepExecutionSplitter(new JobRepositorySupport(), false, "step1", true);
 	}
 
 	@Test

spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapterTests.java

@@ -20,10 +20,12 @@ import static org.junit.Assert.assertEquals;
 import java.io.Serializable;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.locks.ReentrantLock;
 
 import javax.batch.api.partition.PartitionCollector;
 
 import org.junit.Test;
+import org.springframework.batch.core.scope.context.ChunkContext;
 
 public class PartitionCollectorAdapterTests {
 
@@ -43,9 +45,14 @@ public class PartitionCollectorAdapterTests {
 			}
 		});
 
-		adapter.afterChunk(null);
-		adapter.afterChunkError(null);
-		adapter.afterChunk(null);
+		adapter.setPartitionLock(new ReentrantLock());
+
+		ChunkContext context = new ChunkContext(null);
+		context.setComplete();
+
+		adapter.afterChunk(context);
+		adapter.afterChunkError(context);
+		adapter.afterChunk(context);
 
 		assertEquals(3, dataQueue.size());
 		assertEquals("0", dataQueue.remove());

spring-batch-core/src/test/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapterTests.java

@@ -14,6 +14,7 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.StepContribution;
+import org.springframework.batch.core.scope.context.ChunkContext;
 import org.springframework.batch.repeat.RepeatStatus;
 
 public class BatchletAdapterTests {
@@ -38,7 +39,7 @@ public class BatchletAdapterTests {
 
 	@Test
 	public void testExecuteNoExitStatus() throws Exception {
-		assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, null));
+		assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, new ChunkContext(null)));
 
 		verify(delegate).process();
 	}
@@ -47,7 +48,7 @@ public class BatchletAdapterTests {
 	public void testExecuteWithExitStatus() throws Exception {
 		when(delegate.process()).thenReturn("my exit status");
 
-		assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, null));
+		assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, new ChunkContext(null)));
 
 		verify(delegate).process();
 		verify(contribution).setExitStatus(new ExitStatus("my exit status"));

spring-batch-core/src/test/resources/META-INF/batch-jobs/jsrPartitionHandlerRestartWithOverrideJob.xml

@@ -0,0 +1,15 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<job id="job_partitioned_1step" xmlns="http://xmlns.jcp.org/xml/ns/javaee" version="1.0">
+	<step id="step1">
+		<batchlet ref="org.springframework.batch.core.jsr.partition.JsrPartitionHandlerTests$MyBatchlet"/>
+	    <partition>
+	    	<mapper ref="org.springframework.batch.core.jsr.partition.JsrPartitionHandlerTests$MyPartitionMapper">
+	    		<properties>
+	    			<property name="overrideString" value="#{jobParameters['mapper.override']}"/>
+	    		</properties>
+	    	</mapper>
+	    	<collector ref="org.springframework.batch.core.jsr.partition.JsrPartitionHandlerTests$CountingPartitionCollector"/>
+	    	<reducer ref="org.springframework.batch.core.jsr.partition.JsrPartitionHandlerTests$MyPartitionReducer"/>
+	    </partition>
+	</step>
+</job>