BATCH-1924: Updated SimpleFlow to correctly continue processing on the restart of a STOPPED state

spring-batch-core/src/main/java/org/springframework/batch/core/job/flow/JobFlowExecutor.java

@@ -32,6 +32,7 @@ import org.springframework.batch.core.repository.JobRestartException;
  * execute a flow related to a {@link JobExecution}.
  *
  * @author Dave Syer
+ * @author Michael Minella
  *
  */
 public class JobFlowExecutor implements FlowExecutor {
@@ -59,17 +60,30 @@ public class JobFlowExecutor implements FlowExecutor {
 	@Override
 	public String executeStep(Step step) throws JobInterruptedException, JobRestartException,
 	StartLimitExceededException {
+		boolean isRerun = isStepRestart(step);
 		StepExecution stepExecution = stepHandler.handleStep(step, execution);
 		stepExecutionHolder.set(stepExecution);
+
 		if (stepExecution == null) {
 			return  ExitStatus.COMPLETED.getExitCode();
 		}
 		if (stepExecution.isTerminateOnly()) {
 			throw new JobInterruptedException("Step requested termination: "+stepExecution, stepExecution.getStatus());
 		}
+
+		if(isRerun) {
+			stepExecution.getExecutionContext().put("batch.restart", true);
+		}
+
 		return stepExecution.getExitStatus().getExitCode();
 	}
 
+	private boolean isStepRestart(Step step) {
+		int count = jobRepository.getStepExecutionCount(execution.getJobInstance(), step.getName());
+
+		return count > 0;
+	}
+
 	@Override
 	public void abandonStepExecution() {
 		StepExecution lastStepExecution = stepExecutionHolder.get();

spring-batch-core/src/main/java/org/springframework/batch/core/job/flow/support/SimpleFlow.java

@@ -29,6 +29,7 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.batch.core.JobExecutionException;
 import org.springframework.batch.core.Step;
+import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.job.flow.Flow;
 import org.springframework.batch.core.job.flow.FlowExecution;
 import org.springframework.batch.core.job.flow.FlowExecutionException;
@@ -45,6 +46,7 @@ import org.springframework.beans.factory.InitializingBean;
  * if unambiguous.
  *
  * @author Dave Syer
+ * @author Michael Minella
  * @since 2.0
  */
 public class SimpleFlow implements Flow, InitializingBean {
@@ -139,15 +141,16 @@ public class SimpleFlow implements Flow, InitializingBean {
 		State state = stateMap.get(stateName);
 
 		logger.debug("Resuming state="+stateName+" with status="+status);
+		StepExecution stepExecution = null;
 
 		// Terminate if there are no more states
-		while (state != null && status!=FlowExecutionStatus.STOPPED) {
-
+		while (isFlowContinued(state, status, stepExecution)) {
 			stateName = state.getName();
 
 			try {
 				logger.debug("Handling state="+stateName);
 				status = state.handle(executor);
+				stepExecution = executor.getStepExecution();
 			}
 			catch (FlowExecutionException e) {
 				executor.close(new FlowExecution(stateName, status));
@@ -162,7 +165,6 @@ public class SimpleFlow implements Flow, InitializingBean {
 			logger.debug("Completed state="+stateName+" with status="+status);
 
 			state = nextState(stateName, status);
-
 		}
 
 		FlowExecution result = new FlowExecution(stateName, status);
@@ -171,6 +173,22 @@ public class SimpleFlow implements Flow, InitializingBean {
 
 	}
 
+	private boolean isFlowContinued(State state, FlowExecutionStatus status, StepExecution stepExecution) {
+		boolean continued = true;
+
+		continued = state != null && status!=FlowExecutionStatus.STOPPED;
+
+		if(stepExecution != null) {
+			Boolean reRun = (Boolean) stepExecution.getExecutionContext().get("batch.restart");
+
+			if(reRun != null && reRun && status == FlowExecutionStatus.STOPPED && !state.getName().endsWith(stepExecution.getStepName())) {
+				continued = true;
+			}
+		}
+
+		return continued;
+	}
+
 	/**
 	 * @return the next {@link Step} (or null if this is the end)
 	 * @throws JobExecutionException
@@ -187,7 +205,7 @@ public class SimpleFlow implements Flow, InitializingBean {
 		String next = null;
 		String exitCode = status.getName();
 		for (StateTransition stateTransition : set) {
-			if (stateTransition.matches(exitCode)) {
+			if (stateTransition.matches(exitCode) || (exitCode.equals("PENDING") && stateTransition.matches("STOPPED"))) {
 				if (stateTransition.isEnd()) {
 					// End of job
 					return null;
@@ -207,7 +225,9 @@ public class SimpleFlow implements Flow, InitializingBean {
 					getName(), next));
 		}
 
-		return stateMap.get(next);
+		State state = stateMap.get(next);
+
+		return state;
 
 	}
 

spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StopAndRestartFailedJobParserTests.java

@@ -60,7 +60,7 @@ public class StopAndRestartFailedJobParserTests extends AbstractJobParserTests {
 	}
 
 	private JobExecution launchAndAssert(String stepNames) throws JobInstanceAlreadyCompleteException, JobRestartException,
-			JobExecutionAlreadyRunningException {
+	JobExecutionAlreadyRunningException {
 		JobExecution jobExecution = createJobExecution();
 		job.execute(jobExecution);
 		assertEquals(stepNames, stepNamesList.toString());

spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StopRestartOnCompletedStepJobParserTests.java

@@ -54,7 +54,7 @@ public class StopRestartOnCompletedStepJobParserTests extends AbstractJobParserT
 	}
 
 	private void launchAndAssert(String stepNames) throws JobInstanceAlreadyCompleteException, JobRestartException,
-			JobExecutionAlreadyRunningException {
+	JobExecutionAlreadyRunningException {
 		JobExecution jobExecution = createJobExecution();
 		job.execute(jobExecution);
 		assertEquals(stepNames, stepNamesList.toString());

spring-batch-core/src/test/java/org/springframework/batch/core/configuration/xml/StopRestartOnFailedStepJobParserTests.java

@@ -54,7 +54,7 @@ public class StopRestartOnFailedStepJobParserTests extends AbstractJobParserTest
 	}
 
 	private void launchAndAssert(String stepNames) throws JobInstanceAlreadyCompleteException, JobRestartException,
-			JobExecutionAlreadyRunningException {
+	JobExecutionAlreadyRunningException {
 		JobExecution jobExecution = createJobExecution();
 		job.execute(jobExecution);
 		assertEquals(stepNames, stepNamesList.toString());

spring-batch-core/src/test/java/org/springframework/batch/core/step/RestartInPriorStepTests.java

@@ -0,0 +1,110 @@
+package org.springframework.batch.core.step;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.batch.core.BatchStatus;
+import org.springframework.batch.core.ExitStatus;
+import org.springframework.batch.core.Job;
+import org.springframework.batch.core.JobExecution;
+import org.springframework.batch.core.JobParameters;
+import org.springframework.batch.core.StepContribution;
+import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.explore.JobExplorer;
+import org.springframework.batch.core.job.flow.FlowExecutionStatus;
+import org.springframework.batch.core.job.flow.JobExecutionDecider;
+import org.springframework.batch.core.launch.JobLauncher;
+import org.springframework.batch.core.repository.JobRepository;
+import org.springframework.batch.core.scope.context.ChunkContext;
+import org.springframework.batch.core.step.tasklet.Tasklet;
+import org.springframework.batch.repeat.RepeatStatus;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+/**
+ * @author Michael Minella
+ *
+ */
+@ContextConfiguration
+@RunWith(SpringJUnit4ClassRunner.class)
+public class RestartInPriorStepTests {
+
+	@Autowired
+	private JobRepository jobRepository;
+
+	@Autowired
+	private JobExplorer jobExplorer;
+
+	@Autowired
+	private JobLauncher jobLauncher;
+
+	@Autowired
+	private Job job;
+
+	/**
+	 * @throws java.lang.Exception
+	 */
+	@Before
+	public void setUp() throws Exception {
+	}
+
+	@Test
+	public void test() throws Exception {
+		//
+		// Run 1
+		//
+		JobExecution run1 = jobLauncher.run(job, new JobParameters());
+
+		assertEquals(BatchStatus.STOPPED, run1.getStatus());
+		assertEquals(2, run1.getStepExecutions().size());
+
+		//
+		// Run 2
+		//
+		JobExecution run2 = jobLauncher.run(job, new JobParameters());
+
+		assertEquals(BatchStatus.COMPLETED, run2.getStatus());
+		assertEquals(6, run2.getStepExecutions().size());
+	}
+
+	public static class DecidingTasklet implements Tasklet {
+
+		@Override
+		public RepeatStatus execute(StepContribution contribution,
+				ChunkContext chunkContext) throws Exception {
+			Map<String, Object> context = chunkContext.getStepContext().getJobExecutionContext();
+
+			if(context.get("restart") != null) {
+				contribution.setExitStatus(new ExitStatus("ES3"));
+			} else {
+				chunkContext.getStepContext().setAttribute("restart", true);
+				contribution.setExitStatus(new ExitStatus("ES4"));
+			}
+
+			return RepeatStatus.FINISHED;
+		}
+	}
+
+	public static class CompletionDecider implements JobExecutionDecider {
+
+		private int count = 0;
+
+		@Override
+		public FlowExecutionStatus decide(JobExecution jobExecution,
+				StepExecution stepExecution) {
+			count++;
+
+			if(count > 2) {
+				return new FlowExecutionStatus("END");
+			}
+			else {
+				return new FlowExecutionStatus("CONTINUE");
+			}
+		}
+	}
+}

spring-batch-core/src/test/resources/org/springframework/batch/core/step/RestartInPriorStepTests-context.xml

@@ -0,0 +1,87 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans" 
+	xmlns:util="http://www.springframework.org/schema/beans"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xsi:schemaLocation="http://www.springframework.org/schema/beans
+                    http://www.springframework.org/schema/beans/spring-beans.xsd
+					http://www.springframework.org/schema/util
+					http://www.springframework.org/schema/util/spring-util.xsd
+					http://www.springframework.org/schema/batch
+	                http://www.springframework.org/schema/batch/spring-batch-2.2.xsd">
+	                
+	<bean id="step1ItemReader" class="org.springframework.batch.item.support.ListItemReader">
+		<constructor-arg>
+			<util:list>
+				<util:value>A</util:value>
+			</util:list>
+		</constructor-arg>
+	</bean>
+	
+	<bean id="step1ItemWriter" class="org.springframework.batch.core.configuration.xml.DummyItemWriter"/>
+	                
+	<bean id="step2ItemReader" class="org.springframework.batch.item.support.ListItemReader">
+		<constructor-arg>
+			<util:list>
+				<util:value>A</util:value>
+				<util:value>B</util:value>
+			</util:list>
+		</constructor-arg>
+	</bean>
+	
+	<bean id="step2ItemWriter" class="org.springframework.batch.core.configuration.xml.DummyItemWriter"/>
+	                
+	<bean id="step3ItemReader" class="org.springframework.batch.item.support.ListItemReader">
+		<constructor-arg>
+			<util:list>
+				<util:value>A</util:value>
+				<util:value>B</util:value>
+				<util:value>C</util:value>
+			</util:list>
+		</constructor-arg>
+	</bean>
+	
+	<bean id="step3ItemWriter" class="org.springframework.batch.core.configuration.xml.DummyItemWriter"/>
+	
+	<bean id="decidingTasklet" class="org.springframework.batch.core.step.RestartInPriorStepTests$DecidingTasklet"/>
+	
+	<bean id="completionDecider" class="org.springframework.batch.core.step.RestartInPriorStepTests$CompletionDecider"/>
+	                
+	<job id="restartJob" xmlns="http://www.springframework.org/schema/batch">
+		<step id="step1" next="step2">
+			<tasklet allow-start-if-complete="true">
+				<chunk reader="step1ItemReader" writer="step1ItemWriter" commit-interval="10" />
+			</tasklet>
+		</step>
+		<step id="step2">
+			<tasklet ref="decidingTasklet"/>
+			<next on="ES3" to="step3"/>
+			<stop on="ES4" restart="step4"/>
+		</step>
+		<step id="step3" next="step4">
+			<tasklet>
+				<chunk reader="step2ItemReader" writer="step2ItemWriter" commit-interval="10" />
+			</tasklet>
+		</step>
+		<step id="step4" next="completionDecider">
+			<tasklet>
+				<chunk reader="step3ItemReader" writer="step3ItemWriter" commit-interval="10" />
+			</tasklet>
+		</step>
+		<decision decider="completionDecider" id="completionDecider">
+			<next on="CONTINUE" to="step3"/>
+			<end on="END"/>
+		</decision>
+	</job>
+	
+	<bean id="jobRepository" class="org.springframework.batch.core.repository.support.MapJobRepositoryFactoryBean"/>
+	
+	<bean id="jobExplorer" class="org.springframework.batch.core.explore.support.MapJobExplorerFactoryBean">
+		<property name="repositoryFactory" ref="&amp;jobRepository"/>
+	</bean>
+	
+	<bean id="transactionManager" class="org.springframework.batch.support.transaction.ResourcelessTransactionManager"/>
+	
+	<bean id="jobLauncher" class="org.springframework.batch.core.launch.support.SimpleJobLauncher">
+		<property name="jobRepository" ref="jobRepository"/>
+	</bean>
+</beans>