Optimize pipe TsFile cleanup on drop

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -1343,6 +1343,8 @@ public final class DataNodePipeMessages {
       "Finished deleting all stale pipe hardlink dirs by periodical job.";
   public static final String PIPE_HARDLINK_DIR_MOVE_FAILED_DELETING_SYNC =
       "Failed to move pipe hardlink dir {} for periodical deletion, deleting it synchronously.";
+  public static final String PIPE_HARDLINK_DIR_MOVE_FAILED_SKIPPING_PERIODICAL_DELETE =
+      "Failed to move pipe hardlink dir {} for periodical deletion, skip registering the original dir to avoid deleting files of a recreated pipe.";
   public static final String PIPE_SNAPSHOT_DIR_FOUND_DELETING_IT =
       "Pipe snapshot dir found, deleting it: {},";
   public static final String SHRINK_CALLBACK_IS_NOT_SUPPORTED_IN_PIPEFIXEDMEMORYBLOCK =

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -1266,6 +1266,8 @@ public final class DataNodePipeMessages {
       "已通过周期性任务删除全部过期的 Pipe hardlink 目录。";
   public static final String PIPE_HARDLINK_DIR_MOVE_FAILED_DELETING_SYNC =
       "为周期性删除移动 Pipe hardlink 目录 {} 失败，改为同步删除。";
+  public static final String PIPE_HARDLINK_DIR_MOVE_FAILED_SKIPPING_PERIODICAL_DELETE =
+      "为周期性删除移动 Pipe hardlink 目录 {} 失败，跳过注册原目录以避免删除重建 Pipe 的文件。";
   public static final String PIPE_SNAPSHOT_DIR_FOUND_DELETING_IT =
       "发现 Pipe snapshot 目录，正在删除：{},";
   public static final String SHRINK_CALLBACK_IS_NOT_SUPPORTED_IN_PIPEFIXEDMEMORYBLOCK =

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

@@ -59,6 +59,7 @@
 import org.apache.iotdb.db.pipe.metric.overview.PipeTsFileToTabletsMetrics;
 import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
 import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager;
+import org.apache.iotdb.db.pipe.resource.tsfile.PipeTsFileResourceManager;
 import org.apache.iotdb.db.pipe.source.dataregion.DataRegionListeningFilter;
 import org.apache.iotdb.db.pipe.source.dataregion.realtime.listener.PipeInsertionDataNodeListener;
 import org.apache.iotdb.db.pipe.source.schemaregion.SchemaRegionListeningFilter;
@@ -304,13 +305,20 @@ protected void freezeRate(final String pipeName, final long creationTime) {
 
   @Override
   protected boolean dropPipe(final String pipeName, final long creationTime) {
+    final String pipeTsFileResourcePipeName =
+        PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime);
+    PipeDataNodeResourceManager.tsfile().markPipeTsFileDirUnderDeletion(pipeTsFileResourcePipeName);
+
     if (!super.dropPipe(pipeName, creationTime)) {
+      PipeDataNodeResourceManager.tsfile()
+          .unmarkPipeTsFileDirUnderDeletion(pipeTsFileResourcePipeName);
       return false;
     }
 
     final String taskId = pipeName + "_" + creationTime;
     PipeTsFileToTabletsMetrics.getInstance().deregister(taskId);
     PipeDataNodeSinglePipeMetrics.getInstance().deregister(taskId);
+    PipeDataNodeResourceManager.tsfile().cleanPipeTsFileDir(pipeTsFileResourcePipeName);
 
     return true;
   }
@@ -319,6 +327,15 @@ protected boolean dropPipe(final String pipeName, final long creationTime) {
   protected boolean dropPipe(final String pipeName) {
     // Get the pipe meta first because it is removed after super#dropPipe(pipeName)
     final PipeMeta pipeMeta = pipeMetaKeeper.getPipeMeta(pipeName);
+    final String pipeTsFileResourcePipeName =
+        Objects.isNull(pipeMeta)
+            ? null
+            : PipeTsFileResourceManager.getPipeTsFileResourcePipeName(
+                pipeName, pipeMeta.getStaticMeta().getCreationTime());
+    if (Objects.nonNull(pipeTsFileResourcePipeName)) {
+      PipeDataNodeResourceManager.tsfile()
+          .markPipeTsFileDirUnderDeletion(pipeTsFileResourcePipeName);
+    }
 
     // Record whether there are pipe tasks before dropping the pipe
     final boolean hasPipeTasks;
@@ -331,6 +348,10 @@ protected boolean dropPipe(final String pipeName) {
     }
 
     if (!super.dropPipe(pipeName)) {
+      if (Objects.nonNull(pipeTsFileResourcePipeName)) {
+        PipeDataNodeResourceManager.tsfile()
+            .unmarkPipeTsFileDirUnderDeletion(pipeTsFileResourcePipeName);
+      }
       return false;
     }
 
@@ -339,6 +360,7 @@ protected boolean dropPipe(final String pipeName) {
       final String taskId = pipeName + "_" + creationTime;
       PipeTsFileToTabletsMetrics.getInstance().deregister(taskId);
       PipeDataNodeSinglePipeMetrics.getInstance().deregister(taskId);
+      PipeDataNodeResourceManager.tsfile().cleanPipeTsFileDir(pipeTsFileResourcePipeName);
       // When the pipe contains no pipe tasks, there is no corresponding prefetching queue for the
       // subscribed pipe, so the subscription needs to be manually marked as completed.
       if (!hasPipeTasks && PipeStaticMeta.isSubscriptionPipe(pipeName)) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/PipeDataNodeHardlinkOrCopiedFileDirStartupCleaner.java

@@ -37,6 +37,7 @@
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -49,6 +50,9 @@ public class PipeDataNodeHardlinkOrCopiedFileDirStartupCleaner {
       "PipeDataNodeHardlinkOrCopiedFileDirStartupCleaner#cleanTsFileDir()";
   private static final long DELETE_MAX_PATH_COUNT_PER_ROUND = 100_000L;
   private static final long DELETE_MAX_TIME_PER_ROUND_MS = 1_000L;
+  private static final PeriodicalStalePipeDirCleaner PERIODICAL_STALE_PIPE_DIR_CLEANER =
+      new PeriodicalStalePipeDirCleaner();
+  private static final AtomicBoolean PERIODICAL_CLEANUP_JOB_REGISTERED = new AtomicBoolean(false);
 
   /**
    * Delete the data directory and all of its subdirectories that contain the
@@ -71,7 +75,8 @@ private static void cleanTsFileDir(final PeriodicalJobRegistrar periodicalJobReg
         moveAsideAndCollect(pipeHardLinkDir, pipeHardlinkBaseDirName, stalePipeDirs);
       }
     }
-    registerPeriodicalCleanupJob(periodicalJobRegistrar, stalePipeDirs);
+    PERIODICAL_STALE_PIPE_DIR_CLEANER.addStalePipeDirs(stalePipeDirs);
+    registerPeriodicalCleanupJob(periodicalJobRegistrar);
   }
 
   private static void collectInterruptedStalePipeDirs(
@@ -82,7 +87,7 @@ private static void collectInterruptedStalePipeDirs(
         localDataDir.listFiles(
             file ->
                 file.isDirectory()
-                    && file.getName().startsWith(pipeHardlinkBaseDirName + STALE_PIPE_DIR_SUFFIX));
+                    && file.getName().contains(pipeHardlinkBaseDirName + STALE_PIPE_DIR_SUFFIX));
     if (stalePipeDirFiles == null) {
       return;
     }
@@ -137,17 +142,47 @@ private static File moveAside(final File pipeHardLinkDir, final String pipeHardl
   }
 
   private static void registerPeriodicalCleanupJob(
-      final PeriodicalJobRegistrar periodicalJobRegistrar, final List<File> stalePipeDirs) {
-    if (stalePipeDirs.isEmpty()) {
+      final PeriodicalJobRegistrar periodicalJobRegistrar) {
+    if (!PERIODICAL_CLEANUP_JOB_REGISTERED.compareAndSet(false, true)) {
       return;
     }
 
     periodicalJobRegistrar.register(
         PERIODICAL_CLEANUP_JOB_ID,
-        new PeriodicalStalePipeDirCleaner(stalePipeDirs)::cleanOneRound,
+        PERIODICAL_STALE_PIPE_DIR_CLEANER::cleanOneRound,
         PipeConfig.getInstance().getPipeSubtaskExecutorCronHeartbeatEventIntervalSeconds());
   }
 
+  public static void submitStalePipeDirForPeriodicalCleanup(final File stalePipeDir) {
+    if (stalePipeDir == null) {
+      return;
+    }
+
+    final File stalePipeDirToDelete;
+    if (stalePipeDir.isDirectory()) {
+      try {
+        stalePipeDirToDelete = moveAside(stalePipeDir, stalePipeDir.getName());
+        LOGGER.info(
+            DataNodePipeMessages.PIPE_HARDLINK_DIR_FOUND_MOVED_TO_PERIODICAL_DELETE,
+            stalePipeDir,
+            stalePipeDirToDelete);
+      } catch (final IOException e) {
+        LOGGER.warn(
+            DataNodePipeMessages.PIPE_HARDLINK_DIR_MOVE_FAILED_SKIPPING_PERIODICAL_DELETE,
+            stalePipeDir,
+            e);
+        return;
+      }
+    } else {
+      stalePipeDirToDelete = stalePipeDir;
+    }
+
+    LOGGER.info(
+        DataNodePipeMessages.PIPE_STALE_HARDLINK_DIR_FOUND_REGISTERING_PERIODICAL_DELETE,
+        stalePipeDirToDelete);
+    PERIODICAL_STALE_PIPE_DIR_CLEANER.addStalePipeDir(stalePipeDirToDelete);
+  }
+
   private static CleanupRoundResult deleteQuietlyWithThrottle(final File stalePipeDir) {
     if (!stalePipeDir.exists()) {
       return CleanupRoundResult.finished();
@@ -230,48 +265,48 @@ public interface PeriodicalJobRegistrar {
 
   private static class PeriodicalStalePipeDirCleaner {
 
-    private final List<File> stalePipeDirs;
-    private int currentDirIndex;
-    private boolean finished;
+    private final ConcurrentLinkedQueue<File> stalePipeDirs = new ConcurrentLinkedQueue<>();
+    private File currentStalePipeDir;
 
-    private PeriodicalStalePipeDirCleaner(final List<File> stalePipeDirs) {
-      this.stalePipeDirs = stalePipeDirs;
-      currentDirIndex = 0;
-      finished = false;
+    private void addStalePipeDir(final File stalePipeDir) {
+      stalePipeDirs.offer(stalePipeDir);
     }
 
-    private void cleanOneRound() {
-      if (finished) {
-        return;
-      }
+    private void addStalePipeDirs(final List<File> stalePipeDirs) {
+      stalePipeDirs.forEach(this::addStalePipeDir);
+    }
 
+    private void cleanOneRound() {
       long deletedPathCount = 0;
-      while (currentDirIndex < stalePipeDirs.size()) {
-        final File stalePipeDir = stalePipeDirs.get(currentDirIndex);
-        final CleanupRoundResult result = deleteQuietlyWithThrottle(stalePipeDir);
+      while (true) {
+        if (currentStalePipeDir == null) {
+          currentStalePipeDir = stalePipeDirs.poll();
+        }
+        if (currentStalePipeDir == null) {
+          return;
+        }
+
+        final CleanupRoundResult result = deleteQuietlyWithThrottle(currentStalePipeDir);
         deletedPathCount += result.deletedPathCount;
 
         if (result.finished) {
           LOGGER.info(
               DataNodePipeMessages.PIPE_HARDLINK_DIR_PERIODICAL_DELETE_FINISHED,
-              stalePipeDir,
+              currentStalePipeDir,
               result.success);
-          ++currentDirIndex;
+          currentStalePipeDir = null;
           continue;
         }
 
         if (deletedPathCount > 0 || !result.success) {
           LOGGER.info(
               DataNodePipeMessages.PIPE_HARDLINK_DIR_PERIODICAL_DELETE_PROGRESS,
               deletedPathCount,
-              stalePipeDir,
+              currentStalePipeDir,
               result.success);
         }
         return;
       }
-
-      finished = true;
-      LOGGER.info(DataNodePipeMessages.PIPE_HARDLINK_DIR_PERIODICAL_DELETE_ALL_FINISHED);
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResource.java

@@ -70,9 +70,15 @@ public void increaseReferenceCount() {
   }
 
   public boolean decreaseReferenceCount() {
+    return decreaseReferenceCount(true);
+  }
+
+  public boolean decreaseReferenceCount(final boolean deleteFileWhenNoReference) {
     final int finalReferenceCount = referenceCount.addAndGet(-1);
     if (finalReferenceCount == 0) {
-      close();
+      if (deleteFileWhenNoReference) {
+        close();
+      }
       return true;
     }
     if (finalReferenceCount < 0) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java

@@ -24,6 +24,7 @@
 import org.apache.iotdb.commons.utils.FileUtils;
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.i18n.DataNodePipeMessages;
+import org.apache.iotdb.db.pipe.resource.PipeDataNodeHardlinkOrCopiedFileDirStartupCleaner;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
 
 import org.apache.tsfile.enums.TSDataType;
@@ -39,6 +40,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class PipeTsFileResourceManager {
@@ -54,6 +56,9 @@ public class PipeTsFileResourceManager {
   // PipeName -> TsFilePath -> PipeTsFileResource
   private final Map<String, Map<String, PipeTsFileResource>>
       hardlinkOrCopiedFileToPipeTsFileResourceMap = new ConcurrentHashMap<>();
+  private final Map<String, String> pipeNameToPipeTsFileDirPathMap = new ConcurrentHashMap<>();
+  private final Set<String> pipeTsFileResourcePipeNameSetUnderDeletion =
+      ConcurrentHashMap.newKeySet();
   private final PipeTsFileResourceSegmentLock segmentLock = new PipeTsFileResourceSegmentLock();
 
   public static String getPipeTsFileResourcePipeName(
@@ -126,6 +131,8 @@ private File increaseFileReference(
       // file in pipe dir, create a hardlink or copy it to pipe dir, maintain a reference count for
       // the hardlink or copied file, and return the hardlink or copied file.
       if (Objects.nonNull(pipeName)) {
+        pipeNameToPipeTsFileDirPathMap.putIfAbsent(
+            pipeName, hardlinkOrCopiedFile.getParentFile().getPath());
         hardlinkOrCopiedFileToPipeTsFileResourceMap
             .computeIfAbsent(pipeName, k -> new ConcurrentHashMap<>())
             .put(resultFile.getPath(), new PipeTsFileResource(resultFile));
@@ -225,7 +232,8 @@ public void decreaseFileReference(
     try {
       final String filePath = hardlinkOrCopiedFile.getPath();
       final PipeTsFileResource resource = getResourceMap(pipeName).get(filePath);
-      if (resource != null && resource.decreaseReferenceCount()) {
+      if (resource != null
+          && resource.decreaseReferenceCount(shouldDeleteFileWhenNoReference(pipeName))) {
         getResourceMap(pipeName).remove(filePath);
       }
     } finally {
@@ -246,6 +254,35 @@ private void decreasePublicReferenceIfExists(final File file, final @Nullable St
     decreaseFileReference(new File(getCommonFilePath(file)), null);
   }
 
+  private boolean shouldDeleteFileWhenNoReference(
+      final @Nullable String pipeTsFileResourcePipeName) {
+    return Objects.isNull(pipeTsFileResourcePipeName)
+        || !pipeTsFileResourcePipeNameSetUnderDeletion.contains(pipeTsFileResourcePipeName);
+  }
+
+  public void markPipeTsFileDirUnderDeletion(final @Nonnull String pipeTsFileResourcePipeName) {
+    pipeTsFileResourcePipeNameSetUnderDeletion.add(pipeTsFileResourcePipeName);
+  }
+
+  public void unmarkPipeTsFileDirUnderDeletion(final @Nonnull String pipeTsFileResourcePipeName) {
+    pipeTsFileResourcePipeNameSetUnderDeletion.remove(pipeTsFileResourcePipeName);
+  }
+
+  public void cleanPipeTsFileDir(final @Nonnull String pipeTsFileResourcePipeName) {
+    final String pipeTsFileDirPath =
+        pipeNameToPipeTsFileDirPathMap.remove(pipeTsFileResourcePipeName);
+    hardlinkOrCopiedFileToPipeTsFileResourceMap.remove(pipeTsFileResourcePipeName);
+
+    if (Objects.isNull(pipeTsFileDirPath)) {
+      pipeTsFileResourcePipeNameSetUnderDeletion.remove(pipeTsFileResourcePipeName);
+      return;
+    }
+
+    PipeDataNodeHardlinkOrCopiedFileDirStartupCleaner.submitStalePipeDirForPeriodicalCleanup(
+        new File(pipeTsFileDirPath));
+    pipeTsFileResourcePipeNameSetUnderDeletion.remove(pipeTsFileResourcePipeName);
+  }
+
   // Warning: Shall not be called by the assigner
   private String getCommonFilePath(final @Nonnull File file) {
     // If the parent or grandparent is null then this is testing scenario