diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java index 1573792c5..15d1f8bbc 100644 --- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java +++ b/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java @@ -20,12 +20,15 @@ import java.util.Date; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.springframework.batch.core.domain.BatchStatus; +import org.springframework.batch.core.domain.JobInterruptedException; import org.springframework.batch.core.domain.Step; import org.springframework.batch.core.domain.StepExecution; -import org.springframework.batch.core.domain.JobInterruptedException; import org.springframework.batch.core.domain.StepSupport; import org.springframework.batch.core.repository.JobRepository; import org.springframework.batch.core.tasklet.Tasklet; +import org.springframework.batch.execution.scope.SimpleStepContext; +import org.springframework.batch.execution.scope.StepContext; +import org.springframework.batch.execution.scope.StepSynchronizationManager; import org.springframework.batch.io.exception.BatchCriticalException; import org.springframework.batch.repeat.ExitStatus; import org.springframework.batch.repeat.RepeatCallback; @@ -52,14 +55,14 @@ public class TaskletStep extends StepSupport implements InitializingBean { private JobRepository jobRepository; private RepeatListener[] listeners = new RepeatListener[] {}; - + public void setListeners(RepeatListener[] listeners) { - this.listeners = listeners; + this.listeners = listeners; } public void setListener(RepeatListener listener) { listeners = new RepeatListener[] { listener }; - } + } /** * Check mandatory properties. @@ -113,20 +116,24 @@ public class TaskletStep extends StepSupport implements InitializingBean { ExitStatus exitStatus = ExitStatus.FAILED; try { + StepContext parentStepContext = StepSynchronizationManager.getContext(); + final StepContext stepContext = new SimpleStepContext(stepExecution, parentStepContext); + StepSynchronizationManager.register(stepContext); + // We are using the RepeatTemplate as a vehicle for the listener // so it can be set up cheaply here with standard properties. RepeatTemplate template = new RepeatTemplate(); template.setCompletionPolicy(new SimpleCompletionPolicy(1)); template.setListeners(listeners); - exitStatus =template.iterate(new RepeatCallback() { + exitStatus = template.iterate(new RepeatCallback() { public ExitStatus doInIteration(RepeatContext context) throws Exception { return tasklet.execute(); - } + } }); - + updateStatus(stepExecution, BatchStatus.COMPLETED); - + } catch (Exception e) { logger.error("Encountered an error running the tasklet"); @@ -136,7 +143,16 @@ public class TaskletStep extends StepSupport implements InitializingBean { finally { stepExecution.setExitStatus(exitStatus); stepExecution.setEndTime(new Date()); - jobRepository.saveOrUpdate(stepExecution); + try { + jobRepository.saveOrUpdate(stepExecution); + } + catch (Exception e) { + logger.error("Encountered error saving batch meta data. " + + "This job is now in an unknown state and should not be restarted.", e); + } + finally { + StepSynchronizationManager.close(); + } } } diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java index 9a1504ced..8d36ac47e 100644 --- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java +++ b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java @@ -13,6 +13,7 @@ import org.springframework.batch.core.domain.JobSupport; import org.springframework.batch.core.domain.StepExecution; import org.springframework.batch.core.domain.StepSupport; import org.springframework.batch.core.tasklet.Tasklet; +import org.springframework.batch.execution.scope.StepSynchronizationManager; import org.springframework.batch.execution.step.support.JobRepositorySupport; import org.springframework.batch.io.exception.BatchCriticalException; import org.springframework.batch.repeat.ExitStatus; @@ -61,6 +62,14 @@ public class TaskletStepTests extends TestCase { assertNotNull(stepExecution.getEndTime()); } + public void testSuccessfulExecutionWithStepContext() throws Exception { + TaskletStep step = new TaskletStep(new StubTasklet(false, false, true), new JobRepositorySupport()); + step.execute(stepExecution); + assertNotNull(stepExecution.getStartTime()); + assertEquals(ExitStatus.FINISHED, stepExecution.getExitStatus()); + assertNotNull(stepExecution.getEndTime()); + } + public void testFailureExecution() throws Exception { TaskletStep step = new TaskletStep(new StubTasklet(true, false), new JobRepositorySupport()); step.execute(stepExecution); @@ -80,7 +89,6 @@ public class TaskletStepTests extends TestCase { } }); step.execute(stepExecution); - System.err.println(list); assertEquals(2, list.size()); } @@ -97,15 +105,22 @@ public class TaskletStepTests extends TestCase { } } - private class StubTasklet implements Tasklet { + private class StubTasklet implements Tasklet{ private final boolean exitFailure; private final boolean throwException; + private final boolean assertStepContext; + public StubTasklet(boolean exitFailure, boolean throwException) { + this(exitFailure, throwException, false); + } + + public StubTasklet(boolean exitFailure, boolean throwException, boolean assertStepContext) { this.exitFailure = exitFailure; this.throwException = throwException; + this.assertStepContext = assertStepContext; } public ExitStatus execute() throws Exception { @@ -116,6 +131,10 @@ public class TaskletStepTests extends TestCase { if (exitFailure) { return ExitStatus.FAILED; } + + if (assertStepContext) { + assertNotNull(StepSynchronizationManager.getContext()); + } return ExitStatus.FINISHED; }