Load: Support storing piecenode in multiple folders during 2nd phase

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java

@@ -310,8 +310,9 @@ public class IoTDBConfig {
     {IoTDBConstant.DN_DEFAULT_DATA_DIR + File.separator + IoTDBConstant.DATA_FOLDER_NAME}
   };
 
-  private String loadTsFileDir =
-      tierDataDirs[0][0] + File.separator + IoTDBConstant.LOAD_TSFILE_FOLDER_NAME;
+  private String[] loadTsFileDirs = {
+    tierDataDirs[0][0] + File.separator + IoTDBConstant.LOAD_TSFILE_FOLDER_NAME
+  };
 
   /** Strategy of multiple directories. */
   private String multiDirStrategyClassName = null;
@@ -1276,7 +1277,6 @@ public void updatePath() {
   private void formulateFolders() {
     systemDir = addDataHomeDir(systemDir);
     schemaDir = addDataHomeDir(schemaDir);
-    loadTsFileDir = addDataHomeDir(loadTsFileDir);
     consensusDir = addDataHomeDir(consensusDir);
     dataRegionConsensusDir = addDataHomeDir(dataRegionConsensusDir);
     ratisDataRegionSnapshotDir = addDataHomeDir(ratisDataRegionSnapshotDir);
@@ -1324,6 +1324,7 @@ private void formulateDataDirs(String[][] tierDataDirs) {
         }
       }
     }
