Fix infinite loop in FlowBuilder#next()

injae-kim · fmbenhassine · commit 4a67b22e9040 · 2023-11-14T16:43:05.000+01:00
Resolves #4432
diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/job/builder/FlowBuilder.java b/spring-batch-core/src/main/java/org/springframework/batch/core/job/builder/FlowBuilder.java
@@ -48,6 +48,7 @@
  * @author Dave Syer
  * @author Michael Minella
  * @author Mahmoud Ben Hassine
+ * @author Injae Kim
  * @since 2.2
  * @param <Q> the type of object returned by the builder (by default a Flow)
  *
@@ -107,7 +108,8 @@ public Q build() {
 
 	/**
 	 * Transition to the next step on successful completion of the current step. All other
-	 * outcomes are treated as failures.
+	 * outcomes are treated as failures. If no steps are registered yet, then this method
+	 * will behave in the same way as {@link #start(Step)}.
 	 * @param step the next step
 	 * @return this to enable chaining
 	 */
@@ -250,26 +252,32 @@ private void doNext(Object input) {
 		if (this.currentState == null) {
 			doStart(input);
 		}
-		State next = createState(input);
-		addTransition("COMPLETED", next);
-		addTransition("*", failedState);
-		this.currentState = next;
+		else {
+			State next = createState(input);
+			addTransition("COMPLETED", next);
+			addTransition("*", failedState);
+			this.currentState = next;
+		}
 	}
 
 	private void doStart(Object input) {
 		if (this.currentState != null) {
 			doFrom(input);
 		}
-		this.currentState = createState(input);
+		else {
+			this.currentState = createState(input);
+		}
 	}
 
 	private void doFrom(Object input) {
 		if (currentState == null) {
 			doStart(input);
 		}
-		State state = createState(input);
-		tos.put(currentState.getName(), currentState);
-		this.currentState = state;
+		else {
+			State state = createState(input);
+			tos.put(currentState.getName(), currentState);
+			this.currentState = state;
+		}
 	}
 
 	private State createState(Object input) {
diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/job/builder/FlowBuilderTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/job/builder/FlowBuilderTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2022 the original author or authors.
+ * Copyright 2012-2023 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.
@@ -19,6 +19,7 @@
 
 import org.junit.jupiter.api.Test;
 
+import org.springframework.batch.core.BatchStatus;
 import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.JobExecution;
 import org.springframework.batch.core.JobInterruptedException;
@@ -34,26 +35,79 @@
 import org.springframework.batch.core.step.StepSupport;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 
 /**
  * @author Dave Syer
  * @author Michael Minella
  * @author Mahmoud Ben Hassine
+ * @author Injae Kim
  *
  */
 class FlowBuilderTests {
 
 	@Test
-	void test() throws Exception {
+	void testNext() throws Exception {
 		FlowBuilder<Flow> builder = new FlowBuilder<>("flow");
 		JobRepository jobRepository = new JobRepositorySupport();
 		JobExecution execution = jobRepository.createJobExecution("foo", new JobParameters());
-		builder.start(new StepSupport("step") {
-			@Override
-			public void execute(StepExecution stepExecution)
-					throws JobInterruptedException, UnexpectedJobExecutionException {
-			}
-		}).end().start(new JobFlowExecutor(jobRepository, new SimpleStepHandler(jobRepository), execution));
+
+		builder.next(createCompleteStep("stepA"))
+			.end()
+			.start(new JobFlowExecutor(jobRepository, new SimpleStepHandler(jobRepository), execution));
+
+		Iterator<StepExecution> stepExecutions = execution.getStepExecutions().iterator();
+		assertEquals(stepExecutions.next().getStepName(), "stepA");
+		assertFalse(stepExecutions.hasNext());
+	}
+
+	@Test
+	void testMultipleNext() throws Exception {
+		FlowBuilder<Flow> builder = new FlowBuilder<>("flow");
+		JobRepository jobRepository = new JobRepositorySupport();
+		JobExecution execution = jobRepository.createJobExecution("foo", new JobParameters());
+
+		builder.next(createCompleteStep("stepA"))
+			.next(createCompleteStep("stepB"))
+			.next(createCompleteStep("stepC"))
+			.end()
+			.start(new JobFlowExecutor(jobRepository, new SimpleStepHandler(jobRepository), execution));
+
+		Iterator<StepExecution> stepExecutions = execution.getStepExecutions().iterator();
+		assertEquals(stepExecutions.next().getStepName(), "stepA");
+		assertEquals(stepExecutions.next().getStepName(), "stepB");
+		assertEquals(stepExecutions.next().getStepName(), "stepC");
+		assertFalse(stepExecutions.hasNext());
+	}
+
+	@Test
+	void testStart() throws Exception {
+		FlowBuilder<Flow> builder = new FlowBuilder<>("flow");
+		JobRepository jobRepository = new JobRepositorySupport();
+		JobExecution execution = jobRepository.createJobExecution("foo", new JobParameters());
+
+		builder.start(createCompleteStep("stepA"))
+			.end()
+			.start(new JobFlowExecutor(jobRepository, new SimpleStepHandler(jobRepository), execution));
+
+		Iterator<StepExecution> stepExecutions = execution.getStepExecutions().iterator();
+		assertEquals(stepExecutions.next().getStepName(), "stepA");
+		assertFalse(stepExecutions.hasNext());
+	}
+
+	@Test
+	void testFrom() throws Exception {
+		FlowBuilder<Flow> builder = new FlowBuilder<>("flow");
+		JobRepository jobRepository = new JobRepositorySupport();
+		JobExecution execution = jobRepository.createJobExecution("foo", new JobParameters());
+
+		builder.from(createCompleteStep("stepA"))
+			.end()
+			.start(new JobFlowExecutor(jobRepository, new SimpleStepHandler(jobRepository), execution));
+
+		Iterator<StepExecution> stepExecutions = execution.getStepExecutions().iterator();
+		assertEquals(stepExecutions.next().getStepName(), "stepA");
+		assertFalse(stepExecutions.hasNext());
 	}
 
 	@Test
@@ -66,7 +120,7 @@ void testTransitionOrdering() throws Exception {
 			@Override
 			public void execute(StepExecution stepExecution)
 					throws JobInterruptedException, UnexpectedJobExecutionException {
-				stepExecution.setExitStatus(new ExitStatus("FAILED"));
+				stepExecution.setExitStatus(ExitStatus.FAILED);
 			}
 		};
 
@@ -94,10 +148,20 @@ public void execute(StepExecution stepExecution)
 			.start(new JobFlowExecutor(jobRepository, new SimpleStepHandler(jobRepository), execution));
 
 		Iterator<StepExecution> stepExecutions = execution.getStepExecutions().iterator();
-		StepExecution stepExecutionA = stepExecutions.next();
-		assertEquals(stepExecutionA.getStepName(), "stepA");
-		StepExecution stepExecutionC = stepExecutions.next();
-		assertEquals(stepExecutionC.getStepName(), "stepC");
+		assertEquals(stepExecutions.next().getStepName(), "stepA");
+		assertEquals(stepExecutions.next().getStepName(), "stepC");
+		assertFalse(stepExecutions.hasNext());
+	}
+
+	private static StepSupport createCompleteStep(String name) {
+		return new StepSupport(name) {
+			@Override
+			public void execute(StepExecution stepExecution)
+					throws JobInterruptedException, UnexpectedJobExecutionException {
+				stepExecution.upgradeStatus(BatchStatus.COMPLETED);
+				stepExecution.setExitStatus(ExitStatus.COMPLETED);
+			}
+		};
 	}
 
 }
