Fix snapshot pipe auto-drop when some DataNodes fail to initialize

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/coordinator/runtime/heartbeat/PipeHeartbeatParser.java

@@ -39,6 +39,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -156,40 +157,48 @@ private void parseHeartbeatAndSaveMetaChangeLocally(
       final PipeTemporaryMetaInCoordinator temporaryMeta =
           (PipeTemporaryMetaInCoordinator) pipeMetaFromCoordinator.getTemporaryMeta();
 
+      final Set<Integer> expectedDataNodeIds = getExpectedDataNodeIds(pipeMetaFromCoordinator);
+
       // Remove completed pipes
       final Boolean isPipeCompletedFromAgent = pipeHeartbeat.isCompleted(staticMeta);
       if (Boolean.TRUE.equals(isPipeCompletedFromAgent)) {
 
-        temporaryMeta.markDataNodeCompleted(nodeId);
-        PipeLogger.log(
-            LOGGER::info,
-            ManagerMessages.DETECTED_HISTORICAL_PIPE_COMPLETION_REPORT_FROM_DATANODE,
-            nodeId,
-            staticMeta.getPipeName(),
-            pipeHeartbeat.getRemainingEventCount(staticMeta),
-            pipeHeartbeat.getRemainingTime(staticMeta),
-            temporaryMeta.getCompletedDataNodeIds());
-
-        final Set<Integer> uncompletedDataNodeIds =
-            configManager.getNodeManager().getRegisteredDataNodeLocations().keySet();
-        uncompletedDataNodeIds.removeAll(temporaryMeta.getCompletedDataNodeIds());
-        if (uncompletedDataNodeIds.isEmpty()) {
-          PipeLogger.log(
-              LOGGER::info,
-              ManagerMessages.ALL_DATANODES_REPORTED_HISTORICAL_PIPE_COMPLETED,
-              staticMeta.getPipeName(),
-              temporaryMeta.getGlobalRemainingEvents(),
-              temporaryMeta.getGlobalRemainingTime(),
-              staticMeta);
-          pipeTaskInfo.get().removePipeMeta(staticMeta);
+        if (expectedDataNodeIds.contains(nodeId)) {
+          temporaryMeta.markDataNodeCompleted(nodeId);
           PipeLogger.log(
               LOGGER::info,
-              ManagerMessages.DETECTED_COMPLETION_OF_PIPE_STATIC_META_REMOVE_IT,
+              ManagerMessages.DETECTED_HISTORICAL_PIPE_COMPLETION_REPORT_FROM_DATANODE,
+              nodeId,
               staticMeta.getPipeName(),
-              staticMeta);
-          needWriteConsensusOnConfigNodes.set(true);
-          needPushPipeMetaToDataNodes.set(true);
-          continue;
+              pipeHeartbeat.getRemainingEventCount(staticMeta),
+              pipeHeartbeat.getRemainingTime(staticMeta),
+              temporaryMeta.getCompletedDataNodeIds());
+        }
+
+        // Only DataNodes that are expected to run this Pipe participate in the completion
+        // judgment. A DataNode that does not own any target region should not block the Pipe
+        // from being automatically dropped after all expected DataNodes complete.
+        if (!expectedDataNodeIds.isEmpty()) {
+          final Set<Integer> uncompletedDataNodeIds = new HashSet<>(expectedDataNodeIds);
+          uncompletedDataNodeIds.removeAll(temporaryMeta.getCompletedDataNodeIds());
+          if (uncompletedDataNodeIds.isEmpty()) {
+            PipeLogger.log(
+                LOGGER::info,
+                ManagerMessages.ALL_DATANODES_REPORTED_HISTORICAL_PIPE_COMPLETED,
+                staticMeta.getPipeName(),
+                temporaryMeta.getGlobalRemainingEvents(),
+                temporaryMeta.getGlobalRemainingTime(),
+                staticMeta);
+            pipeTaskInfo.get().removePipeMeta(staticMeta);
+            PipeLogger.log(
+                LOGGER::info,
+                ManagerMessages.DETECTED_COMPLETION_OF_PIPE_STATIC_META_REMOVE_IT,
+                staticMeta.getPipeName(),
+                staticMeta);
+            needWriteConsensusOnConfigNodes.set(true);
+            needPushPipeMetaToDataNodes.set(true);
+            continue;
+          }
         }
       }
 
@@ -331,4 +340,22 @@ private void parseHeartbeatAndSaveMetaChangeLocally(
       }
     }
   }
