diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/RemoteStepExecutionAggregator.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/RemoteStepExecutionAggregator.java index 88e0afd42..49610dddb 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/RemoteStepExecutionAggregator.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/RemoteStepExecutionAggregator.java @@ -1,5 +1,5 @@ /* - * Copyright 2006-2013 the original author or authors. + * Copyright 2006-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. @@ -16,9 +16,12 @@ package org.springframework.batch.core.partition.support; -import java.util.ArrayList; import java.util.Collection; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; +import org.springframework.batch.core.JobExecution; import org.springframework.batch.core.StepExecution; import org.springframework.batch.core.explore.JobExplorer; import org.springframework.beans.factory.InitializingBean; @@ -88,14 +91,16 @@ public class RemoteStepExecutionAggregator implements StepExecutionAggregator, I if (executions == null) { return; } - Collection updates = new ArrayList<>(); - for (StepExecution stepExecution : executions) { + Set stepExecutionIds = executions.stream().map(stepExecution -> { Long id = stepExecution.getId(); Assert.state(id != null, "StepExecution has null id. It must be saved first: " + stepExecution); - StepExecution update = jobExplorer.getStepExecution(stepExecution.getJobExecutionId(), id); - Assert.state(update != null, "Could not reload StepExecution from JobRepository: " + stepExecution); - updates.add(update); - } + return id; + }).collect(Collectors.toSet()); + JobExecution jobExecution = jobExplorer.getJobExecution(result.getJobExecutionId()); + Assert.state(jobExecution != null, + "Could not load JobExecution from JobRepository for id " + result.getJobExecutionId()); + List updates = jobExecution.getStepExecutions().stream() + .filter(stepExecution -> stepExecutionIds.contains(stepExecution.getId())).collect(Collectors.toList()); delegate.aggregate(result, updates); } diff --git a/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java b/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java index 334634684..2ac196a49 100644 --- a/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java +++ b/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2009-2022 the original author or authors. + * Copyright 2009-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. @@ -16,7 +16,6 @@ package org.springframework.batch.integration.partition; import java.util.HashSet; -import java.util.Iterator; import java.util.List; import java.util.Set; import java.util.concurrent.Callable; @@ -29,6 +28,7 @@ import javax.sql.DataSource; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; +import org.springframework.batch.core.JobExecution; import org.springframework.batch.core.Step; import org.springframework.batch.core.StepExecution; import org.springframework.batch.core.explore.JobExplorer; @@ -256,20 +256,13 @@ public class MessageChannelPartitionHandler extends AbstractPartitionHandler imp Callable> callback = new Callable>() { @Override public Set call() throws Exception { - - for (Iterator stepExecutionIterator = split.iterator(); stepExecutionIterator - .hasNext();) { - StepExecution curStepExecution = stepExecutionIterator.next(); - - if (!result.contains(curStepExecution)) { - StepExecution partitionStepExecution = jobExplorer - .getStepExecution(managerStepExecution.getJobExecutionId(), curStepExecution.getId()); - - if (!partitionStepExecution.getStatus().isRunning()) { - result.add(partitionStepExecution); - } - } - } + Set currentStepExecutionIds = split.stream().map(StepExecution::getId) + .collect(Collectors.toSet()); + JobExecution jobExecution = jobExplorer.getJobExecution(managerStepExecution.getJobExecutionId()); + jobExecution.getStepExecutions().stream() + .filter(stepExecution -> currentStepExecutionIds.contains(stepExecution.getId())) + .filter(stepExecution -> !result.contains(stepExecution)) + .filter(stepExecution -> !stepExecution.getStatus().isRunning()).forEach(result::add); if (logger.isDebugEnabled()) { logger.debug(String.format("Currently waiting on %s partitions to finish", split.size())); diff --git a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java index 2720b9295..97200b2d7 100644 --- a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java +++ b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2020-2022 the original author or authors. + * Copyright 2020-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. @@ -16,6 +16,7 @@ package org.springframework.batch.integration.partition; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.HashSet; @@ -174,8 +175,12 @@ class MessageChannelPartitionHandlerTests { stepExecutions.add(partition2); stepExecutions.add(partition3); when(stepExecutionSplitter.split(any(StepExecution.class), eq(1))).thenReturn(stepExecutions); - when(jobExplorer.getStepExecution(eq(5L), any(Long.class))).thenReturn(partition2, partition1, partition3, - partition3, partition3, partition3, partition4); + JobExecution runningJobExecution = new JobExecution(5L, new JobParameters()); + runningJobExecution.addStepExecutions(Arrays.asList(partition2, partition1, partition3)); + JobExecution completedJobExecution = new JobExecution(5L, new JobParameters()); + completedJobExecution.addStepExecutions(Arrays.asList(partition2, partition1, partition4)); + when(jobExplorer.getJobExecution(5L)).thenReturn(runningJobExecution, runningJobExecution, runningJobExecution, + completedJobExecution); // set messageChannelPartitionHandler.setMessagingOperations(operations); @@ -220,7 +225,9 @@ class MessageChannelPartitionHandlerTests { stepExecutions.add(partition2); stepExecutions.add(partition3); when(stepExecutionSplitter.split(any(StepExecution.class), eq(1))).thenReturn(stepExecutions); - when(jobExplorer.getStepExecution(eq(5L), any(Long.class))).thenReturn(partition2, partition1, partition3); + JobExecution runningJobExecution = new JobExecution(5L, new JobParameters()); + runningJobExecution.addStepExecutions(Arrays.asList(partition2, partition1, partition3)); + when(jobExplorer.getJobExecution(5L)).thenReturn(runningJobExecution); // set messageChannelPartitionHandler.setMessagingOperations(operations);