Fix/min folder occupied space cache

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/directories/strategy/DirectoryStrategyTest.java

@@ -140,12 +140,47 @@ public void testMinFolderOccupiedSpaceFirstStrategy()
 
     PowerMockito.when(JVMCommonUtils.getOccupiedSpace(dataDirList.get(minIndex)))
         .thenReturn(Long.MAX_VALUE);
+    // The occupied space is cached, so the new value is reflected only after a cache refresh.
+    minFolderOccupiedSpaceFirstStrategy.invalidateCache();
     minIndex = getIndexOfMinOccupiedSpace();
     for (int i = 0; i < dataDirList.size(); i++) {
       assertEquals(minIndex, minFolderOccupiedSpaceFirstStrategy.nextFolderIndex());
     }
   }
 
+  @Test
+  public void testMinFolderOccupiedSpaceFirstStrategyCachesOccupiedSpace()
+      throws DiskSpaceInsufficientException, IOException {
+    MinFolderOccupiedSpaceFirstStrategy strategy = new MinFolderOccupiedSpaceFirstStrategy();
+    // Disable the time-based refresh so the count-based refresh is the only trigger under test.
+    strategy.setRefreshIntervalMs(Long.MAX_VALUE);
+    strategy.setRefreshSelectionThreshold(3);
+    strategy.setFolders(dataDirList);
+
+    int minIndex = getIndexOfMinOccupiedSpace();
+    // The first selection builds the cache via a single round of getOccupiedSpace calls.
+    assertEquals(minIndex, strategy.nextFolderIndex());
+    assertEquals(1, strategy.getSelectionsSinceRefresh());
+
+    // Mutate the occupied space of every folder so that a different folder becomes the least
+    // occupied one. While the cache is still valid the selection must not change, which proves the
+    // strategy is no longer re-walking the directory tree on every selection.
+    for (int i = 0; i < dataDirList.size(); i++) {
+      boolean available = !fullDirIndexSet.contains(i);
+      PowerMockito.when(JVMCommonUtils.getOccupiedSpace(dataDirList.get(i)))
+          .thenReturn(available ? (long) (dataDirList.size() - i) : Long.MAX_VALUE);
+    }
+    assertEquals(minIndex, strategy.nextFolderIndex());
+    assertEquals(minIndex, strategy.nextFolderIndex());
+    assertEquals(3, strategy.getSelectionsSinceRefresh());
+
+    // The threshold has been reached, so the next selection resets the state, recomputes the
+    // occupied space and reflects the mutated values.
+    int newMinIndex = getIndexOfMinOccupiedSpace();
+    assertEquals(newMinIndex, strategy.nextFolderIndex());
+    assertEquals(1, strategy.getSelectionsSinceRefresh());
+  }
+
   private int getIndexOfMinOccupiedSpace() throws IOException {
     int index = -1;
     long minOccupied = Long.MAX_VALUE;

iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template

@@ -765,6 +765,18 @@ topology_probing_timeout_ratio=0.5
 # Datatype: double(percentage)
 disk_space_warning_threshold=0.05
 
+# The refresh interval in milliseconds for the occupied-space cache used by
+# MinFolderOccupiedSpaceFirstStrategy. The default is 60000 (60 seconds).
+# effectiveMode: restart
+# Datatype: long
+min_folder_occupied_space_cache_refresh_interval_ms=60000
+
+# The number of folder selections after which MinFolderOccupiedSpaceFirstStrategy refreshes its
+# occupied-space cache. The default is 1000.
+# effectiveMode: restart
+# Datatype: int
+min_folder_occupied_space_cache_refresh_selection_threshold=1000
+
 # Purpose: for data partition repair
 # The number of threads used for parallel scanning in the partition table recovery
 # effectiveMode: restart

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java

@@ -185,6 +185,12 @@ public class CommonConfig {
   /** Disk Monitor. */
   private double diskSpaceWarningThreshold = 0.05;
 
+  /** Refresh interval for MinFolderOccupiedSpaceFirstStrategy occupied-space cache. */
+  private long minFolderOccupiedSpaceCacheRefreshIntervalMs = 60_000L;
+
+  /** Refresh selection threshold for MinFolderOccupiedSpaceFirstStrategy occupied-space cache. */
+  private int minFolderOccupiedSpaceCacheRefreshSelectionThreshold = 1000;
+
   /** Time partition origin in milliseconds. */
   private long timePartitionOrigin = 0;
 
@@ -781,6 +787,26 @@ public void setDiskSpaceWarningThreshold(double diskSpaceWarningThreshold) {
     this.diskSpaceWarningThreshold = diskSpaceWarningThreshold;
   }
 
+  public long getMinFolderOccupiedSpaceCacheRefreshIntervalMs() {
+    return minFolderOccupiedSpaceCacheRefreshIntervalMs;
+  }
+
+  public void setMinFolderOccupiedSpaceCacheRefreshIntervalMs(
+      long minFolderOccupiedSpaceCacheRefreshIntervalMs) {
+    this.minFolderOccupiedSpaceCacheRefreshIntervalMs =
+        minFolderOccupiedSpaceCacheRefreshIntervalMs;
+  }
+
+  public int getMinFolderOccupiedSpaceCacheRefreshSelectionThreshold() {
+    return minFolderOccupiedSpaceCacheRefreshSelectionThreshold;
+  }
+
+  public void setMinFolderOccupiedSpaceCacheRefreshSelectionThreshold(
+      int minFolderOccupiedSpaceCacheRefreshSelectionThreshold) {
+    this.minFolderOccupiedSpaceCacheRefreshSelectionThreshold =
+        minFolderOccupiedSpaceCacheRefreshSelectionThreshold;
+  }
+
   public boolean isReadOnly() {
     return status == NodeStatus.ReadOnly;
   }

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonDescriptor.java

@@ -230,6 +230,31 @@ public void loadCommonProps(TrimProperties properties) throws IOException {
                     String.valueOf(config.getDiskSpaceWarningThreshold()))
                 .trim()));
 
