diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/AbstractPartitionHandler.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/AbstractPartitionHandler.java index bf507a510..0d4538511 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/AbstractPartitionHandler.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/AbstractPartitionHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2006-2021 the original author or authors. + * Copyright 2006-2022 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. @@ -35,7 +35,7 @@ import org.springframework.batch.core.partition.StepExecutionSplitter; */ public abstract class AbstractPartitionHandler implements PartitionHandler { - private int gridSize = 1; + protected int gridSize = 1; /** * Executes the specified {@link StepExecution} instances and returns an updated view 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 2b9bccfa9..23a905388 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-2021 the original author or authors. + * Copyright 2009-2022 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. @@ -15,14 +15,14 @@ */ package org.springframework.batch.integration.partition; -import java.util.ArrayList; -import java.util.Collection; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Set; import java.util.concurrent.Callable; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; import javax.sql.DataSource; @@ -35,6 +35,7 @@ import org.springframework.batch.core.explore.JobExplorer; import org.springframework.batch.core.explore.support.JobExplorerFactoryBean; import org.springframework.batch.core.partition.PartitionHandler; import org.springframework.batch.core.partition.StepExecutionSplitter; +import org.springframework.batch.core.partition.support.AbstractPartitionHandler; import org.springframework.batch.core.repository.JobRepository; import org.springframework.batch.poller.DirectPoller; import org.springframework.batch.poller.Poller; @@ -85,12 +86,10 @@ import org.springframework.util.CollectionUtils; * */ @MessageEndpoint -public class MessageChannelPartitionHandler implements PartitionHandler, InitializingBean { +public class MessageChannelPartitionHandler extends AbstractPartitionHandler implements InitializingBean { private static Log logger = LogFactory.getLog(MessageChannelPartitionHandler.class); - private int gridSize = 1; - private MessagingTemplate messagingGateway; private String stepName; @@ -187,18 +186,6 @@ public class MessageChannelPartitionHandler implements PartitionHandler, Initial this.messagingGateway = messagingGateway; } - /** - * Passed to the {@link StepExecutionSplitter} in the - * {@link #handle(StepExecutionSplitter, StepExecution)} method, instructing it how - * many {@link StepExecution} instances are required, ideally. The - * {@link StepExecutionSplitter} is allowed to ignore the grid size in the case of a - * restart, since the input data partitions must be preserved. - * @param gridSize the number of step executions that will be created - */ - public void setGridSize(int gridSize) { - this.gridSize = gridSize; - } - /** * The name of the {@link Step} that will be used to execute the partitioned * {@link StepExecution}. This is a regular Spring Batch step, with all the business @@ -234,19 +221,17 @@ public class MessageChannelPartitionHandler implements PartitionHandler, Initial * * @see PartitionHandler#handle(StepExecutionSplitter, StepExecution) */ - public Collection handle(StepExecutionSplitter stepExecutionSplitter, - final StepExecution managerStepExecution) throws Exception { + @Override + protected Set doHandle(StepExecution managerStepExecution, Set partitionStepExecutions) throws Exception { - final Set split = stepExecutionSplitter.split(managerStepExecution, gridSize); - - if (CollectionUtils.isEmpty(split)) { - return split; + if (CollectionUtils.isEmpty(partitionStepExecutions)) { + return partitionStepExecutions; } int count = 0; - for (StepExecution stepExecution : split) { - Message request = createMessage(count++, split.size(), + for (StepExecution stepExecution : partitionStepExecutions) { + Message request = createMessage(count++, partitionStepExecutions.size(), new StepExecutionRequest(stepName, stepExecution.getJobExecutionId(), stepExecution.getId()), replyChannel); if (logger.isDebugEnabled()) { @@ -259,17 +244,17 @@ public class MessageChannelPartitionHandler implements PartitionHandler, Initial return receiveReplies(replyChannel); } else { - return pollReplies(managerStepExecution, split); + return pollReplies(managerStepExecution, partitionStepExecutions); } } - private Collection pollReplies(final StepExecution managerStepExecution, + private Set pollReplies(final StepExecution managerStepExecution, final Set split) throws Exception { - final Collection result = new ArrayList<>(split.size()); + final Set result = new HashSet<>(split.size()); - Callable> callback = new Callable>() { + Callable> callback = new Callable>() { @Override - public Collection call() throws Exception { + public Set call() throws Exception { for (Iterator stepExecutionIterator = split.iterator(); stepExecutionIterator .hasNext();) { @@ -298,8 +283,8 @@ public class MessageChannelPartitionHandler implements PartitionHandler, Initial } }; - Poller> poller = new DirectPoller<>(pollInterval); - Future> resultsFuture = poller.poll(callback); + Poller> poller = new DirectPoller<>(pollInterval); + Future> resultsFuture = poller.poll(callback); if (timeout >= 0) { return resultsFuture.get(timeout, TimeUnit.MILLISECONDS); @@ -309,9 +294,8 @@ public class MessageChannelPartitionHandler implements PartitionHandler, Initial } } - private Collection receiveReplies(PollableChannel currentReplyChannel) { - @SuppressWarnings("unchecked") - Message> message = (Message>) messagingGateway + private Set receiveReplies(PollableChannel currentReplyChannel) { + Message> message = (Message>) messagingGateway .receive(currentReplyChannel); if (message == null) { @@ -321,7 +305,7 @@ public class MessageChannelPartitionHandler implements PartitionHandler, Initial logger.debug("Received replies: " + message); } - return message.getPayload(); + return new HashSet<>(message.getPayload()); } private Message createMessage(int sequenceNumber, int sequenceSize, 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 cdeb91367..2720b9295 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 @@ -84,7 +84,7 @@ class MessageChannelPartitionHandlerTests { HashSet stepExecutions = new HashSet<>(); stepExecutions.add(new StepExecution("step1", new JobExecution(5L))); when(stepExecutionSplitter.split(any(StepExecution.class), eq(1))).thenReturn(stepExecutions); - when(message.getPayload()).thenReturn(Collections.emptyList()); + when(message.getPayload()).thenReturn(Collections.emptySet()); when(operations.receive((PollableChannel) any())).thenReturn(message); // set messageChannelPartitionHandler.setMessagingOperations(operations); @@ -112,7 +112,7 @@ class MessageChannelPartitionHandlerTests { HashSet stepExecutions = new HashSet<>(); stepExecutions.add(new StepExecution("step1", new JobExecution(5L))); when(stepExecutionSplitter.split(any(StepExecution.class), eq(1))).thenReturn(stepExecutions); - when(message.getPayload()).thenReturn(Collections.emptyList()); + when(message.getPayload()).thenReturn(Collections.emptySet()); when(operations.receive(replyChannel)).thenReturn(message); // set messageChannelPartitionHandler.setMessagingOperations(operations);