Load: Fix active load reliability issues

integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileIT.java

@@ -20,8 +20,11 @@
 package org.apache.iotdb.db.it;
 
 import org.apache.iotdb.commons.auth.entity.PrivilegeType;
+import org.apache.iotdb.commons.path.MeasurementPath;
 import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant;
 import org.apache.iotdb.db.it.utils.TestUtils;
+import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
+import org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry;
 import org.apache.iotdb.it.env.EnvFactory;
 import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
 import org.apache.iotdb.it.framework.IoTDBTestRunner;
@@ -985,6 +988,62 @@ public void testLoadWithEmptyTsFile() throws Exception {
     }
   }
 
+  @Test
+  public void testAsyncLoadKeepsSameNamedTsFilesWithModsIsolated() throws Exception {
+    registerSchema();
+
+    long expectedPointCount = 0;
+    // Before each file group had its own transfer directory, these same-named TsFiles and mods
+    // were renamed independently in one shared directory and could be paired with the wrong file.
+    for (int i = 0; i < 2; i++) {
+      final File sourceDir = new File(tmpDir, "source-" + i);
+      Assert.assertTrue(sourceDir.mkdirs());
+      try (final TsFileGenerator generator =
+          new TsFileGenerator(new File(sourceDir, "1-0-0-0.tsfile"))) {
+        generator.resetRandom(i);
+        generator.registerTimeseries(
+            SchemaConfig.DEVICE_0, Collections.singletonList(SchemaConfig.MEASUREMENT_00));
+        generator.generateData(SchemaConfig.DEVICE_0, 20, 1, false, TimeUnit.SECONDS.toMillis(i));
+        // Each group contributes 20 points and its own mods deletes exactly one. Losing either
+        // TsFile-to-mods pairing therefore leaves 39 points instead of the expected 38.
+        expectedPointCount += generator.getTotalNumber() - 1;
+      }
+      try (final ModificationFile modificationFile =
+          new ModificationFile(
+              new File(sourceDir, "1-0-0-0.tsfile" + ModificationFile.FILE_SUFFIX), false)) {
+        modificationFile.write(
+            new TreeDeletionEntry(
+                new MeasurementPath(
+                    SchemaConfig.DEVICE_0, SchemaConfig.MEASUREMENT_00.getMeasurementName()),
+                TimeUnit.SECONDS.toMillis(i) + 1,
+                TimeUnit.SECONDS.toMillis(i) + 1));
+      }
+    }
+
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute(
+          String.format(
+              "load \"%s\" with ('async'='true','database-level'='2','on-success'='delete')",
+              tmpDir.getAbsolutePath()));
+    }
+
+    TestUtils.assertDataEventuallyOnEnv(
+        EnvFactory.getEnv(),
+        "select count("
+            + SchemaConfig.MEASUREMENT_00.getMeasurementName()
+            + ") from "
+            + SchemaConfig.DEVICE_0,
+        Collections.singletonMap(
+            "count("
+                + SchemaConfig.DEVICE_0
+                + "."
+                + SchemaConfig.MEASUREMENT_00.getMeasurementName()
+                + ")",
+            Long.toString(expectedPointCount)),
+        30);
+  }
+
   @Test
   public void testLoadTsFileWithWrongTimestampPrecision() throws Exception {
     try (final TsFileGenerator generator =

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadSingleTsFileNode.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.commons.utils.RetryUtils;
 import org.apache.iotdb.commons.utils.TimePartitionUtils;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
@@ -241,17 +242,24 @@ public String toString() {
   }
 
   public void clean() {
+    if (!deleteAfterLoad) {
+      return;
+    }
+    deleteFile(tsFile);
+    deleteFile(new File(LoadUtil.getTsFileResourcePath(tsFile.getAbsolutePath())));
+    deleteFile(ModificationFile.getExclusiveMods(tsFile));
+    deleteFile(new File(LoadUtil.getTsFileModsV1Path(tsFile.getAbsolutePath())));
+  }
+
+  private void deleteFile(final File file) {
     try {
-      if (deleteAfterLoad) {
-        Files.deleteIfExists(tsFile.toPath());
-        Files.deleteIfExists(
-            new File(LoadUtil.getTsFileResourcePath(tsFile.getAbsolutePath())).toPath());
-        Files.deleteIfExists(ModificationFile.getExclusiveMods(tsFile).toPath());
-        Files.deleteIfExists(
-            new File(LoadUtil.getTsFileModsV1Path(tsFile.getAbsolutePath())).toPath());
-      }
-    } catch (final IOException e) {
-      LOGGER.warn(DataNodeQueryMessages.DELETE_AFTER_LOADING_ERROR, tsFile, e);
+      RetryUtils.retryOnException(
+          () -> {
+            Files.deleteIfExists(file.toPath());
+            return null;
+          });
+    } catch (final Exception e) {
+      LOGGER.warn(DataNodeQueryMessages.DELETE_AFTER_LOADING_ERROR, file, e);
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java

@@ -800,8 +800,7 @@ private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException
                       singleTsFileNode
                           .getTsFileResource()
                           .getTsFile()))); // can not just remove, because of deletion
-          dataSize -= pieceNode.getDataSize();
-          block.reduceMemoryUsage(pieceNode.getDataSize());
+          releaseMemoryUsage(pieceNode.getDataSize());
 
           if (!isDispatchSuccess) {
             // Currently there is no retry, so return directly
@@ -886,7 +885,7 @@ private boolean sendAllTsFileData() throws LoadFileException {
       boolean isAllSuccess = true;
       for (Map.Entry<TConsensusGroupId, Pair<TRegionReplicaSet, LoadTsFilePieceNode>> entry :
           regionId2ReplicaSetAndNode.entrySet()) {
-        block.reduceMemoryUsage(entry.getValue().getRight().getDataSize());
+        releaseMemoryUsage(entry.getValue().getRight().getDataSize());
         if (isAllSuccess
             && !scheduler.dispatchOnePieceNode(
                 entry.getValue().getRight(), entry.getValue().getLeft())) {
@@ -900,7 +899,17 @@ private boolean sendAllTsFileData() throws LoadFileException {
       return isAllSuccess;
     }
 
+    private void releaseMemoryUsage(final long memorySize) {
+      dataSize -= memorySize;
+      block.reduceMemoryUsage(memorySize);
+    }
+
     private void clear() {
+      if (dataSize > 0) {
+        block.reduceMemoryUsage(dataSize);
+        dataSize = 0;
+      }
+      nonDirectionalChunkData.clear();
       regionId2ReplicaSetAndNode.clear();
     }
   }

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

@@ -108,21 +108,22 @@ private void scan() throws IOException {
           FileUtils.streamFiles(listeningDirFile, true, (String[]) null)) {
         try {
           fileStream
+              .map(file -> new File(LoadUtil.getTsFilePath(file.getAbsolutePath())))
+              .distinct()
               .filter(file -> !activeLoadTsFileLoader.isFilePendingOrLoading(file))
               .filter(File::exists)
-              .map(file -> LoadUtil.getTsFilePath(file.getAbsolutePath()))
-              .filter(this::isTsFileCompleted)
+              .filter(file -> isTsFileCompleted(file.getAbsolutePath()))
               .limit(currentAllowedPendingSize)
               .forEach(
-                  filePath -> {
-                    final File tsFile = new File(filePath);
+                  tsFile -> {
                     final Map<String, String> attributes =
                         ActiveLoadPathHelper.parseAttributes(tsFile, listeningDirFile);
 
                     final File parentFile = tsFile.getParentFile();
                     final boolean isTableModel =
                         ActiveLoadPathHelper.containsDatabaseName(attributes)
-                            || (parentFile != null
+                            || (attributes.isEmpty()
+                                && parentFile != null
                                 && !Objects.equals(
                                     parentFile.getAbsoluteFile(),
                                     listeningDirFile.getAbsoluteFile()));

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

@@ -61,9 +61,9 @@ public synchronized ActiveLoadEntry dequeueFromPending() {
   }
 
   public synchronized void removeFromLoading(final String file) {
-    loadingFileSet.remove(file);
-
-    ActiveLoadingFilesNumberMetricsSet.getInstance().increaseLoadingFileCounter(-1);
+    if (loadingFileSet.remove(file)) {
+      ActiveLoadingFilesNumberMetricsSet.getInstance().increaseLoadingFileCounter(-1);
+    }
   }
 
   public synchronized boolean isFilePendingOrLoading(final String file) {
@@ -78,6 +78,20 @@ public synchronized boolean isEmpty() {
     return pendingFileQueue.isEmpty() && loadingFileSet.isEmpty();
   }
 
+  public synchronized void clear() {
+    final int loadingFileCount = loadingFileSet.size();
+    clearPending();
+    loadingFileSet.clear();
+    ActiveLoadingFilesNumberMetricsSet.getInstance().increaseLoadingFileCounter(-loadingFileCount);
+  }
+
+  public synchronized void clearPending() {
+    final int pendingFileCount = pendingFileSet.size();
+    pendingFileSet.clear();
+    pendingFileQueue.clear();
+    ActiveLoadingFilesNumberMetricsSet.getInstance().increaseQueuingFileCounter(-pendingFileCount);
+  }
+
   public static class ActiveLoadEntry {
     private final String file;
     private final String pendingDir;

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

@@ -160,12 +160,15 @@ private void adjustExecutorIfNecessary() {
   public void stop() {
     final WrappedThreadPoolExecutor executor = activeLoadExecutor.getAndSet(null);
     if (executor == null) {
+      pendingQueue.clearPending();
       return;
     }
 
-    executor.shutdownNow();
+    boolean isTerminated = false;
     try {
-      if (!executor.awaitTermination(30, TimeUnit.SECONDS)) {
+      executor.shutdownNow();
+      isTerminated = executor.awaitTermination(30, TimeUnit.SECONDS);
+      if (!isTerminated) {
         LOGGER.warn(
             StorageEngineMessages.STILL_NOT_EXIT_AFTER_30S,
             ThreadName.ACTIVE_LOAD_TSFILE_LOADER.getName());
@@ -175,6 +178,12 @@ public void stop() {
           StorageEngineMessages.STILL_NOT_EXIT_AFTER_30S,
           ThreadName.ACTIVE_LOAD_TSFILE_LOADER.getName());
       Thread.currentThread().interrupt();
+    } finally {
+      if (isTerminated) {
+        pendingQueue.clear();
+      } else {
+        pendingQueue.clearPending();
+      }
     }
   }
 
@@ -213,6 +222,7 @@ private void tryLoadPendingTsFiles() {
           handleOtherException(loadEntry.get(), e);
         } finally {
           pendingQueue.removeFromLoading(loadEntry.get().getFile());
+          cleanupEmptyDirectories(loadEntry.get());
         }
       }
     } finally {
@@ -364,6 +374,32 @@ private void removeToFailDir(final String filePath) {
     }
   }
 
+  private void cleanupEmptyDirectories(final ActiveLoadPendingQueue.ActiveLoadEntry entry) {
+    final File pendingDir =
+        entry.getPendingDir() == null
+            ? ActiveLoadPathHelper.findPendingDirectory(new File(entry.getFile()))
+            : new File(entry.getPendingDir());
+    if (pendingDir == null) {
+      return;
+    }
+
+    final Path pendingPath = pendingDir.toPath().toAbsolutePath().normalize();
+    Path currentPath = new File(entry.getFile()).toPath().toAbsolutePath().normalize().getParent();
+    while (currentPath != null
+        && currentPath.startsWith(pendingPath)
+        && !currentPath.equals(pendingPath)) {
+      try {
+        Files.delete(currentPath);
+      } catch (final IOException e) {
+        if (Files.exists(currentPath)) {
+          LOGGER.debug(StorageEngineMessages.FAILED_DELETE_FOLDER_CLEANING_UP, currentPath, e);
+        }
+        return;
+      }
+      currentPath = currentPath.getParent();
+    }
+  }
+
   public boolean isFilePendingOrLoading(final File file) {
     return pendingQueue.isFilePendingOrLoading(file.getAbsolutePath());
   }