+    formulateLoadTsFileDirs(tierDataDirs);
   }
 
   void reloadDataDirs(String[][] tierDataDirs) throws LoadConfigurationException {
@@ -1427,7 +1428,6 @@ public void setTierDataDirs(String[][] tierDataDirs) {
     // TODO(szywilliam): rewrite the logic here when ratis supports complete snapshot semantic
     setRatisDataRegionSnapshotDir(
         tierDataDirs[0][0] + File.separator + IoTDBConstant.SNAPSHOT_FOLDER_NAME);
-    setLoadTsFileDir(tierDataDirs[0][0] + File.separator + IoTDBConstant.LOAD_TSFILE_FOLDER_NAME);
   }
 
   public String getRpcAddress() {
@@ -1462,12 +1462,16 @@ public void setSystemDir(String systemDir) {
     this.systemDir = systemDir;
   }
 
-  public String getLoadTsFileDir() {
-    return loadTsFileDir;
+  public String[] getLoadTsFileDirs() {
+    return this.loadTsFileDirs;
   }
 
-  public void setLoadTsFileDir(String loadTsFileDir) {
-    this.loadTsFileDir = loadTsFileDir;
+  public void formulateLoadTsFileDirs(String[][] tierDataDirs) {
+    loadTsFileDirs = new String[tierDataDirs.length];
+    for (int i = 0; i < tierDataDirs.length; i++) {
+      loadTsFileDirs[i] =
+          tierDataDirs[i][0] + File.separator + IoTDBConstant.LOAD_TSFILE_FOLDER_NAME;
+    }
   }
 
   public String getSchemaDir() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java

@@ -895,32 +895,6 @@ public void loadProperties(Properties properties) throws BadNodeUrlException, IO
       conf.setIntoOperationExecutionThreadCount(2);
     }
 
-    conf.setMaxAllocateMemoryRatioForLoad(
-        Double.parseDouble(
-            properties.getProperty(
-                "max_allocate_memory_ratio_for_load",
-                String.valueOf(conf.getMaxAllocateMemoryRatioForLoad()))));
-    conf.setLoadTsFileAnalyzeSchemaBatchFlushTimeSeriesNumber(
-        Integer.parseInt(
-            properties.getProperty(
-                "load_tsfile_analyze_schema_batch_flush_time_series_number",
-                String.valueOf(conf.getLoadTsFileAnalyzeSchemaBatchFlushTimeSeriesNumber()))));
-    conf.setLoadTsFileAnalyzeSchemaMemorySizeInBytes(
-        Long.parseLong(
-            properties.getProperty(
-                "load_tsfile_analyze_schema_memory_size_in_bytes",
-                String.valueOf(conf.getLoadTsFileAnalyzeSchemaMemorySizeInBytes()))));
-    conf.setLoadCleanupTaskExecutionDelayTimeSeconds(
-        Long.parseLong(
-            properties.getProperty(
-                "load_clean_up_task_execution_delay_time_seconds",
-                String.valueOf(conf.getLoadCleanupTaskExecutionDelayTimeSeconds()))));
-    conf.setLoadWriteThroughputBytesPerSecond(
-        Double.parseDouble(
-            properties.getProperty(
-                "load_write_throughput_bytes_per_second",
-                String.valueOf(conf.getLoadWriteThroughputBytesPerSecond()))));
-
     conf.setExtPipeDir(properties.getProperty("ext_pipe_dir", conf.getExtPipeDir()).trim());
 
     // At the same time, set TSFileConfig
@@ -1023,6 +997,9 @@ public void loadProperties(Properties properties) throws BadNodeUrlException, IO
     // CQ
     loadCQProps(properties);
 
+    // Load TsFile
+    loadLoadTsFileProps(properties);
+
     // Pipe
     loadPipeProps(properties);
 
@@ -2092,6 +2069,34 @@ private void initSchemaMemoryAllocate(Properties properties) {
     LOGGER.info("allocateMemoryForPartitionCache = {}", conf.getAllocateMemoryForPartitionCache());
   }
 
+  private void loadLoadTsFileProps(Properties properties) {
+    conf.setMaxAllocateMemoryRatioForLoad(
+        Double.parseDouble(
+            properties.getProperty(
+                "max_allocate_memory_ratio_for_load",
+                String.valueOf(conf.getMaxAllocateMemoryRatioForLoad()))));
+    conf.setLoadTsFileAnalyzeSchemaBatchFlushTimeSeriesNumber(
+        Integer.parseInt(
+            properties.getProperty(
+                "load_tsfile_analyze_schema_batch_flush_time_series_number",
+                String.valueOf(conf.getLoadTsFileAnalyzeSchemaBatchFlushTimeSeriesNumber()))));
+    conf.setLoadTsFileAnalyzeSchemaMemorySizeInBytes(
+        Long.parseLong(
+            properties.getProperty(
+                "load_tsfile_analyze_schema_memory_size_in_bytes",
+                String.valueOf(conf.getLoadTsFileAnalyzeSchemaMemorySizeInBytes()))));
+    conf.setLoadCleanupTaskExecutionDelayTimeSeconds(
+        Long.parseLong(
+            properties.getProperty(
+                "load_clean_up_task_execution_delay_time_seconds",
+                String.valueOf(conf.getLoadCleanupTaskExecutionDelayTimeSeconds()))));
+    conf.setLoadWriteThroughputBytesPerSecond(
+        Double.parseDouble(
+            properties.getProperty(
+                "load_write_throughput_bytes_per_second",
+                String.valueOf(conf.getLoadWriteThroughputBytesPerSecond()))));
+  }
+
   @SuppressWarnings("squid:S3518") // "proportionSum" can't be zero
   private void loadUDFProps(Properties properties) {
     String initialByteArrayLengthForMemoryControl =

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/load/LoadTsFileManager.java

@@ -29,6 +29,7 @@
 import org.apache.iotdb.commons.utils.FileUtils;
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.exception.DiskSpaceInsufficientException;
 import org.apache.iotdb.db.exception.LoadFileException;
 import org.apache.iotdb.db.pipe.agent.PipeAgent;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode;
@@ -38,6 +39,8 @@
 import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
 import org.apache.iotdb.db.storageengine.dataregion.utils.TsFileResourceUtils;
+import org.apache.iotdb.db.storageengine.rescon.disk.FolderManager;
+import org.apache.iotdb.db.storageengine.rescon.disk.strategy.DirectoryStrategyType;
 import org.apache.iotdb.metrics.utils.MetricLevel;
 
 import org.apache.tsfile.common.constant.TsFileConstant;
@@ -51,12 +54,16 @@
 import java.nio.file.DirectoryNotEmptyException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * {@link LoadTsFileManager} is used for dealing with {@link LoadTsFilePieceNode} and {@link
@@ -74,16 +81,16 @@ public class LoadTsFileManager {
       "%s TsFileWriterManager has been closed.";
   private static final String MESSAGE_DELETE_FAIL = "failed to delete {}.";
 
-  private final File loadDir;
+  private static final AtomicReference<String[]> LOAD_BASE_DIRS =
+      new AtomicReference<>(CONFIG.getLoadTsFileDirs());
+  private static final AtomicReference<FolderManager> FOLDER_MANAGER = new AtomicReference<>();
 
   private final Map<String, TsFileWriterManager> uuid2WriterManager = new ConcurrentHashMap<>();
 
   private final Map<String, CleanupTask> uuid2CleanupTask = new ConcurrentHashMap<>();
   private final PriorityBlockingQueue<CleanupTask> cleanupTaskQueue = new PriorityBlockingQueue<>();
 
   public LoadTsFileManager() {
-    this.loadDir = SystemFileFactory.INSTANCE.getFile(CONFIG.getLoadTsFileDir());
-
     registerCleanupTaskExecutor();
     recover();
   }
@@ -130,28 +137,43 @@ private void cleanupTasks() {
   }
 
   private void recover() {
-    if (!loadDir.exists()) {
-      return;
-    }
-
-    final File[] files = loadDir.listFiles();
-    if (files == null) {
-      return;
-    }
-    for (final File taskDir : files) {
-      String uuid = taskDir.getName();
-      TsFileWriterManager writerManager = new TsFileWriterManager(taskDir);
-
-      uuid2WriterManager.put(uuid, writerManager);
-      writerManager.close();
-
-      synchronized (uuid2CleanupTask) {
-        final CleanupTask cleanupTask =
-            new CleanupTask(uuid, CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000);
-        uuid2CleanupTask.put(uuid, cleanupTask);
-        cleanupTaskQueue.add(cleanupTask);
+    if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
+      synchronized (FOLDER_MANAGER) {
+        if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
+          LOAD_BASE_DIRS.set(CONFIG.getLoadTsFileDirs());
+        }
       }
     }
+
+    final File[] baseDirs = Arrays.stream(LOAD_BASE_DIRS.get()).map(File::new).toArray(File[]::new);
+    final File[] files =
+        Arrays.stream(baseDirs)
+            .filter(File::exists)
+            .flatMap(
+                dir -> {
+                  final File[] listedFiles = dir.listFiles();
+                  return listedFiles != null ? Arrays.stream(listedFiles) : Stream.empty();
+                })
+            .toArray(File[]::new);
+
+    Arrays.stream(files)
+        .parallel()
+        .forEach(
+            taskDir -> {
+              final String uuid = taskDir.getName();
+              final TsFileWriterManager writerManager = new TsFileWriterManager(taskDir);
+
+              uuid2WriterManager.put(uuid, writerManager);
+              writerManager.close();
+
+              synchronized (uuid2CleanupTask) {
+                final CleanupTask cleanupTask =
+                    new CleanupTask(
+                        uuid, CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000);
+                uuid2CleanupTask.put(uuid, cleanupTask);
+                cleanupTaskQueue.add(cleanupTask);
+              }
+            });
   }
 
   public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNode, String uuid)
@@ -168,10 +190,26 @@ public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNo
     final Optional<CleanupTask> cleanupTask = Optional.of(uuid2CleanupTask.get(uuid));
     cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning);
     try {
+      final AtomicReference<Exception> exception = new AtomicReference<>();
       final TsFileWriterManager writerManager =
           uuid2WriterManager.computeIfAbsent(
               uuid,
-              o -> new TsFileWriterManager(SystemFileFactory.INSTANCE.getFile(loadDir, uuid)));
+              o -> {
+                try {
+                  return new TsFileWriterManager(new File(getNextFolder(), uuid));
+                } catch (DiskSpaceInsufficientException e) {
+                  exception.set(e);
+                  return null;
+                }
+              });
+      if (exception.get() != null || writerManager == null) {
+        throw new IOException(
+            "Failed to create TsFileWriterManager for uuid "
+                + uuid
+                + " because of insufficient disk space.",
+            exception.get());
+      }
+
       for (TsFileData tsFileData : pieceNode.getAllTsFileData()) {
         if (!tsFileData.isModification()) {
           ChunkData chunkData = (ChunkData) tsFileData;
@@ -186,6 +224,33 @@ public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNo
     }
   }
 
+  private String getNextFolder() throws DiskSpaceInsufficientException {
+    if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
+      synchronized (FOLDER_MANAGER) {
+        if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
+          LOAD_BASE_DIRS.set(CONFIG.getLoadTsFileDirs());
+          FOLDER_MANAGER.set(
+              new FolderManager(
+                  Arrays.asList(LOAD_BASE_DIRS.get()), DirectoryStrategyType.SEQUENCE_STRATEGY));
+          return FOLDER_MANAGER.get().getNextFolder();
+        }
+      }
+    }
+
+    if (FOLDER_MANAGER.get() == null) {
+      synchronized (FOLDER_MANAGER) {
+        if (FOLDER_MANAGER.get() == null) {
+          FOLDER_MANAGER.set(
+              new FolderManager(
+                  Arrays.asList(LOAD_BASE_DIRS.get()), DirectoryStrategyType.SEQUENCE_STRATEGY));
+          return FOLDER_MANAGER.get().getNextFolder();
+        }
+      }
+    }
+
+    return FOLDER_MANAGER.get().getNextFolder();
+  }
+
   public boolean loadAll(String uuid, boolean isGeneratedByPipe, ProgressIndex progressIndex)
       throws IOException, LoadFileException {
     if (!uuid2WriterManager.containsKey(uuid)) {
@@ -229,17 +294,20 @@ private void forceCloseWriterManager(String uuid) {
       writerManager.close();
     }
 
-    final Path loadDirPath = loadDir.toPath();
-    if (!Files.exists(loadDirPath)) {
-      return;
-    }
-    try {
-      Files.deleteIfExists(loadDirPath);
-      LOGGER.info("Load dir {} was deleted.", loadDirPath);
-    } catch (DirectoryNotEmptyException e) {
-      LOGGER.info("Load dir {} is not empty, skip deleting.", loadDirPath);
-    } catch (IOException e) {
-      LOGGER.info(MESSAGE_DELETE_FAIL, loadDirPath);
+    for (final Path loadDirPath :
+        Arrays.stream(LOAD_BASE_DIRS.get())
+            .map(File::new)
+            .filter(File::exists)
+            .map(File::toPath)
+            .collect(Collectors.toList())) {
+      try {
+        Files.deleteIfExists(loadDirPath);
+        LOGGER.info("Load dir {} was deleted.", loadDirPath);
+      } catch (DirectoryNotEmptyException e) {
+        LOGGER.info("Load dir {} is not empty, skip deleting.", loadDirPath);
+      } catch (Exception e) {
+        LOGGER.info(MESSAGE_DELETE_FAIL, loadDirPath);
+      }
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/disk/FolderManager.java

@@ -79,4 +79,8 @@ public String getNextFolder() throws DiskSpaceInsufficientException {
       throw e;
     }
   }
+
+  public List<String> getFolders() {
+    return folders;
+  }
 }