+    long minFolderOccupiedSpaceCacheRefreshIntervalMs =
+        Long.parseLong(
+            properties
+                .getProperty(
+                    "min_folder_occupied_space_cache_refresh_interval_ms",
+                    String.valueOf(config.getMinFolderOccupiedSpaceCacheRefreshIntervalMs()))
+                .trim());
+    if (minFolderOccupiedSpaceCacheRefreshIntervalMs > 0) {
+      config.setMinFolderOccupiedSpaceCacheRefreshIntervalMs(
+          minFolderOccupiedSpaceCacheRefreshIntervalMs);
+    }
+
+    int minFolderOccupiedSpaceCacheRefreshSelectionThreshold =
+        Integer.parseInt(
+            properties
+                .getProperty(
+                    "min_folder_occupied_space_cache_refresh_selection_threshold",
+                    String.valueOf(
+                        config.getMinFolderOccupiedSpaceCacheRefreshSelectionThreshold()))
+                .trim());
+    if (minFolderOccupiedSpaceCacheRefreshSelectionThreshold > 0) {
+      config.setMinFolderOccupiedSpaceCacheRefreshSelectionThreshold(
+          minFolderOccupiedSpaceCacheRefreshSelectionThreshold);
+    }
+
     config.setTimestampPrecision(
         properties.getProperty("timestamp_precision", config.getTimestampPrecision()).trim());
 

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/disk/strategy/MinFolderOccupiedSpaceFirstStrategy.java

@@ -18,20 +18,56 @@
  */
 package org.apache.iotdb.commons.disk.strategy;
 
+import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.exception.DiskSpaceInsufficientException;
 import org.apache.iotdb.commons.i18n.UtilMessages;
 import org.apache.iotdb.commons.utils.JVMCommonUtils;
+import org.apache.iotdb.commons.utils.TestOnly;
 
 import java.io.IOException;
 
