From 6c15634b1e3880165f8d42ff3cebe7e5284f864c Mon Sep 17 00:00:00 2001
From: shuwenwei <55970239+shuwenwei@users.noreply.github.com>
Date: Fri, 26 Jan 2024 19:00:37 +0800
Subject: [PATCH] Fix estimating compaction memory without acquire read lock
 (#11972)

---
 .../apache/iotdb/db/it/IoTDBRepairDataIT.java |  2 ++
 .../AbstractCompactionEstimator.java          | 10 ++++++++++
 .../AbstractCrossSpaceEstimator.java          |  3 +++
 .../AbstractInnerSpaceEstimator.java          | 19 ++++++++++++++++---
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBRepairDataIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBRepairDataIT.java
index 5241531aec08..78038802dcd4 100644
--- a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBRepairDataIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBRepairDataIT.java
@@ -44,6 +44,7 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.iotdb.db.utils.constant.TestConstant.TIMESTAMP_STR;
 import static org.junit.Assert.assertNotNull;
@@ -81,6 +82,7 @@ public void testRepairData() {
         if (sorted) {
           return;
         }
+        Thread.sleep(TimeUnit.SECONDS.toMillis(1));
       }
       Assert.fail();
     } catch (Exception e) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCompactionEstimator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCompactionEstimator.java
index 2211ce53e12d..77517edbc08f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCompactionEstimator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCompactionEstimator.java
@@ -23,6 +23,7 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.storageengine.dataregion.flush.CompressionRatio;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
+import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.DeviceTimeIndex;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.FileTimeIndex;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.ITimeIndex;
@@ -63,6 +64,15 @@ public abstract class AbstractCompactionEstimator {
 
   protected abstract long calculatingDataMemoryCost(CompactionTaskInfo taskInfo) throws IOException;
 
+  protected boolean isAllSourceFileExist(List<TsFileResource> resources) {
+    for (TsFileResource resource : resources) {
+      if (resource.getStatus() == TsFileResourceStatus.DELETED) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   protected CompactionTaskInfo calculatingCompactionTaskInfo(List<TsFileResource> resources)
       throws IOException {
     List<FileInfo> fileInfoList = new ArrayList<>();
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCrossSpaceEstimator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCrossSpaceEstimator.java
index 0c072f4002f6..2c5a6acd2b57 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCrossSpaceEstimator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractCrossSpaceEstimator.java
@@ -45,6 +45,9 @@ public long estimateCrossCompactionMemory(
 
     long cost = 0;
     try {
+      if (!isAllSourceFileExist(resources)) {
+        return -1L;
+      }
       CompactionTaskInfo taskInfo = calculatingCompactionTaskInfo(resources);
       cost += calculatingMetadataMemoryCost(taskInfo);
       cost += calculatingDataMemoryCost(taskInfo);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractInnerSpaceEstimator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractInnerSpaceEstimator.java
index df28012bd127..ae2c6f2e016b 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractInnerSpaceEstimator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/selector/estimator/AbstractInnerSpaceEstimator.java
@@ -34,9 +34,22 @@ public long estimateInnerCompactionMemory(List<TsFileResource> resources) throws
     if (!config.isEnableCompactionMemControl()) {
       return 0;
     }
-    CompactionTaskInfo taskInfo = calculatingCompactionTaskInfo(resources);
-    long cost = calculatingMetadataMemoryCost(taskInfo);
-    cost += calculatingDataMemoryCost(taskInfo);
+
+    if (!CompactionEstimateUtils.addReadLock(resources)) {
+      return -1L;
+    }
+    long cost = 0;
+    try {
+      if (!isAllSourceFileExist(resources)) {
+        return -1L;
+      }
+
+      CompactionTaskInfo taskInfo = calculatingCompactionTaskInfo(resources);
+      cost = calculatingMetadataMemoryCost(taskInfo);
+      cost += calculatingDataMemoryCost(taskInfo);
+    } finally {
+      CompactionEstimateUtils.releaseReadLock(resources);
+    }
     return cost;
   }
 }