This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/dev/pr-12195-bd09f6d9ae97de7230a6659962bb06f0b8eef10e in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit ae9132fe0767eee612259c293dc28b73a48dfae4 Author: Daniel <[email protected]> AuthorDate: Mon Sep 14 14:37:01 2026 +0000 [Fix][Zeta] Checkpoint restore remap duplicates subtask state for parallelism > 1 (#12195) Co-authored-by: Claude Fable 5.1 <[email protected]> --- .../server/checkpoint/CheckpointCoordinator.java | 49 ++++- .../checkpoint/CheckpointCoordinatorTest.java | 198 +++++++++++++++++++++ 2 files changed, 246 insertions(+), 1 deletion(-) diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinator.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinator.java index 0fd4cf039d..d4dfa45436 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinator.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinator.java @@ -112,6 +112,18 @@ public class CheckpointCoordinator { */ private final Map<Long, Integer> pipelineTasks; + /** + * Current parallelism of every action in this pipeline, keyed by {@link ActionStateKey}. + * + * <p>Used exclusively by the checkpoint-state remap in {@link #restoreTaskState(TaskLocation)}. + * It is intentionally kept separate from {@link #pipelineTasks}: that map is keyed by {@link + * TaskLocation#getTaskVertexId()}, which encodes the subtask's own parallelism index and is + * therefore unique per subtask (see the {@link TaskLocation} constructor), so grouping by it + * can never recover how many subtasks currently run a given action. See {@link + * #getActionParallelism(Map)} for how this is derived instead. + */ + private final Map<ActionStateKey, Integer> actionParallelism; + private final Map<Long, SeaTunnelTaskState> pipelineTaskStatus; private final CheckpointPlan plan; @@ -238,6 +250,7 @@ public class CheckpointCoordinator { this.scheduler = MDCTracer.tracing(scheduler); this.serializer = new ProtoStuffSerializer(); this.pipelineTasks = getPipelineTasks(plan.getPipelineSubtasks()); + this.actionParallelism = getActionParallelism(plan.getSubtaskActions()); this.pipelineTaskStatus = new ConcurrentHashMap<>(); this.checkpointIdCounter = checkpointIdCounter; this.readyToCloseStartingTask = new CopyOnWriteArraySet<>(); @@ -378,7 +391,6 @@ public class CheckpointCoordinator { if (!latestCompletedCheckpoint.isRestored()) { latestCompletedCheckpoint.setRestored(true); } - final Integer currentParallelism = pipelineTasks.get(taskLocation.getTaskVertexId()); plan.getSubtaskActions() .get(taskLocation) .forEach( @@ -396,6 +408,12 @@ public class CheckpointCoordinator { states.add(actionState.getCoordinatorState()); return; } + // The remap step must be the CURRENT parallelism of this + // specific action, not of the task/vertex: a single subtask can + // carry several chained actions, and pipelineTasks is keyed by + // TaskLocation#getTaskVertexId(), which is unique per subtask (see + // actionParallelism's javadoc), so it can never be reused here. + final int currentParallelism = actionParallelism.get(tuple.f0()); for (int i = tuple.f1(); i < actionState.getParallelism(); i += currentParallelism) { @@ -829,6 +847,35 @@ public class CheckpointCoordinator { .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().size())); } + /** + * Derives the current parallelism of every action from the subtask-to-action mapping of the + * checkpoint plan. + * + * <p>{@code + * org.apache.seatunnel.engine.server.dag.physical.PhysicalPlanGenerator#fillCheckpointPlan} + * (and the enumerator/committer task wiring around it) records, for every subtask, one {@code + * (ActionStateKey, index)} tuple per action that subtask participates in, where {@code index} + * is that subtask's own parallelism index (or {@link CheckpointPlan#COORDINATOR_INDEX} for + * coordinator-only tasks such as the split enumerator). Counting, per action, how many tuples + * carry a real (non-coordinator) index therefore yields exactly the number of subtasks + * currently running that action, i.e. its current parallelism. + * + * <p>This must NOT be approximated via {@link #getPipelineTasks(Set)}: that map is keyed by + * {@link TaskLocation#getTaskVertexId()}, which folds in the subtask's own parallelism index + * and is therefore unique per subtask (see the {@link TaskLocation} constructor), so grouping + * by it always yields a group size of 1 regardless of the action's real parallelism. + * + * @param subtaskActions the subtask-to-action mapping of the checkpoint plan + * @return the current parallelism of every action key found in {@code subtaskActions} + */ + public static Map<ActionStateKey, Integer> getActionParallelism( + Map<TaskLocation, Set<Tuple2<ActionStateKey, Integer>>> subtaskActions) { + return subtaskActions.values().stream() + .flatMap(Set::stream) + .filter(tuple -> !COORDINATOR_INDEX.equals(tuple.f1())) + .collect(Collectors.groupingBy(Tuple2::f0, Collectors.summingInt(tuple -> 1))); + } + @SneakyThrows public PassiveCompletableFuture<CompletedCheckpoint> startSavepoint() { LOG.info("start save point for job id: {}.", jobId); diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinatorTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinatorTest.java index cf33840e91..b6a2e396f9 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinatorTest.java +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinatorTest.java @@ -18,6 +18,7 @@ package org.apache.seatunnel.engine.server.checkpoint; import org.apache.seatunnel.common.utils.ReflectionUtils; +import org.apache.seatunnel.engine.checkpoint.storage.PipelineState; import org.apache.seatunnel.engine.checkpoint.storage.api.CheckpointStorage; import org.apache.seatunnel.engine.common.config.server.CheckpointConfig; import org.apache.seatunnel.engine.common.config.server.CheckpointStorageConfig; @@ -25,8 +26,10 @@ import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture; import org.apache.seatunnel.engine.core.checkpoint.CheckpointIDCounter; import org.apache.seatunnel.engine.core.checkpoint.CheckpointType; import org.apache.seatunnel.engine.core.job.RestoreMode; +import org.apache.seatunnel.engine.serializer.protobuf.ProtoStuffSerializer; import org.apache.seatunnel.engine.server.AbstractSeaTunnelServerTest; import org.apache.seatunnel.engine.server.checkpoint.monitor.CheckpointMonitorService; +import org.apache.seatunnel.engine.server.checkpoint.operation.NotifyTaskRestoreOperation; import org.apache.seatunnel.engine.server.checkpoint.operation.TaskAcknowledgeOperation; import org.apache.seatunnel.engine.server.common.SeaTunnelEngineContext; import org.apache.seatunnel.engine.server.execution.TaskGroupLocation; @@ -37,6 +40,7 @@ import org.apache.seatunnel.engine.server.task.statemachine.SeaTunnelTaskState; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.MockedStatic; import org.mockito.Mockito; @@ -45,6 +49,7 @@ import com.hazelcast.map.IMap; import com.hazelcast.spi.impl.NodeEngine; import com.hazelcast.spi.impl.operationservice.impl.InvocationFuture; +import java.nio.charset.StandardCharsets; import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; @@ -65,6 +70,7 @@ import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; import static org.apache.seatunnel.engine.common.Constant.IMAP_RUNNING_JOB_STATE; @@ -1176,6 +1182,198 @@ public class CheckpointCoordinatorTest executorService.shutdownNow(); } } + + /** + * Regression for the per-subtask state remap in {@code CheckpointCoordinator#restoreTaskState}: + * for every old/new parallelism pair in 1..4 x 1..4, every subtask state recorded in the + * checkpoint must be handed to exactly one subtask of the restored plan (no duplicate, no + * drop), a subtask must only receive checkpointed indexes congruent to its own index modulo the + * new parallelism, and the coordinator task must receive exactly the coordinator state. + * + * <p>Before the fix the remap step was looked up through {@code + * TaskLocation#getTaskVertexId()}, which is unique per subtask, so the step degenerated to 1 + * and subtask j received every checkpointed state with index >= j, duplicating splits after any + * restore with parallelism > 1. + */ + @Test + void testRestoreTaskStateDeliversEveryCheckpointedSubtaskStateExactlyOnce() { + ExecutorService executorService = Executors.newCachedThreadPool(); + try { + for (int oldParallelism = 1; oldParallelism <= 4; oldParallelism++) { + for (int newParallelism = 1; newParallelism <= 4; newParallelism++) { + assertRestoreRemap(executorService, oldParallelism, newParallelism); + } + } + } finally { + executorService.shutdownNow(); + } + } + + /** + * Restores a checkpoint taken at {@code oldParallelism} into a plan with {@code newParallelism} + * subtasks through the real {@code restoreTaskState} path and asserts the remap contract on the + * {@link NotifyTaskRestoreOperation}s sent to the member nodes. + */ + private void assertRestoreRemap( + ExecutorService executorService, int oldParallelism, int newParallelism) { + String scenario = "oldParallelism=" + oldParallelism + ", newParallelism=" + newParallelism; + ActionStateKey actionKey = new ActionStateKey("ActionStateKey - remap-source"); + + // Checkpoint taken at the old parallelism: one state per old subtask plus the coordinator + // (split enumerator) state, indexed exactly as PendingCheckpoint#acknowledgeTask records + // them (subtask index -> subtaskStates, COORDINATOR_INDEX -> coordinatorState). + ActionState actionState = new ActionState(actionKey, oldParallelism); + actionState.reportState( + CheckpointPlan.COORDINATOR_INDEX, + new ActionSubtaskState( + actionKey, + CheckpointPlan.COORDINATOR_INDEX, + Collections.singletonList("coordinator".getBytes(StandardCharsets.UTF_8)))); + for (int index = 0; index < oldParallelism; index++) { + actionState.reportState( + index, + new ActionSubtaskState( + actionKey, + index, + Collections.singletonList( + String.valueOf(index).getBytes(StandardCharsets.UTF_8)))); + } + Map<ActionStateKey, ActionState> taskStates = new HashMap<>(); + taskStates.put(actionKey, actionState); + long now = System.currentTimeMillis(); + CompletedCheckpoint completedCheckpoint = + new CompletedCheckpoint( + 1L, + 1, + 1L, + now, + CheckpointType.SAVEPOINT_TYPE, + now, + taskStates, + new HashMap<>()); + PipelineState pipelineState = + PipelineState.builder() + .jobId("1") + .pipelineId(1) + .checkpointId(1L) + .states(new ProtoStuffSerializer().serialize(completedCheckpoint)) + .build(); + + // Plan of the restored job at the new parallelism, shaped like the output of + // PhysicalPlanGenerator: the coordinator task is registered with COORDINATOR_INDEX and + // every parallelism index gets its own task group (hence its own task id) registered + // with that index. + TaskLocation coordinatorTask = new TaskLocation(new TaskGroupLocation(1L, 1, 1), 0, 0); + Map<TaskLocation, Set<Tuple2<ActionStateKey, Integer>>> subtaskActions = new HashMap<>(); + subtaskActions.put( + coordinatorTask, + Collections.singleton(Tuple2.tuple2(actionKey, CheckpointPlan.COORDINATOR_INDEX))); + List<TaskLocation> subtasks = new ArrayList<>(); + for (int index = 0; index < newParallelism; index++) { + TaskLocation subtask = + new TaskLocation(new TaskGroupLocation(1L, 1, 2 + index), 0, index); + subtasks.add(subtask); + subtaskActions.put(subtask, Collections.singleton(Tuple2.tuple2(actionKey, index))); + } + Map<ActionStateKey, Integer> pipelineActions = new HashMap<>(); + pipelineActions.put(actionKey, newParallelism); + Set<TaskLocation> pipelineSubtasks = new HashSet<>(subtasks); + pipelineSubtasks.add(coordinatorTask); + CheckpointPlan plan = + CheckpointPlan.builder() + .pipelineId(1) + .pipelineSubtasks(pipelineSubtasks) + .startingSubtasks(Collections.singleton(coordinatorTask)) + .pipelineActions(pipelineActions) + .subtaskActions(subtaskActions) + .build(); + + int derivedParallelism = + CheckpointCoordinator.getActionParallelism(plan.getSubtaskActions()).get(actionKey); + Assertions.assertEquals( + newParallelism, + derivedParallelism, + scenario + ": the remap step must be the per-action parallelism"); + + CheckpointConfig checkpointConfig = new CheckpointConfig(); + checkpointConfig.setStorage(new CheckpointStorageConfig()); + CheckpointManager mockManager = Mockito.mock(CheckpointManager.class); + Mockito.doReturn(Mockito.mock(InvocationFuture.class)) + .when(mockManager) + .sendOperationToMemberNode(Mockito.any(TaskOperation.class)); + @SuppressWarnings("unchecked") + IMap<Object, Object> mockIMap = Mockito.mock(IMap.class); + CheckpointCoordinator coordinator = + new CheckpointCoordinator( + mockManager, + Mockito.mock(CheckpointStorage.class), + checkpointConfig, + 1L, + plan, + Mockito.mock(CheckpointIDCounter.class), + pipelineState, + executorService, + mockIMap, + true, + null); + + ReflectionUtils.invoke(coordinator, "restoreTaskState", coordinatorTask); + for (TaskLocation subtask : subtasks) { + ReflectionUtils.invoke(coordinator, "restoreTaskState", subtask); + } + + ArgumentCaptor<TaskOperation> captor = ArgumentCaptor.forClass(TaskOperation.class); + Mockito.verify(mockManager, Mockito.times(newParallelism + 1)) + .sendOperationToMemberNode(captor.capture()); + Map<TaskLocation, List<Integer>> restoredIndexes = new HashMap<>(); + for (TaskOperation operation : captor.getAllValues()) { + Assertions.assertInstanceOf(NotifyTaskRestoreOperation.class, operation, scenario); + @SuppressWarnings("unchecked") + List<ActionSubtaskState> restored = + (List<ActionSubtaskState>) + ReflectionUtils.getField(operation, "restoredState") + .orElseThrow( + () -> + new IllegalStateException( + "restoredState field not found")); + restoredIndexes.put( + operation.getTaskLocation(), + restored.stream() + .map(ActionSubtaskState::getIndex) + .collect(Collectors.toList())); + } + + Assertions.assertEquals( + Collections.singletonList(CheckpointPlan.COORDINATOR_INDEX), + restoredIndexes.get(coordinatorTask), + scenario + ": the coordinator task must receive exactly the coordinator state"); + + List<Integer> delivered = new ArrayList<>(); + for (TaskLocation subtask : subtasks) { + List<Integer> indexes = restoredIndexes.get(subtask); + Assertions.assertNotNull(indexes, scenario + ": subtask " + subtask.getTaskIndex()); + for (Integer index : indexes) { + Assertions.assertEquals( + subtask.getTaskIndex(), + index % newParallelism, + scenario + + ": subtask " + + subtask.getTaskIndex() + + " received checkpointed index " + + index); + } + delivered.addAll(indexes); + } + Collections.sort(delivered); + List<Integer> expected = new ArrayList<>(); + for (int index = 0; index < oldParallelism; index++) { + expected.add(index); + } + Assertions.assertEquals( + expected, + delivered, + scenario + ": every checkpointed subtask state must be restored exactly once"); + } } class TestCheckpointManager extends CheckpointManager {