+/**
+ * Selects the folder with the least occupied space.
+ *
+ * <p>Computing the occupied space of a folder requires a full {@code Files.walk} of its directory
+ * tree, which is very expensive when a folder holds a huge number of (small) files, e.g. while a
+ * snapshot consisting of hundreds of thousands of tiny files is being received. Re-walking on every
+ * selection turned the per-file cost into a full-tree scan and made the overall cost quadratic.
+ *
+ * <p>To avoid that, the occupied space of every folder is cached and only recomputed periodically.
+ * Between two refreshes an incremental selection counter is maintained; once enough folders have
+ * been selected (or enough time has elapsed) the cached state is reset and the occupied space is
+ * recomputed. This keeps the selection semantics (pick the least occupied folder) while bounding
+ * the number of full directory scans.
+ */
 public class MinFolderOccupiedSpaceFirstStrategy extends DirectoryStrategy {
 
+  private long refreshIntervalMs =
+      CommonDescriptor.getInstance().getConfig().getMinFolderOccupiedSpaceCacheRefreshIntervalMs();
+  private int refreshSelectionThreshold =
+      CommonDescriptor.getInstance()
+          .getConfig()
+          .getMinFolderOccupiedSpaceCacheRefreshSelectionThreshold();
+
+  /** Cached occupied space per folder, captured at the last refresh. */
+  private long[] cachedOccupiedSpace;
+
+  /** Incremental count of selections made since the last refresh. */
+  private int selectionsSinceRefresh;
+
+  /** Timestamp (ms) of the last refresh; a negative value means the cache must be (re)built. */
+  private long lastRefreshTimeMs = -1;
+
   @Override
   public int nextFolderIndex() throws DiskSpaceInsufficientException {
     return getMinOccupiedSpaceFolder();
   }
 
-  private int getMinOccupiedSpaceFolder() throws DiskSpaceInsufficientException {
+  private synchronized int getMinOccupiedSpaceFolder() throws DiskSpaceInsufficientException {
+    if (needRefresh()) {
+      refreshOccupiedSpace();
+    }
+
     int minIndex = -1;
     long minSpace = Long.MAX_VALUE;
 
@@ -43,14 +79,7 @@ private int getMinOccupiedSpaceFolder() throws DiskSpaceInsufficientException {
       if (!JVMCommonUtils.hasSpace(folder)) {
         continue;
       }
-
-      long space = 0;
-      try {
-        space = JVMCommonUtils.getOccupiedSpace(folder);
-      } catch (IOException e) {
-        LOGGER.error(UtilMessages.CANNOT_CALCULATE_OCCUPIED_SPACE, folder, e);
-        continue;
-      }
+      long space = cachedOccupiedSpace[i];
       if (space < minSpace) {
         minSpace = space;
         minIndex = i;
@@ -61,6 +90,61 @@ private int getMinOccupiedSpaceFolder() throws DiskSpaceInsufficientException {
       throw new DiskSpaceInsufficientException(folders);
     }
 
+    selectionsSinceRefresh++;
     return minIndex;
   }
+
+  private boolean needRefresh() {
+    if (cachedOccupiedSpace == null || cachedOccupiedSpace.length != folders.size()) {
+      return true;
+    }
+    if (lastRefreshTimeMs < 0 || selectionsSinceRefresh >= refreshSelectionThreshold) {
+      return true;
+    }
+    return System.currentTimeMillis() - lastRefreshTimeMs >= refreshIntervalMs;
+  }
+
+  /** Recompute the occupied space of every folder and reset the incremental state. */
+  private void refreshOccupiedSpace() {
+    if (cachedOccupiedSpace == null || cachedOccupiedSpace.length != folders.size()) {
+      cachedOccupiedSpace = new long[folders.size()];
+    }
+    for (int i = 0; i < folders.size(); i++) {
+      String folder = folders.get(i);
+      if (isUnavailableFolder(folder) || !JVMCommonUtils.hasSpace(folder)) {
+        // Folder is not a selection candidate; keep it deprioritized without paying for a walk.
+        cachedOccupiedSpace[i] = Long.MAX_VALUE;
+        continue;
+      }
+      try {
+        cachedOccupiedSpace[i] = JVMCommonUtils.getOccupiedSpace(folder);
+      } catch (IOException e) {
+        LOGGER.error(UtilMessages.CANNOT_CALCULATE_OCCUPIED_SPACE, folder, e);
+        cachedOccupiedSpace[i] = Long.MAX_VALUE;
+      }
+    }
+    selectionsSinceRefresh = 0;
+    lastRefreshTimeMs = System.currentTimeMillis();
+  }
+
+  @TestOnly
+  public void setRefreshIntervalMs(long refreshIntervalMs) {
+    this.refreshIntervalMs = refreshIntervalMs;
+  }
+
+  @TestOnly
+  public void setRefreshSelectionThreshold(int refreshSelectionThreshold) {
+    this.refreshSelectionThreshold = refreshSelectionThreshold;
+  }
+
+  /** Forces the next selection to recompute the occupied space of every folder. */
+  @TestOnly
+  public synchronized void invalidateCache() {
+    this.lastRefreshTimeMs = -1;
+  }
+
+  @TestOnly
+  public synchronized int getSelectionsSinceRefresh() {
+    return selectionsSinceRefresh;
+  }
 }