[To rel/0.13][IOTDB-4645] Optimization of estimating memory usage of cross space compaction task (#7595)

Author: MrQuansy

server/src/main/java/org/apache/iotdb/db/engine/compaction/cross/rewrite/selector/AbstractCompactionEstimator.java

@@ -0,0 +1,64 @@
+/*
+ * 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.db.engine.compaction.cross.rewrite.selector;
+
+import org.apache.iotdb.db.engine.storagegroup.TsFileResource;
+import org.apache.iotdb.tsfile.read.TsFileSequenceReader;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Estimate the memory cost of one compaction task with specific source files based on its
+ * corresponding implementation.
+ */
+public abstract class AbstractCompactionEstimator implements AutoCloseable {
+
+  protected Map<TsFileResource, TsFileSequenceReader> fileReaderCache = new HashMap<>();
+
+  /**
+   * Estimate the memory cost of compacting the unseq file and its corresponding overlapped seq
+   * files in cross space compaction task.
+   */
+  public abstract long estimateCrossCompactionMemory(
+      List<TsFileResource> seqResources, TsFileResource unseqResource) throws IOException;
+
+  /** Estimate the memory cost of compacting the source files in inner space compaction task. */
+  public abstract long estimateInnerCompactionMemory(List<TsFileResource> resources);
+
+  /** Construct a new or get an existing TsFileSequenceReader of a TsFile. */
+  protected TsFileSequenceReader getFileReader(TsFileResource tsFileResource) throws IOException {
+    TsFileSequenceReader reader = fileReaderCache.get(tsFileResource);
+    if (reader == null) {
+      reader = new TsFileSequenceReader(tsFileResource.getTsFilePath(), true, false);
+      fileReaderCache.put(tsFileResource, reader);
+    }
+    return reader;
+  }
+
+  @Override
+  public void close() throws Exception {
+    for (TsFileSequenceReader sequenceReader : fileReaderCache.values()) {
+      sequenceReader.close();
+    }
+    fileReaderCache.clear();
+  }
+}

server/src/main/java/org/apache/iotdb/db/engine/compaction/cross/rewrite/selector/AbstractCrossSpaceEstimator.java

@@ -0,0 +1,38 @@
+/*
+ * 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.db.engine.compaction.cross.rewrite.selector;
+
+import org.apache.iotdb.db.engine.storagegroup.TsFileResource;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Estimate the memory cost of one cross space compaction task with specific source files based on
+ * its corresponding implementation.
+ */
+public abstract class AbstractCrossSpaceEstimator extends AbstractCompactionEstimator {
+  public abstract long estimateCrossCompactionMemory(
+      List<TsFileResource> seqResources, TsFileResource unseqResource) throws IOException;
+
+  public long estimateInnerCompactionMemory(List<TsFileResource> resources) {
+    throw new RuntimeException(
+        "This kind of estimator cannot be used to estimate inner space compaction task");
+  }
+}

server/src/main/java/org/apache/iotdb/db/engine/compaction/cross/rewrite/selector/RewriteCompactionFileSelector.java

@@ -24,7 +24,6 @@
 import org.apache.iotdb.db.engine.storagegroup.TsFileNameGenerator;
 import org.apache.iotdb.db.engine.storagegroup.TsFileResource;
 import org.apache.iotdb.db.exception.MergeException;
-import org.apache.iotdb.db.utils.MergeUtils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -73,13 +72,16 @@ public class RewriteCompactionFileSelector implements ICrossSpaceMergeFileSelect
   private boolean[] seqSelected;
   private int seqSelectedNum;
 
+  private AbstractCompactionEstimator compactionEstimator;
+
   public RewriteCompactionFileSelector(CrossSpaceCompactionResource resource, long memoryBudget) {
     this.resource = resource;
     this.memoryBudget = memoryBudget;
     this.maxCrossCompactionFileNum =
         IoTDBDescriptor.getInstance().getConfig().getMaxCrossCompactionCandidateFileNum();
     this.maxCrossCompactionFileSize =
         IoTDBDescriptor.getInstance().getConfig().getMaxCrossCompactionCandidateFileSize();
+    this.compactionEstimator = new RewriteCrossCompactionEstimator();
   }
 
   /**
@@ -113,10 +115,7 @@ public List[] select() throws MergeException {
           "Selecting merge candidates from {} seqFile, {} unseqFiles",
           resource.getSeqFiles().size(),
           resource.getUnseqFiles().size());
-      select(false);
-      if (selectedUnseqFiles.isEmpty()) {
-        select(true);
-      }
+      selectFiles();
       resource.setSeqFiles(selectedSeqFiles);
       resource.setUnseqFiles(selectedUnseqFiles);
       resource.removeOutdatedSeqReaders();
@@ -126,6 +125,12 @@ public List[] select() throws MergeException {
       }
     } catch (IOException e) {
       throw new MergeException(e);
+    } finally {
+      try {
+        compactionEstimator.close();
+      } catch (Exception e) {
+        throw new MergeException(e);
+      }
     }
     if (logger.isInfoEnabled()) {
       logger.info(
@@ -147,10 +152,9 @@ public List[] select() throws MergeException {
    * exceed the memory overhead preset by the system for the compaction thread, put them into the
    * selectedSeqFiles and selectedUnseqFiles.
    *
-   * @param useTightBound whether is tight estimate or loop estimate
    * @throws IOException
    */
-  void select(boolean useTightBound) throws IOException {
+  void selectFiles() throws IOException {
     tmpSelectedSeqFiles = new HashSet<>();
     seqSelected = new boolean[resource.getSeqFiles().size()];
     seqSelectedNum = 0;
@@ -189,16 +193,17 @@ void select(boolean useTightBound) throws IOException {
         }
       }
 
+      List<TsFileResource> tmpSelectedSeqFileResources = new ArrayList<>();
       for (int seqIndex : tmpSelectedSeqFiles) {
+        TsFileResource tsFileResource = resource.getSeqFiles().get(seqIndex);
+        tmpSelectedSeqFileResources.add(tsFileResource);
         totalSize += resource.getSeqFiles().get(seqIndex).getTsFileSize();
       }
       totalSize += unseqFile.getTsFileSize();
 
       tempMaxSeqFileCost = maxSeqFileCost;
       long newCost =
-          useTightBound
-              ? calculateTightMemoryCost(unseqFile, tmpSelectedSeqFiles, startTime, timeLimit)
-              : calculateLooseMemoryCost(unseqFile, tmpSelectedSeqFiles, startTime, timeLimit);
+          compactionEstimator.estimateCrossCompactionMemory(tmpSelectedSeqFileResources, unseqFile);
       if (!updateSelectedFiles(newCost, unseqFile)) {
         // older unseq files must be merged before newer ones
         break;
@@ -342,110 +347,6 @@ private void selectOverlappedSeqFiles(TsFileResource unseqFile) {
     }
   }
 
-  private long calculateMemoryCost(
-      TsFileResource tmpSelectedUnseqFile,
-      Collection<Integer> tmpSelectedSeqFiles,
-      IFileQueryMemMeasurement unseqMeasurement,
-      IFileQueryMemMeasurement seqMeasurement,
-      long startTime,
-      long timeLimit)
-      throws IOException {
-    long cost = 0;
-    Long fileCost = unseqMeasurement.measure(tmpSelectedUnseqFile);
-    cost += fileCost;
-
-    for (Integer seqFileIdx : tmpSelectedSeqFiles) {
-      TsFileResource seqFile = resource.getSeqFiles().get(seqFileIdx);
-      fileCost = seqMeasurement.measure(seqFile);
-      if (fileCost > tempMaxSeqFileCost) {
-        // only one file will be read at the same time, so only the largest one is recorded here
-        cost -= tempMaxSeqFileCost;
-        cost += fileCost;
-        tempMaxSeqFileCost = fileCost;
-      }
-      // but writing data into a new file may generate the same amount of metadata in memory
-      cost += calculateMetadataSize(seqFile);
-      long timeConsumption = System.currentTimeMillis() - startTime;
-      if (timeConsumption > timeLimit) {
-        return Long.MAX_VALUE;
-      }
-    }
-    return cost;
-  }
-
-  private long calculateLooseMemoryCost(
-      TsFileResource tmpSelectedUnseqFile,
-      Collection<Integer> tmpSelectedSeqFiles,
-      long startTime,
-      long timeLimit)
-      throws IOException {
-    return calculateMemoryCost(
-        tmpSelectedUnseqFile,
-        tmpSelectedSeqFiles,
-        TsFileResource::getTsFileSize,
-        this::calculateMetadataSize,
-        startTime,
-        timeLimit);
-  }
-
-  private long calculateTightMemoryCost(
-      TsFileResource tmpSelectedUnseqFile,
-      Collection<Integer> tmpSelectedSeqFiles,
-      long startTime,
-      long timeLimit)
-      throws IOException {
-    return calculateMemoryCost(
-        tmpSelectedUnseqFile,
-        tmpSelectedSeqFiles,
-        this::calculateTightUnseqMemoryCost,
-        this::calculateTightSeqMemoryCost,
-        startTime,
-        timeLimit);
-  }
-
-  private long calculateMetadataSize(TsFileResource seqFile) throws IOException {
-    Long cost = fileMetaSizeMap.get(seqFile);
-    if (cost == null) {
-      cost = MergeUtils.getFileMetaSize(seqFile, resource.getFileReader(seqFile));
-      fileMetaSizeMap.put(seqFile, cost);
-      logger.debug(LOG_FILE_COST, seqFile, cost);
-    }
-    return cost;
-  }
-
-  private long calculateTightFileMemoryCost(
-      TsFileResource seqFile, IFileQueryMemMeasurement measurement) throws IOException {
-    Long cost = maxSeriesQueryCostMap.get(seqFile);
-    if (cost == null) {
-      long[] chunkNums =
-          MergeUtils.findTotalAndLargestSeriesChunkNum(seqFile, resource.getFileReader(seqFile));
-      long totalChunkNum = chunkNums[0];
-      long maxChunkNum = chunkNums[1];
-      cost = measurement.measure(seqFile) * maxChunkNum / totalChunkNum;
-      maxSeriesQueryCostMap.put(seqFile, cost);
-      logger.debug(LOG_FILE_COST, seqFile, cost);
-    }
-    return cost;
-  }
-
-  // this method traverses all ChunkMetadata to find out which series has the most chunks and uses
-  // its proportion to all series to get a maximum estimation
-  private long calculateTightSeqMemoryCost(TsFileResource seqFile) throws IOException {
-    long singleSeriesCost = calculateTightFileMemoryCost(seqFile, this::calculateMetadataSize);
-    long multiSeriesCost = concurrentMergeNum * singleSeriesCost;
-    long maxCost = calculateMetadataSize(seqFile);
-    return Math.min(multiSeriesCost, maxCost);
-  }
-
-  // this method traverses all ChunkMetadata to find out which series has the most chunks and uses
-  // its proportion among all series to get a maximum estimation
-  private long calculateTightUnseqMemoryCost(TsFileResource unseqFile) throws IOException {
-    long singleSeriesCost = calculateTightFileMemoryCost(unseqFile, TsFileResource::getTsFileSize);
-    long multiSeriesCost = concurrentMergeNum * singleSeriesCost;
-    long maxCost = unseqFile.getTsFileSize();
-    return Math.min(multiSeriesCost, maxCost);
-  }
-
   @Override
   public int getConcurrentMergeNum() {
     return concurrentMergeNum;