Terminates the job if a flag is set on afterStep

spring-projects · Sep 23, 2023 · 5c5a70d · 5c5a70d
1 parent 6eeaf4d
commit 5c5a70d
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/step/AbstractStep.java b/spring-batch-core/src/main/java/org/springframework/batch/core/step/AbstractStep.java
@@ -63,6 +63,7 @@
  * @author Michael Minella
  * @author Chris Schaefer
  * @author Mahmoud Ben Hassine
+ * @author SungBeom Um
  */
 public abstract class AbstractStep implements Step, InitializingBean, BeanNameAware {
 
@@ -268,6 +269,11 @@ public final void execute(StepExecution stepExecution)
 				exitStatus = exitStatus.and(stepExecution.getExitStatus());
 				stepExecution.setExitStatus(exitStatus);
 				exitStatus = exitStatus.and(getCompositeListener().afterStep(stepExecution));
+
+				// Check if someone is trying to stop us
+				if (stepExecution.isTerminateOnly()) {
+					stepExecution.upgradeStatus(BatchStatus.STOPPED);
+				}
 			}
 			catch (Exception e) {
 				logger.error(String.format("Exception in afterStep callback in step %s in job %s", name,

diff --git a/...ng-batch-core/src/test/java/org/springframework/batch/core/step/NonAbstractStepTests.java b/...ng-batch-core/src/test/java/org/springframework/batch/core/step/NonAbstractStepTests.java
@@ -358,6 +358,42 @@ public void updateExecutionContext(StepExecution stepExecution) {
 		assertEquals(ExitStatus.UNKNOWN, execution.getExitStatus());
 	}
 
+	@Test
+	void testStoppedOnAfterStep() throws Exception {
+		tested = new EventTrackingStep();
+		tested.setJobRepository(repository);
+		StepExecutionListener listener3 = new EventTrackingListener("listener3") {
+			@Override
+			public ExitStatus afterStep(StepExecution stepExecution) {
+				super.afterStep(stepExecution);
+				stepExecution.setTerminateOnly();
+				stepExecution.setExitStatus(ExitStatus.STOPPED);
+				return stepExecution.getExitStatus();
+			}
+		};
+		tested.setStepExecutionListeners(new StepExecutionListener[] { listener1, listener2, listener3 });
+
+		tested.execute(execution);
+		assertEquals(BatchStatus.STOPPED, execution.getStatus());
+
+		int i = 0;
+		assertEquals("listener1#beforeStep", events.get(i++));
+		assertEquals("listener2#beforeStep", events.get(i++));
+		assertEquals("listener3#beforeStep", events.get(i++));
+		assertEquals("open", events.get(i++));
+		assertEquals("doExecute", events.get(i++));
+		assertEquals("listener3#afterStep(COMPLETED)", events.get(i++));
+		assertEquals("listener2#afterStep(STOPPED)", events.get(i++));
+		assertEquals("listener1#afterStep(STOPPED)", events.get(i++));
+		assertEquals("close", events.get(i++));
+		assertEquals(9, events.size());
+
+		assertEquals("STOPPED", execution.getExitStatus().getExitCode());
+
+		assertTrue(repository.saved.containsKey("afterStep"),
+			"Execution context modifications made by listener should be persisted");
+	}
+
 	/**
 	 * JobRepository is a required property.
 	 */