Refactor write memory accounting into named memory blocks (#18848)

* Refactor write memory accounting into named memory blocks

* Address write memory accounting review feedback

Author: hongzhi-gao

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

@@ -1670,6 +1670,15 @@ private StorageEngineMessages() {}
   public static final String MESSAGE_THE_ASSOCIATED_RESOURCE_FILE_OF_ARG_IS_NOT_FOUND_IN_THE_SNAPSHOT_CB9152B5 = "The associated resource file of {} is not found in the snapshot";
   public static final String MESSAGE_EVICTED_NON_EXISTING_EXISTING_SERIES_COUNT_ARG_ARG_ARG_TOTAL_REQUEST_ARG_3026ADBD = "Evicted non-existing/existing series count: {}/{}({}), total request: {}";
 
+  public static final String EXCEPTION_WRITE_MEMORY_BLOCK_ARG_IS_CLOSED_9E2561DF =
+      "Write memory block %s is closed";
+  public static final String EXCEPTION_WRITE_MEMORY_SIZE_MUST_BE_NON_NEGATIVE_ARG_2CBCF358 =
+      "Write memory size must be non-negative: %s";
+  public static final String EXCEPTION_CANNOT_RELEASE_ARG_BYTES_FROM_WRITE_MEMORY_BLOCK_ARG_WITH_ARG_BYTES_3783A614 =
+      "Cannot release %s bytes from write memory block %s with %s bytes";
+  public static final String EXCEPTION_WRITE_MEMORY_BLOCK_ARG_ALREADY_BELONGS_TO_ANOTHER_MANAGER_A05A8229 =
+      "Write memory block %s already belongs to another manager";
+
   public static final String EXCEPTION_CHUNK_LAYOUT_IS_MISSING_E87C71C0 = "Chunk layout is missing";
   public static final String EXCEPTION_CHUNK_HEADER_LIST_IS_EMPTY_5DE68C78 =
       "Chunk header list is empty.";

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java

@@ -1591,6 +1591,15 @@ private StorageEngineMessages() {}
   public static final String MESSAGE_THE_ASSOCIATED_RESOURCE_FILE_OF_ARG_IS_NOT_FOUND_IN_THE_SNAPSHOT_CB9152B5 = "在快照中未找到 {} 关联的资源文件";
   public static final String MESSAGE_EVICTED_NON_EXISTING_EXISTING_SERIES_COUNT_ARG_ARG_ARG_TOTAL_REQUEST_ARG_3026ADBD = "淘汰的不存在/已存在时间序列数量：{}/{}({})，总请求数：{}";
 
+  public static final String EXCEPTION_WRITE_MEMORY_BLOCK_ARG_IS_CLOSED_9E2561DF =
+      "写入内存块 %s 已关闭";
+  public static final String EXCEPTION_WRITE_MEMORY_SIZE_MUST_BE_NON_NEGATIVE_ARG_2CBCF358 =
+      "写入内存大小不能为负数：%s";
+  public static final String EXCEPTION_CANNOT_RELEASE_ARG_BYTES_FROM_WRITE_MEMORY_BLOCK_ARG_WITH_ARG_BYTES_3783A614 =
+      "无法释放 %s 字节：写入内存块 %s 当前占用 %s 字节";
+  public static final String EXCEPTION_WRITE_MEMORY_BLOCK_ARG_ALREADY_BELONGS_TO_ANOTHER_MANAGER_A05A8229 =
+      "写入内存块 %s 已属于其他管理器";
+
   public static final String EXCEPTION_CHUNK_LAYOUT_IS_MISSING_E87C71C0 = "缺少 Chunk layout";
   public static final String EXCEPTION_CHUNK_HEADER_LIST_IS_EMPTY_5DE68C78 = "Chunk header 列表为空。";
   public static final String EXCEPTION_FAILED_TO_ENCODE_CHUNK_USING_ALIGNEDCHUNKWRITERIMPL_5AAF9AA8 =

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java

@@ -372,7 +372,7 @@ public class DataRegion implements IDataRegionForQuery {
               .build();
 
   /** database info for mem control. */
-  private final DataRegionInfo dataRegionInfo = new DataRegionInfo(this);
+  private final DataRegionInfo dataRegionInfo;
 
   /** whether it's ready from recovery. */
   private boolean isReady = false;
@@ -431,6 +431,7 @@ public DataRegion(
     this.dataRegionIdString = dataRegionIdString;
     this.dataRegionId = new DataRegionId(Integer.parseInt(dataRegionIdString));
     this.databaseName = databaseName;
+    this.dataRegionInfo = new DataRegionInfo(this);
     this.fileFlushPolicy = fileFlushPolicy;
     this.isTableModel = isTableModelDatabase(databaseName);
     this.delayAnalyzer =
@@ -512,6 +513,7 @@ public DataRegion(String databaseName, String dataRegionIdString) {
     this.isTableModel = isTableModelDatabase(databaseName);
     this.dataRegionIdString = dataRegionIdString;
     this.dataRegionId = new DataRegionId(Integer.parseInt(this.dataRegionIdString));
+    this.dataRegionInfo = new DataRegionInfo(this);
     this.tsFileManager = new TsFileManager(databaseName, dataRegionIdString, "");
     this.partitionMaxFileVersions = new HashMap<>();
     partitionMaxFileVersions.put(0L, 0L);
@@ -1073,7 +1075,9 @@ private void recoverUnsealedTsFileCallBack(UnsealedTsFileRecoverPerformer recove
         tsFileResource.removeResourceFile();
         tsFileProcessor.setTimeRangeId(timePartitionId);
         writer.makeMetadataVisible();
-        TsFileProcessorInfo tsFileProcessorInfo = new TsFileProcessorInfo(dataRegionInfo);
+        TsFileProcessorInfo tsFileProcessorInfo =
+            new TsFileProcessorInfo(
+                dataRegionInfo, tsFileProcessor.getTsFileResource().getTsFilePath());
         tsFileProcessor.setTsFileProcessorInfo(tsFileProcessorInfo);
         this.dataRegionInfo.initTsFileProcessorInfo(tsFileProcessor);
         // get chunkMetadata size
@@ -2322,7 +2326,9 @@ private TsFileProcessor getTsFileProcessor(
             this::flushCallback,
             sequence);
 
-    TsFileProcessorInfo tsFileProcessorInfo = new TsFileProcessorInfo(dataRegionInfo);
+    TsFileProcessorInfo tsFileProcessorInfo =
+        new TsFileProcessorInfo(
+            dataRegionInfo, tsFileProcessor.getTsFileResource().getTsFilePath());
     tsFileProcessor.setTsFileProcessorInfo(tsFileProcessorInfo);
     this.dataRegionInfo.initTsFileProcessorInfo(tsFileProcessor);
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java

@@ -21,6 +21,7 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor;
 import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo;
+import org.apache.iotdb.db.storageengine.rescon.memory.WriteMemoryManager;
 
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -32,11 +33,7 @@ public class DataRegionInfo {
 
   private final DataRegion dataRegion;
 
-  /**
-   * The total database memory cost, including unsealed TsFileResource, ChunkMetadata, WAL,
-   * primitive arrays and TEXT values
-   */
-  private final AtomicLong memoryCost;
+  private final WriteMemoryManager writeMemoryManager;
 
   /** The threshold of reporting it's size to SystemInfo */
   private final long storageGroupSizeReportThreshold =
@@ -57,7 +54,13 @@ public class DataRegionInfo {
 
   public DataRegionInfo(DataRegion dataRegion) {
     this.dataRegion = dataRegion;
-    memoryCost = new AtomicLong();
+    this.writeMemoryManager =
+        new WriteMemoryManager(
+            "DataRegion["
+                + dataRegion.getDatabaseName()
+                + "/"
+                + dataRegion.getDataRegionIdString()
+                + "]");
   }
 
   public DataRegion getDataRegion() {
@@ -69,16 +72,12 @@ public void initTsFileProcessorInfo(TsFileProcessor tsFileProcessor) {
     reportedTsps.add(tsFileProcessor);
   }
 
-  public void addStorageGroupMemCost(long cost) {
-    memoryCost.getAndAdd(cost);
-  }
-
-  public void releaseStorageGroupMemCost(long cost) {
-    memoryCost.getAndAdd(-cost);
+  public WriteMemoryManager getWriteMemoryManager() {
+    return writeMemoryManager;
   }
 
   public long getMemCost() {
-    return memoryCost.get();
+    return writeMemoryManager.getMemoryUsageInBytes();
   }
 
   public List<TsFileProcessor> getAllReportedTsp() {
@@ -87,7 +86,7 @@ public List<TsFileProcessor> getAllReportedTsp() {
 
   public boolean needToReportToSystem() {
     boolean needToReport =
-        memoryCost.get() - lastReportedSize.get() > storageGroupSizeReportThreshold
+        getMemCost() - lastReportedSize.get() > storageGroupSizeReportThreshold
             || needToReportToSystem.get();
     // report once and then reset flag to false
     if (needToReportToSystem.get()) {