Skip to content

Commit cc2d57f

Browse files
quaffquaff
authored
Fix step execution context is not persisted and restored (#5147)
1. Step execution context is not persisted in `SimpleStepExecutionSplitter::split` 2. Step execution context is not restored in `SimpleJobRepository::getStepExecution` Closes GH-5138 Signed-off-by: Yanming Zhou <[email protected]>
1 parent dcfd42e commit cc2d57f

File tree

5 files changed

+54
-21
lines changed

5 files changed

+54
-21
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
*
4646
* @author Dave Syer
4747
* @author Mahmoud Ben Hassine
48+
* @author Yanming Zhou
4849
* @since 2.0
4950
*/
5051
public class SimpleStepExecutionSplitter implements StepExecutionSplitter {
@@ -138,13 +139,15 @@ public Set<StepExecution> split(StepExecution stepExecution, int gridSize) throw
138139
if (lastStepExecution == null) { // fresh start
139140
StepExecution currentStepExecution = jobRepository.createStepExecution(stepName, jobExecution);
140141
currentStepExecution.setExecutionContext(context.getValue());
142+
jobRepository.updateExecutionContext(currentStepExecution);
141143
set.add(currentStepExecution);
142144
}
143145
else { // restart
144146
if (lastStepExecution.getStatus() != BatchStatus.COMPLETED
145147
&& shouldStart(allowStartIfComplete, stepExecution, lastStepExecution)) {
146148
StepExecution currentStepExecution = jobRepository.createStepExecution(stepName, jobExecution);
147149
currentStepExecution.setExecutionContext(lastStepExecution.getExecutionContext());
150+
jobRepository.updateExecutionContext(currentStepExecution);
148151
set.add(currentStepExecution);
149152
}
150153
}

‎spring-batch-core/src/main/java/org/springframework/batch/core/repository/explore/support/SimpleJobExplorer.java‎

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
* @author Mahmoud Ben Hassine
4545
* @author Parikshit Dutta
4646
* @author Glenn Renfro
47+
* @author Yanming Zhou
4748
* @see JobExplorer
4849
* @see JobInstanceDao
4950
* @see JobExecutionDao
@@ -157,9 +158,9 @@ public long getJobInstanceCount(String jobName) throws NoSuchJobException {
157158
public List<JobExecution> getJobExecutions(JobInstance jobInstance) {
158159
List<JobExecution> executions = jobExecutionDao.findJobExecutions(jobInstance);
159160
for (JobExecution jobExecution : executions) {
160-
getJobExecutionDependencies(jobExecution);
161+
fillJobExecutionDependencies(jobExecution);
161162
for (StepExecution stepExecution : jobExecution.getStepExecutions()) {
162-
getStepExecutionDependencies(stepExecution);
163+
fillStepExecutionDependencies(stepExecution);
163164
}
164165
}
165166
return executions;
@@ -170,9 +171,9 @@ public List<JobExecution> getJobExecutions(JobInstance jobInstance) {
170171
public JobExecution getLastJobExecution(JobInstance jobInstance) {
171172
JobExecution lastJobExecution = jobExecutionDao.getLastJobExecution(jobInstance);
172173
if (lastJobExecution != null) {
173-
getJobExecutionDependencies(lastJobExecution);
174+
fillJobExecutionDependencies(lastJobExecution);
174175
for (StepExecution stepExecution : lastJobExecution.getStepExecutions()) {
175-
getStepExecutionDependencies(stepExecution);
176+
fillStepExecutionDependencies(stepExecution);
176177
}
177178
}
178179
return lastJobExecution;
@@ -198,7 +199,7 @@ public List<JobExecution> findJobExecutions(JobInstance jobInstance) {
198199
JobExecution jobExecution = jobExecutionDao.getLastJobExecution(jobInstance);
199200

200201
if (jobExecution != null) {
201-
getJobExecutionDependencies(jobExecution);
202+
fillJobExecutionDependencies(jobExecution);
202203
}
203204
return jobExecution;
204205
}
@@ -207,9 +208,9 @@ public List<JobExecution> findJobExecutions(JobInstance jobInstance) {
207208
public Set<JobExecution> findRunningJobExecutions(@Nullable String jobName) {
208209
Set<JobExecution> executions = jobExecutionDao.findRunningJobExecutions(jobName);
209210
for (JobExecution jobExecution : executions) {
210-
getJobExecutionDependencies(jobExecution);
211+
fillJobExecutionDependencies(jobExecution);
211212
for (StepExecution stepExecution : jobExecution.getStepExecutions()) {
212-
getStepExecutionDependencies(stepExecution);
213+
fillStepExecutionDependencies(stepExecution);
213214
}
214215
}
215216
return executions;
@@ -222,20 +223,18 @@ public JobExecution getJobExecution(long executionId) {
222223
if (jobExecution == null) {
223224
return null;
224225
}
225-
getJobExecutionDependencies(jobExecution);
226+
fillJobExecutionDependencies(jobExecution);
226227
for (StepExecution stepExecution : jobExecution.getStepExecutions()) {
227-
getStepExecutionDependencies(stepExecution);
228+
fillStepExecutionDependencies(stepExecution);
228229
}
229230
return jobExecution;
230231
}
231232

232233
/*
233-
* Find all dependencies for a JobExecution, including JobInstance (which requires
234+
* Fill all dependencies for a JobExecution, including JobInstance (which requires
234235
* JobParameters) plus StepExecutions
235236
*/
236-
// TODO rename to something more representative of what it does (side effect on the
237-
// parameter)
238-
private void getJobExecutionDependencies(JobExecution jobExecution) {
237+
protected void fillJobExecutionDependencies(JobExecution jobExecution) {
239238
JobInstance jobInstance = jobInstanceDao.getJobInstance(jobExecution);
240239
jobExecution.setJobInstance(jobInstance);
241240
jobExecution.addStepExecutions(stepExecutionDao.getStepExecutions(jobExecution));
@@ -257,9 +256,9 @@ public StepExecution getStepExecution(long jobExecutionId, long executionId) {
257256
if (jobExecution == null) {
258257
return null;
259258
}
260-
getJobExecutionDependencies(jobExecution);
259+
fillJobExecutionDependencies(jobExecution);
261260
StepExecution stepExecution = stepExecutionDao.getStepExecution(jobExecution, executionId);
262-
getStepExecutionDependencies(stepExecution);
261+
fillStepExecutionDependencies(stepExecution);
263262
return stepExecution;
264263
}
265264

@@ -268,8 +267,7 @@ public StepExecution getStepExecution(long jobExecutionId, long executionId) {
268267
StepExecution latest = stepExecutionDao.getLastStepExecution(jobInstance, stepName);
269268

270269
if (latest != null) {
271-
ExecutionContext stepExecutionContext = ecDao.getExecutionContext(latest);
272-
latest.setExecutionContext(stepExecutionContext);
270+
fillStepExecutionDependencies(latest);
273271
ExecutionContext jobExecutionContext = ecDao.getExecutionContext(latest.getJobExecution());
274272
latest.getJobExecution().setExecutionContext(jobExecutionContext);
275273
}
@@ -287,7 +285,7 @@ public long getStepExecutionCount(JobInstance jobInstance, String stepName) thro
287285
return stepExecutionDao.countStepExecutions(jobInstance, stepName);
288286
}
289287

290-
private void getStepExecutionDependencies(StepExecution stepExecution) {
288+
protected void fillStepExecutionDependencies(StepExecution stepExecution) {
291289
if (stepExecution != null) {
292290
stepExecution.setExecutionContext(ecDao.getExecutionContext(stepExecution));
293291
}

‎spring-batch-core/src/main/java/org/springframework/batch/core/repository/support/SimpleJobRepository.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
* @author Baris Cubukcuoglu
5353
* @author Parikshit Dutta
5454
* @author Mark John Moreno
55+
* @author Yanming Zhou
5556
* @see JobRepository
5657
* @see JobInstanceDao
5758
* @see JobExecutionDao
@@ -82,7 +83,13 @@ public List<JobInstance> findJobInstances(String jobName) {
8283
@Nullable
8384
@Override
8485
public StepExecution getStepExecution(long executionId) {
85-
return this.stepExecutionDao.getStepExecution(executionId);
86+
StepExecution stepExecution = this.stepExecutionDao.getStepExecution(executionId);
87+
if (stepExecution != null) {
88+
fillStepExecutionDependencies(stepExecution);
89+
ExecutionContext jobExecutionContext = this.ecDao.getExecutionContext(stepExecution.getJobExecution());
90+
stepExecution.getJobExecution().setExecutionContext(jobExecutionContext);
91+
}
92+
return stepExecution;
8693
}
8794

8895
/**

‎spring-batch-core/src/test/java/org/springframework/batch/core/partition/PartitionStepTests.java‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import java.time.LocalDateTime;
1919
import java.util.Arrays;
2020
import java.util.Collection;
21+
import java.util.HashMap;
22+
import java.util.Map;
2123
import java.util.Set;
2224
import java.util.concurrent.atomic.AtomicBoolean;
2325

@@ -41,10 +43,12 @@
4143
import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseBuilder;
4244

4345
import static org.junit.jupiter.api.Assertions.assertEquals;
46+
import static org.junit.jupiter.api.Assertions.assertNotNull;
4447

4548
/**
4649
* @author Dave Syer
4750
* @author Mahmoud Ben Hassine
51+
* @author Yanming Zhou
4852
*
4953
*/
5054
class PartitionStepTests {
@@ -71,12 +75,24 @@ void setUp() throws Exception {
7175
@Test
7276
void testVanillaStepExecution() throws Exception {
7377
SimpleStepExecutionSplitter stepExecutionSplitter = new SimpleStepExecutionSplitter(jobRepository,
74-
step.getName(), new SimplePartitioner());
78+
step.getName(), gridSize -> {
79+
Map<String, ExecutionContext> map = new HashMap<>(gridSize);
80+
for (int i = 0; i < gridSize; i++) {
81+
ExecutionContext context = new ExecutionContext();
82+
context.putString("foo", "foo" + i);
83+
map.put("partition" + i, context);
84+
}
85+
return map;
86+
});
7587
stepExecutionSplitter.setAllowStartIfComplete(true);
7688
step.setStepExecutionSplitter(stepExecutionSplitter);
7789
step.setPartitionHandler((stepSplitter, stepExecution) -> {
7890
Set<StepExecution> executions = stepSplitter.split(stepExecution, 2);
7991
for (StepExecution execution : executions) {
92+
// Query from repository to ensure it's persisted
93+
ExecutionContext context = jobRepository.getStepExecution(execution.getId()).getExecutionContext();
94+
assertNotNull(context.getString("foo"));
95+
8096
execution.setStatus(BatchStatus.COMPLETED);
8197
execution.setExitStatus(ExitStatus.COMPLETED);
8298
jobRepository.update(execution);
@@ -144,7 +160,9 @@ void testRestartStepExecution() throws Exception {
144160
else {
145161
for (StepExecution execution : executions) {
146162
// On restart the execution context should have been restored
147-
assertEquals(execution.getStepName(), execution.getExecutionContext().getString("foo"));
163+
// Query from repository to ensure it's persisted
164+
ExecutionContext context = jobRepository.getStepExecution(execution.getId()).getExecutionContext();
165+
assertEquals(execution.getStepName(), context.getString("foo"));
148166
}
149167
}
150168
for (StepExecution execution : executions) {

‎spring-batch-core/src/test/java/org/springframework/batch/core/repository/support/SimpleJobRepositoryIntegrationTests.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
* @author Robert Kasanicky
4949
* @author Dimitrios Liapis
5050
* @author Mahmoud Ben Hassine
51+
* @author Yanming Zhou
5152
*/
5253
// TODO rename to JdbcJobRepositoryIntegrationTests and update to new domain model
5354
// TODO should add a mongodb similar test suite
@@ -168,13 +169,19 @@ void testSaveExecutionContext() throws Exception {
168169
JobExecution jobExec = jobRepository.createJobExecution(jobInstance, jobParameters, new ExecutionContext());
169170
jobExec.setStartTime(LocalDateTime.now());
170171
jobExec.setExecutionContext(ctx);
172+
jobRepository.updateExecutionContext(jobExec);
171173
Step step = new StepSupport("step1");
172174
StepExecution stepExec = jobRepository.createStepExecution(step.getName(), jobExec);
173175
stepExec.setExecutionContext(ctx);
176+
jobRepository.updateExecutionContext(stepExec);
174177

175178
StepExecution retrievedStepExec = jobRepository.getLastStepExecution(jobExec.getJobInstance(), step.getName());
176179
assertEquals(stepExec, retrievedStepExec);
177180
assertEquals(ctx, retrievedStepExec.getExecutionContext());
181+
182+
retrievedStepExec = jobRepository.getStepExecution(stepExec.getId());
183+
assertEquals(stepExec, retrievedStepExec);
184+
assertEquals(ctx, retrievedStepExec.getExecutionContext());
178185
}
179186

180187
/*

0 commit comments

Comments
 (0)