+
+  // Returns the DataNodes that must complete this Pipe. It derives the expected set from the pipe's
+  // runtime metadata instead of all registered DataNodes, so DataNodes that do not own any target
+  // region are ignored during the auto-drop completion check.
+  private Set<Integer> getExpectedDataNodeIds(final PipeMeta pipeMeta) {
+    final Set<Integer> registeredDataNodeIds =
+        configManager.getNodeManager().getRegisteredDataNodeLocations().keySet();
+    final Set<Integer> expectedDataNodeIds = new HashSet<>();
+    for (final Map.Entry<Integer, PipeTaskMeta> entry :
+        pipeMeta.getRuntimeMeta().getConsensusGroupId2TaskMetaMap().entrySet()) {
+      // The ConfigRegion task is led by a ConfigNode, not by a DataNode.
+      if (entry.getKey() != Integer.MIN_VALUE
+          && registeredDataNodeIds.contains(entry.getValue().getLeaderNodeId())) {
+        expectedDataNodeIds.add(entry.getValue().getLeaderNodeId());
+      }
+    }
+    return expectedDataNodeIds;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java

@@ -542,11 +542,9 @@ private PipeMetaReport collectPipeMetaReport(
         final PipeStaticMeta staticMeta = pipeMeta.getStaticMeta();
 
         final Map<Integer, PipeTask> pipeTaskMap = pipeTaskManager.getPipeTasks(staticMeta);
+        final Set<Integer> expectedDataRegionIds = getExpectedDataRegionIds(pipeMeta);
         final boolean isAllDataRegionCompleted =
-            pipeTaskMap == null
-                || pipeTaskMap.entrySet().stream()
-                    .filter(entry -> dataRegionIds.contains(entry.getKey()))
-                    .allMatch(entry -> ((PipeDataNodeTask) entry.getValue()).isCompleted());
+            isAllExpectedDataRegionCompleted(pipeTaskMap, expectedDataRegionIds);
         final boolean isCompleted =
             isAllDataRegionCompleted && includeDataAndNeedDrop(pipeMeta, includeQueryMode);
         final Pair<Long, Double> remainingEventAndTime =
@@ -584,6 +582,58 @@ private PipeMetaReport collectPipeMetaReport(
     return report;
   }
 
+  // Returns whether every expected DataRegion has a completed local PipeTask on this DataNode.
+  // An empty expected set means this DataNode does not need to transfer history and is completed.
+  // A missing PipeTaskMap or a missing expected DataRegion means initialization failed.
+  static boolean isAllExpectedDataRegionCompleted(
+      final Map<Integer, PipeTask> pipeTaskMap, final Set<Integer> expectedDataRegionIds) {
+    if (expectedDataRegionIds.isEmpty()) {
+      // This DataNode does not own any target DataRegion for the pipe, so there is no local
+      // history transfer to wait for.
+      return true;
+    }
+    return pipeTaskMap != null
+        && expectedDataRegionIds.stream()
+            .allMatch(
+                dataRegionId -> {
+                  final PipeTask pipeTask = pipeTaskMap.get(dataRegionId);
+                  return pipeTask instanceof PipeDataNodeTask
+                      && ((PipeDataNodeTask) pipeTask).isCompleted();
+                });
+  }
+
+  // Returns the DataRegion ids that this DataNode is expected to transfer for the given pipe.
+  // A region is included only when it is owned by this DataNode, is led by this DataNode according
+  // to the pipe's runtime metadata, and is selected by the pipe's source parameters. This expected
+  // set is used instead of the already-created PipeTask map so that a failed task initialization is
+  // not silently treated as a completed region.
+  private Set<Integer> getExpectedDataRegionIds(final PipeMeta pipeMeta) {
+    final PipeStaticMeta staticMeta = pipeMeta.getStaticMeta();
+    final PipeParameters sourceParameters = staticMeta.getSourceParameters();
+    final Set<Integer> localDataRegionIds =
+        StorageEngine.getInstance().getAllDataRegionIds().stream()
+            .map(DataRegionId::getId)
+            .collect(Collectors.toSet());
+    final Set<Integer> expectedDataRegionIds = new HashSet<>();
+    for (final Map.Entry<Integer, PipeTaskMeta> entry :
+        pipeMeta.getRuntimeMeta().getConsensusGroupId2TaskMetaMap().entrySet()) {
+      final int regionId = entry.getKey();
+      if (entry.getValue().getLeaderNodeId() != CONFIG.getDataNodeId()
+          || !localDataRegionIds.contains(regionId)) {
+        continue;
+      }
+      try {
+        if (DataRegionListeningFilter.shouldDataRegionBeListened(
+            sourceParameters, new DataRegionId(regionId), staticMeta.getPipeType())) {
+          expectedDataRegionIds.add(regionId);
+        }
+      } catch (final IllegalPathException e) {
+        throw new PipeException(e.toString());
+      }
+    }
+    return expectedDataRegionIds;
+  }
+
   private boolean includeDataAndNeedDrop(final PipeMeta pipeMeta, final boolean includeQueryMode)
       throws IllegalPathException {
     final PipeParameters sourceParameters = pipeMeta.getStaticMeta().getSourceParameters();

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentTest.java

@@ -23,6 +23,7 @@
 import org.apache.iotdb.commons.consensus.index.ProgressIndex;
 import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
 import org.apache.iotdb.commons.consensus.index.impl.SimpleProgressIndex;
+import org.apache.iotdb.commons.pipe.agent.task.PipeTask;
 import org.apache.iotdb.commons.pipe.agent.task.PipeTaskAgent;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMetaKeeper;
@@ -35,18 +36,63 @@
 
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
+import static org.mockito.Mockito.when;
+
 public class PipeDataNodeTaskAgentTest {
 
   private static final int LOCAL_NODE_ID = 1;
   private static final int REGION_ID = 7;
 
+  @Test
+  public void testIsAllExpectedDataRegionCompleted() {
+    final PipeDataNodeTask completedTask = Mockito.mock(PipeDataNodeTask.class);
+    when(completedTask.isCompleted()).thenReturn(true);
+    final PipeDataNodeTask uncompletedTask = Mockito.mock(PipeDataNodeTask.class);
+    when(uncompletedTask.isCompleted()).thenReturn(false);
+
+    final Map<Integer, PipeTask> pipeTaskMap = new HashMap<>();
+    pipeTaskMap.put(1, completedTask);
+    pipeTaskMap.put(2, completedTask);
+
+    final Set<Integer> completedExpectedRegionIds = new HashSet<>(Arrays.asList(1, 2));
+    Assert.assertTrue(
+        PipeDataNodeTaskAgent.isAllExpectedDataRegionCompleted(
+            pipeTaskMap, completedExpectedRegionIds));
+
+    // A DataRegion that should be transferred is missing its local PipeTask.
+    final Set<Integer> partiallyMissingExpectedRegionIds = new HashSet<>(Arrays.asList(1, 2, 3));
+    Assert.assertFalse(
+        PipeDataNodeTaskAgent.isAllExpectedDataRegionCompleted(
+            pipeTaskMap, partiallyMissingExpectedRegionIds));
+
+    // No local target DataRegion means this DataNode does not need to transfer history.
+    Assert.assertTrue(
+        PipeDataNodeTaskAgent.isAllExpectedDataRegionCompleted(null, Collections.emptySet()));
+
+    // A non-empty expected set with a missing PipeTaskMap means initialization failed.
+    Assert.assertFalse(
+        PipeDataNodeTaskAgent.isAllExpectedDataRegionCompleted(
+            null, partiallyMissingExpectedRegionIds));
+
+    // An uncompleted PipeTask must not be reported as completed.
+    pipeTaskMap.put(3, uncompletedTask);
+    Assert.assertFalse(
+        PipeDataNodeTaskAgent.isAllExpectedDataRegionCompleted(
+            pipeTaskMap, partiallyMissingExpectedRegionIds));
+  }
+
   @Test
   public void testGetPipeTaskProgressIndexReportsMissingTaskMeta() throws Exception {
     final PipeDataNodeTaskAgent taskAgent = new PipeDataNodeTaskAgent();