[Fix] Guard stopped Pipe polling and retain LOAD payloads until replicated

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtask.java

@@ -176,16 +176,22 @@ private void initSink(final Map<String, String> sinkAttributes) throws Exception
    */
   @Override
   protected boolean executeOnce() throws Exception {
-    if (isClosed.get()) {
+    if (isClosed.get() || shouldStopSubmittingSelf.get()) {
       return false;
     }
 
     final Event event = lastEvent != null ? lastEvent : source.supply();
-    // Record the last event for retry when exception occurs
+    // Keep the event for retry or restart if supply() returns after the task is stopped.
     setLastEvent(event);
 
+    // close() may have released the last event before supply() returned.
+    if (isClosed.get()) {
+      clearReferenceCountAndReleaseLastEvent(event);
+      return false;
+    }
+
     try {
-      if (event == null) {
+      if (event == null || shouldStopSubmittingSelf.get()) {
         return false;
       }
 

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtaskTest.java

@@ -0,0 +1,150 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.confignode.manager.pipe.agent.task;
+
+import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
+import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
+import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeAbstractSinkSubtask;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
+import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
+import org.apache.iotdb.pipe.api.PipeConnector;
+import org.apache.iotdb.pipe.api.PipeExtractor;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
+
+public class PipeConfigNodeSubtaskTest {
+
+  private PipeConfigNodeSubtask subtask;
+  private PipeExtractor source;
+  private PipeConnector sink;
+  private boolean subtaskClosed;
+
+  @Before
+  public void setUp() throws Exception {
+    source = mock(PipeExtractor.class);
+    sink = mock(PipeConnector.class);
+    subtask =
+        new PipeConfigNodeSubtask(
+            "PipeConfigNodeSubtaskTest",
+            System.currentTimeMillis(),
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.singletonMap(
+                PipeSinkConstant.CONNECTOR_KEY,
+                BuiltinPipePlugin.DO_NOTHING_CONNECTOR.getPipePluginName()),
+            new PipeTaskMeta(MinimumProgressIndex.INSTANCE, Integer.MIN_VALUE));
+    final Field sourceField = PipeConfigNodeSubtask.class.getDeclaredField("source");
+    sourceField.setAccessible(true);
+    sourceField.set(subtask, source);
+    final Field sinkField = PipeAbstractSinkSubtask.class.getDeclaredField("outputPipeSink");
+    sinkField.setAccessible(true);
+    sinkField.set(subtask, sink);
+  }
+
+  @After
+  public void tearDown() {
+    if (!subtaskClosed) {
+      subtask.close();
+    }
+  }
+
+  @Test
+  public void testStoppedSubtaskDoesNotPollSource() throws Exception {
+    assertFalse(subtask.executeOnce());
+    verifyZeroInteractions(source, sink);
+  }
+
+  @Test
+  public void testStopDuringSupplyRetainsEventUntilRestart() throws Exception {
+    assertStopDuringSupply(false);
+  }
+
+  @Test
+  public void testDropDuringSupplyReleasesEvent() throws Exception {
+    assertStopDuringSupply(true);
+  }
+
+  private void assertStopDuringSupply(final boolean drop) throws Exception {
+    final EnrichedEvent event = mock(EnrichedEvent.class);
+    final CountDownLatch supplyStarted = new CountDownLatch(1);
+    final CountDownLatch supplyReleased = new CountDownLatch(1);
+    final ExecutorService worker = Executors.newSingleThreadExecutor();
+    when(source.supply())
+        .thenAnswer(
+            invocation -> {
+              supplyStarted.countDown();
+              assertTrue(supplyReleased.await(10, TimeUnit.SECONDS));
+              return event;
+            });
+    subtask.allowSubmittingSelf();
+    try {
+      final Future<Boolean> execution = worker.submit(() -> subtask.executeOnce());
+      assertTrue(supplyStarted.await(10, TimeUnit.SECONDS));
+      assertTrue(subtask.disallowSubmittingSelf());
+      if (drop) {
+        subtask.close();
+        subtaskClosed = true;
+      }
+      supplyReleased.countDown();
+
+      assertFalse(execution.get(10, TimeUnit.SECONDS));
+      assertFalse(subtask.executeOnce());
+      verify(sink, never()).transfer(event);
+      verify(event, never()).decreaseReferenceCount(anyString(), anyBoolean());
+
+      if (drop) {
+        verify(event).clearReferenceCount(anyString());
+      } else {
+        verifyZeroInteractions(sink);
+        subtask.allowSubmittingSelf();
+        assertTrue(subtask.executeOnce());
+        verify(source).supply();
+        verify(sink).transfer(event);
+        verify(event).decreaseReferenceCount(anyString(), eq(true));
+      }
+    } finally {
+      supplyReleased.countDown();
+      worker.shutdownNow();
+      assertTrue(worker.awaitTermination(10, TimeUnit.SECONDS));
+    }
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRetention.java

@@ -62,14 +62,7 @@ final class LoadTaskRetention {
 
   /** Checks whether the staged directory must be retained for follower WAL synchronization. */
   boolean mustRetain(final long searchIndex) {
-    // IoTConsensusV2 replicates through pipe, so no replica reads these bytes back out of the WAL
-    // and there is no index to wait for: the staged directory can go as soon as the task ends.
-    if (LoadConsensusCompat.usesIoTConsensusV2()) {
-      return false;
-    }
-    final long safelyDeletedSearchIndex = currentSafelyDeletedSearchIndex();
-    return safelyDeletedSearchIndex != ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX
-        && safelyDeletedSearchIndex < searchIndex;
+    return currentSafelyDeletedSearchIndex() < searchIndex;
   }
 
   /** Takes ownership of a staged directory, postponing deletion until WAL watermarks allow it. */
@@ -122,14 +115,14 @@ void release() {
     }
 
     final long safelyDeletedSearchIndex = currentSafelyDeletedSearchIndex();
-    final long boundary =
-        safelyDeletedSearchIndex == ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX
-            ? Long.MAX_VALUE
-            : safelyDeletedSearchIndex;
+    // The default V1 watermark means no replica progress has been confirmed yet.
+    if (safelyDeletedSearchIndex == ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX) {
+      return;
+    }
 
     for (final Map.Entry<String, RetainedTaskDir> entry : loadId2RetainedTaskDir.entrySet()) {
       final RetainedTaskDir retained = entry.getValue();
-      if (retained.reachedSearchIndex() <= boundary
+      if (retained.reachedSearchIndex() <= safelyDeletedSearchIndex
           && loadId2RetainedTaskDir.remove(entry.getKey(), retained)) {
         LOGGER.info(
             StorageEngineMessages
@@ -165,10 +158,15 @@ private void ensureWalListenerRegistered() {
   }
 
   private long currentSafelyDeletedSearchIndex() {
+    // V2 replicates through pipe. Other protocols without a V1 WAL (for example Ratis) do not
+    // expand staged references from this node's WAL either, so neither has a watermark to wait for.
+    if (LoadConsensusCompat.usesIoTConsensusV2()) {
+      return Long.MAX_VALUE;
+    }
     return dataRegion
         .getWALNode()
         .map(IWALNode::getSafelyDeletedSearchIndex)
-        .orElse(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX);
+        .orElse(Long.MAX_VALUE);
   }
 
   // -------------------------------------------------------------------------

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleaner.java

@@ -350,14 +350,22 @@ private static boolean isComplete(final File taskDir) {
   }
 
   private static boolean replicasReached(final DataRegion dataRegion, final long searchIndex) {
-    final long safelyDeletedSearchIndex =
-        dataRegion
-            .getWALNode()
-            .map(wal -> wal.getSafelyDeletedSearchIndex())
-            .orElse(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX);
-
-    return safelyDeletedSearchIndex == ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX
-        || safelyDeletedSearchIndex >= searchIndex;
+    // IoTConsensusV2 and consensus implementations without a WAL replicate the staged payload
+    // directly, so no follower can read it back from this node's WAL. IoTConsensus V1, however,
+    // must report a real watermark before its staged references are safe to delete.
+    if (LoadConsensusCompat.usesIoTConsensusV2()) {
+      return true;
+    }
+    return dataRegion
+        .getWALNode()
+        .map(
+            wal -> {
+              final long safelyDeletedSearchIndex = wal.getSafelyDeletedSearchIndex();
+              return safelyDeletedSearchIndex
+                      != ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX
+                  && safelyDeletedSearchIndex >= searchIndex;
+            })
+        .orElse(true);
   }
 
   // -------------------------------------------------------------------------

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleanerTest.java

@@ -22,6 +22,7 @@
 import org.apache.iotdb.commons.conf.ConfigurationFileUtils;
 import org.apache.iotdb.commons.conf.TrimProperties;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.consensus.iot.log.ConsensusReqReader;
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode;
@@ -186,6 +187,47 @@ public void testAnIdleTaskIsReclaimedOnlyOnceItsPiecesAreReplicated() throws Exc
     }
   }
 
+  @Test
+  public void testAFinishedTaskWaitsForTheFirstWatermark() throws Exception {
+    final String loadId = "unreported-watermark-task";
+    final File taskDir = new File(tempDir, loadId);
+    Assert.assertTrue(taskDir.mkdirs());
+    stageOneChunk(taskDir, 5L);
+    final LoadTsFileCleaner cleaner = LoadTsFileCleaner.getInstance();
+    cleaner.register(loadId, dataRegion, taskDir, 7L, LoadTsFileConsensusOp.ABORT);
+
+    walWatermark.set(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX);
+    cleaner.sweepRegisteredTasks();
+    Assert.assertTrue("an unreported watermark does not confirm replication", taskDir.exists());
+
+    walWatermark.set(5L);
+    cleaner.sweepRegisteredTasks();
+    Assert.assertTrue("the terminal command is not replicated yet", taskDir.exists());
+
+    walWatermark.set(7L);
+    cleaner.sweepRegisteredTasks();
+    Assert.assertFalse("the terminal command is replicated", taskDir.exists());
+  }
+
+  @Test
+  public void testIoTConsensusV2DoesNotWaitForAWALWatermark() throws Exception {
+    final String loadId = "v2-finished-task";
+    final File taskDir = new File(tempDir, loadId);
+    Assert.assertTrue(taskDir.mkdirs());
+    stageOneChunk(taskDir, 5L);
+    final LoadTsFileCleaner cleaner = LoadTsFileCleaner.getInstance();
+    cleaner.register(loadId, dataRegion, taskDir, 7L, LoadTsFileConsensusOp.ABORT);
+    walWatermark.set(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX);
+
+    LoadConsensusCompat.setIoTConsensusV2ForTest(Boolean.TRUE);
+    try {
+      cleaner.sweepRegisteredTasks();
+      Assert.assertFalse("V2 does not expand staged references from the WAL", taskDir.exists());
+    } finally {
+      LoadConsensusCompat.setIoTConsensusV2ForTest(null);
+    }
+  }
+
   @Test
   public void testAFinishedCommittedTaskIsReleasedOnceItsStagedFilesAreComplete() throws Exception {
     final String loadId = "finished-task";