[Improve][Zeta] Filter tasks and pipelines by state (#9926)

Author: dybyte

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinator.java

@@ -49,6 +49,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.hazelcast.jet.datamodel.Tuple2;
 import com.hazelcast.map.IMap;
 import com.hazelcast.spi.impl.operationservice.impl.InvocationFuture;
 import lombok.Getter;
@@ -58,6 +59,7 @@
 import java.time.Instant;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -787,26 +789,45 @@ private Set<Long> getNotYetAcknowledgedTasks() {
     }
 
     private Map<ActionStateKey, ActionState> getActionStates() {
-        // TODO: some tasks have completed and will not submit state again.
-        return plan.getPipelineActions().entrySet().stream()
+        Map<ActionStateKey, Integer> pipelineActions = new HashMap<>(plan.getPipelineActions());
+        Set<ActionStateKey> closedActionKeys =
+                plan.getSubtaskActions().entrySet().stream()
+                        .filter(
+                                entry ->
+                                        SeaTunnelTaskState.CLOSED.equals(
+                                                this.pipelineTaskStatus.get(
+                                                        entry.getKey().getTaskID())))
+                        .flatMap(entry -> entry.getValue().stream().map(Tuple2::f0))
+                        .collect(Collectors.toSet());
+        pipelineActions.keySet().removeAll(closedActionKeys);
+
+        return pipelineActions.entrySet().stream()
                 .collect(
                         Collectors.toMap(
                                 Map.Entry::getKey,
                                 entry -> new ActionState(entry.getKey(), entry.getValue())));
     }
 
     private Map<Long, TaskStatistics> getTaskStatistics() {
-        // TODO: some tasks have completed and don't need to be ack
-        return this.pipelineTasks.entrySet().stream()
+        Map<Long, Integer> tasks = new HashMap<>(this.pipelineTasks);
+        for (Long taskId : this.pipelineTasks.keySet()) {
+            if (SeaTunnelTaskState.CLOSED.equals(this.pipelineTaskStatus.get(taskId))) {
+                tasks.remove(taskId);
+            }
+        }
+        return tasks.entrySet().stream()
                 .collect(
                         Collectors.toMap(
                                 Map.Entry::getKey,
                                 entry -> new TaskStatistics(entry.getKey(), entry.getValue())));
     }
 
     public InvocationFuture<?>[] triggerCheckpoint(CheckpointBarrier checkpointBarrier) {
-        // TODO: some tasks have completed and don't need to trigger
         return plan.getStartingSubtasks().stream()
+                .filter(
+                        taskLocation ->
+                                !SeaTunnelTaskState.CLOSED.equals(
+                                        this.pipelineTaskStatus.get(taskLocation.getTaskID())))
                 .map(
                         taskLocation ->
                                 new CheckpointBarrierTriggerOperation(
@@ -1100,4 +1121,9 @@ public String getCheckpointStateImapKey() {
     public PendingCheckpoint getSavepointPendingCheckpoint() {
         return savepointPendingCheckpoint;
     }
+
+    @VisibleForTesting
+    public Map<Long, SeaTunnelTaskState> getPipelineTaskStatus() {
+        return pipelineTaskStatus;
+    }
 }

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalPlanGenerator.java

@@ -159,14 +159,24 @@ public Tuple2<PhysicalPlan, Map<Integer, CheckpointPlan>> generate() {
                                 .getJobConfig()
                                 .getEnvOptions()
                                 .get(EnvCommonOptions.NODE_TAG_FILTER.key());
-        // TODO Determine which tasks do not need to be restored according to state
         CopyOnWriteArrayList<PassiveCompletableFuture<PipelineStatus>>
                 waitForCompleteBySubPlanList = new CopyOnWriteArrayList<>();
 
+        List<Pipeline> unclosedPipelines = new ArrayList<>();
+        for (Pipeline pipeline : this.pipelines) {
+            PipelineLocation pipelineLocation =
+                    new PipelineLocation(jobImmutableInformation.getJobId(), pipeline.getId());
+            PipelineStatus pipelineStatus =
+                    (PipelineStatus) runningJobStateIMap.get(pipelineLocation);
+            if (!PipelineStatus.FINISHED.equals(pipelineStatus)) {
+                unclosedPipelines.add(pipeline);
+            }
+        }
+
         Map<Integer, CheckpointPlan> checkpointPlans = new HashMap<>();
-        final int totalPipelineNum = pipelines.size();
+        final int totalPipelineNum = unclosedPipelines.size();
         Stream<SubPlan> subPlanStream =
-                pipelines.stream()
+                unclosedPipelines.stream()
                         .map(
                                 pipeline -> {
                                     this.pipelineTasks.clear();

seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/checkpoint/CheckpointCoordinatorTest.java

@@ -18,7 +18,7 @@
 package org.apache.seatunnel.engine.server.checkpoint;
 
 import org.apache.seatunnel.common.utils.ReflectionUtils;
-import org.apache.seatunnel.engine.checkpoint.storage.exception.CheckpointStorageException;
+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;
 import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture;
@@ -28,19 +28,28 @@
 import org.apache.seatunnel.engine.server.execution.TaskGroupLocation;
 import org.apache.seatunnel.engine.server.execution.TaskLocation;
 import org.apache.seatunnel.engine.server.master.JobMaster;
+import org.apache.seatunnel.engine.server.task.operation.TaskOperation;
+import org.apache.seatunnel.engine.server.task.statemachine.SeaTunnelTaskState;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
+import com.hazelcast.jet.datamodel.Tuple2;
+import com.hazelcast.map.IMap;
+import com.hazelcast.spi.impl.NodeEngine;
 import com.hazelcast.spi.impl.operationservice.impl.InvocationFuture;
 
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -54,7 +63,7 @@ public class CheckpointCoordinatorTest
         extends AbstractSeaTunnelServerTest<CheckpointCoordinatorTest> {
 
     @Test
-    void testACKNotExistPendingCheckpoint() throws CheckpointStorageException {
+    void testACKNotExistPendingCheckpoint() {
         CheckpointConfig checkpointConfig = new CheckpointConfig();
         checkpointConfig.setStorage(new CheckpointStorageConfig());
         Map<Integer, CheckpointPlan> planMap = new HashMap<>();
@@ -80,8 +89,7 @@ void testACKNotExistPendingCheckpoint() throws CheckpointStorageException {
 
     @Test
     void testSchedulerThreadShouldNotBeInterruptedBeforeJobMasterCleaned()
-            throws CheckpointStorageException, ExecutionException, InterruptedException,
-                    TimeoutException {
+            throws ExecutionException, InterruptedException, TimeoutException {
         CheckpointConfig checkpointConfig = new CheckpointConfig();
         // quickly fail the checkpoint
         checkpointConfig.setCheckpointTimeout(5000);
@@ -122,8 +130,7 @@ protected void handleCheckpointError(int pipelineId, boolean neverRestore) {
 
     @Test
     void testCheckpointContinuesWorkAfterClockDrift()
-            throws CheckpointStorageException, ExecutionException, InterruptedException,
-                    TimeoutException {
+            throws ExecutionException, InterruptedException, TimeoutException {
         CheckpointConfig checkpointConfig = new CheckpointConfig();
         checkpointConfig.setStorage(new CheckpointStorageConfig());
         checkpointConfig.setCheckpointTimeout(5000);
@@ -169,9 +176,7 @@ protected void handleCheckpointError(int pipelineId, boolean neverRestore) {
     }
 
     @Test
-    void testCheckpointMinPause()
-            throws CheckpointStorageException, ExecutionException, InterruptedException,
-                    TimeoutException {
+    void testCheckpointMinPause() {
         CheckpointConfig checkpointConfig = new CheckpointConfig();
         checkpointConfig.setStorage(new CheckpointStorageConfig());
         checkpointConfig.setCheckpointInterval(10000); // 10 seconds
@@ -280,4 +285,113 @@ public CheckpointCoordinator getCheckpointCoordinator(int pipelineId) {
             executorService.shutdownNow();
         }
     }
+
+    @Test
+    void testFilteringClosedTasksAndActions() {
+        CheckpointConfig checkpointConfig = new CheckpointConfig();
+        checkpointConfig.setStorage(new CheckpointStorageConfig());
+        Map<Integer, CheckpointPlan> planMap = new HashMap<>();
+        planMap.put(1, CheckpointPlan.builder().pipelineId(1).build());
+        TestCheckpointManager checkpointManager =
+                new TestCheckpointManager(
+                        1L,
+                        nodeEngine,
+                        planMap,
+                        checkpointConfig,
+                        server.getCheckpointService().getCheckpointStorage(),
+                        instance.getExecutorService("test"),
+                        nodeEngine.getHazelcastInstance().getMap(IMAP_RUNNING_JOB_STATE));
+
+        TaskGroupLocation group1 = new TaskGroupLocation(1L, 1, 1);
+        TaskLocation task1 = new TaskLocation(group1, 1, 1);
+        TaskLocation task2 = new TaskLocation(group1, 2, 1);
+
+        ActionStateKey actionKey1 = new ActionStateKey("action1");
+        ActionStateKey actionKey2 = new ActionStateKey("action2");
+
+        Map<TaskLocation, Set<Tuple2<ActionStateKey, Integer>>> subtaskActions = new HashMap<>();
+        subtaskActions.put(task1, new HashSet<>(Arrays.asList(Tuple2.tuple2(actionKey1, 0))));
+        subtaskActions.put(task2, new HashSet<>(Arrays.asList(Tuple2.tuple2(actionKey2, 0))));
+
+        Map<ActionStateKey, Integer> pipelineActions = new HashMap<>();
+        pipelineActions.put(actionKey1, 1);
+        pipelineActions.put(actionKey2, 1);
+
+        CheckpointPlan plan =
+                CheckpointPlan.builder()
+                        .pipelineId(1)
+                        .pipelineSubtasks(new HashSet<>(Arrays.asList(task1, task2)))
+                        .startingSubtasks(new HashSet<>(Arrays.asList(task1, task2)))
+                        .subtaskActions(subtaskActions)
+                        .pipelineActions(pipelineActions)
+                        .build();
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        CheckpointCoordinator coordinator =
+                new CheckpointCoordinator(
+                        checkpointManager,
+                        null,
+                        checkpointConfig,
+                        1L,
+                        plan,
+                        null,
+                        null,
+                        executor,
+                        Mockito.mock(com.hazelcast.map.IMap.class),
+                        false);
+
+        Map<Long, SeaTunnelTaskState> taskStatus = coordinator.getPipelineTaskStatus();
+        taskStatus.put(task1.getTaskID(), SeaTunnelTaskState.RUNNING);
+        taskStatus.put(task2.getTaskID(), SeaTunnelTaskState.CLOSED);
+
+        Map<ActionStateKey, ActionState> actionStates =
+                (Map<ActionStateKey, ActionState>)
+                        ReflectionUtils.invoke(coordinator, "getActionStates");
+        Assertions.assertTrue(actionStates.containsKey(actionKey1));
+        Assertions.assertFalse(actionStates.containsKey(actionKey2));
+
+        Map<Long, TaskStatistics> stats =
+                (Map<Long, TaskStatistics>)
+                        ReflectionUtils.invoke(coordinator, "getTaskStatistics");
+        Assertions.assertTrue(stats.containsKey(task1.getTaskID()));
+        Assertions.assertFalse(stats.containsKey(task2.getTaskID()));
+
+        CheckpointBarrier barrier =
+                new CheckpointBarrier(
+                        1L, System.currentTimeMillis(), CheckpointType.CHECKPOINT_TYPE);
+        coordinator.triggerCheckpoint(barrier);
+        Assertions.assertEquals(1, checkpointManager.operations.size());
+
+        executor.shutdownNow();
+    }
+}
+
+class TestCheckpointManager extends CheckpointManager {
+    public List<TaskOperation> operations = new ArrayList<>();
+
+    public TestCheckpointManager(
+            long jobId,
+            NodeEngine nodeEngine,
+            Map<Integer, CheckpointPlan> checkpointPlanMap,
+            CheckpointConfig checkpointConfig,
+            CheckpointStorage checkpointStorage,
+            ExecutorService executorService,
+            IMap<Object, Object> runningJobStateIMap) {
+        super(
+                jobId,
+                false,
+                nodeEngine,
+                null,
+                checkpointPlanMap,
+                checkpointConfig,
+                checkpointStorage,
+                executorService,
+                runningJobStateIMap);
+    }
+
+    @Override
+    protected InvocationFuture<?> sendOperationToMemberNode(TaskOperation operation) {
+        this.operations.add(operation);
+        return null;
+    }
 }

seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/master/JobMasterTest.java

@@ -32,6 +32,7 @@
 import org.apache.seatunnel.engine.server.dag.physical.PhysicalVertex;
 import org.apache.seatunnel.engine.server.dag.physical.PipelineLocation;
 import org.apache.seatunnel.engine.server.dag.physical.SubPlan;
+import org.apache.seatunnel.engine.server.dag.physical.UnknownPhysicalPlanException;
 import org.apache.seatunnel.engine.server.execution.TaskGroupLocation;
 import org.apache.seatunnel.engine.server.resourcemanager.resource.SlotProfile;
 import org.apache.seatunnel.engine.server.service.slot.SlotService;
@@ -248,6 +249,19 @@ public void testCloseIdleTask() throws InterruptedException {
         assertCloseIdleTask(jobMaster);
     }
 
+    @Test
+    void testFilteringFinishedPipelinesInPhysicalPlanGenerator() throws Exception {
+        long jobId = instance.getFlakeIdGenerator(Constant.SEATUNNEL_ID_GENERATOR_NAME).newId();
+        JobMaster jobMaster = newJobInstanceWithRunningState(jobId);
+
+        jobMaster
+                .getRunningJobStateIMap()
+                .put(new PipelineLocation(jobId, 1), PipelineStatus.FINISHED);
+        Assertions.assertThrows(
+                UnknownPhysicalPlanException.class,
+                () -> jobMaster.init(System.currentTimeMillis(), false));
+    }
+
     private void assertCloseIdleTask(JobMaster jobMaster) {
         SlotService slotService = server.getSlotService();
         Assertions.assertEquals(4, slotService.getWorkerProfile().getAssignedSlots().length);