add memory control when caching last values

Author: jt2594838

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

@@ -1158,6 +1158,8 @@ public class IoTDBConfig {
    */
   private boolean cacheLastValuesForLoad = true;
 
+  private long cacheLastValuesMemoryBudgetInByte = 4 * 1024 * 1024;
+
   IoTDBConfig() {}
 
   public int getMaxLogEntriesNumPerBatch() {
@@ -4096,4 +4098,12 @@ public boolean isCacheLastValuesForLoad() {
   public void setCacheLastValuesForLoad(boolean cacheLastValuesForLoad) {
     this.cacheLastValuesForLoad = cacheLastValuesForLoad;
   }
+
+  public long getCacheLastValuesMemoryBudgetInByte() {
+    return cacheLastValuesMemoryBudgetInByte;
+  }
+
+  public void setCacheLastValuesMemoryBudgetInByte(long cacheLastValuesMemoryBudgetInByte) {
+    this.cacheLastValuesMemoryBudgetInByte = cacheLastValuesMemoryBudgetInByte;
+  }
 }

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

@@ -2272,6 +2272,12 @@ private void loadLoadTsFileProps(TrimProperties properties) throws IOException {
         Boolean.parseBoolean(
             properties.getProperty(
                 "cache_last_values_for_load", String.valueOf(conf.isCacheLastValuesForLoad()))));
+
+    conf.setCacheLastValuesMemoryBudgetInByte(
+        Long.parseLong(
+            properties.getProperty(
+                "cache_last_values_memory_budget_in_byte",
+                String.valueOf(conf.getCacheLastValuesMemoryBudgetInByte()))));
   }
 
   private void loadLoadTsFileHotModifiedProp(TrimProperties properties) throws IOException {

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

@@ -49,6 +49,7 @@
 import org.apache.tsfile.read.reader.page.TimePageReader;
 import org.apache.tsfile.read.reader.page.ValuePageReader;
 import org.apache.tsfile.utils.Pair;
+import org.apache.tsfile.utils.RamUsageEstimator;
 import org.apache.tsfile.utils.TsPrimitiveType;
 import org.apache.tsfile.write.writer.TsFileIOWriter;
 import org.slf4j.Logger;
@@ -430,20 +431,21 @@ public static void updateTsFileResource(
             : tsFileResource.getTimeIndex();
     Map<IDeviceID, List<Pair<String, TimeValuePair>>> deviceLastValues =
         tsFileResource.getLastValues();
+    long lastValueMemCost = 0;
     if (cacheLastValue && deviceLastValues == null) {
       deviceLastValues = new HashMap<>(device2Metadata.size());
     }
     for (Map.Entry<IDeviceID, List<TimeseriesMetadata>> entry : device2Metadata.entrySet()) {
       List<Pair<String, TimeValuePair>> seriesLastValues = null;
-      if (cacheLastValue) {
+      if (deviceLastValues != null) {
         seriesLastValues = new ArrayList<>(entry.getValue().size());
       }
 
       for (TimeseriesMetadata timeseriesMetaData : entry.getValue()) {
         newTimeIndex.updateStartTime(
             entry.getKey(), timeseriesMetaData.getStatistics().getStartTime());
         newTimeIndex.updateEndTime(entry.getKey(), timeseriesMetaData.getStatistics().getEndTime());
-        if (cacheLastValue) {
+        if (deviceLastValues != null) {
           if (timeseriesMetaData.getTsDataType() != TSDataType.BLOB) {
             TsPrimitiveType value;
             value =
@@ -461,10 +463,33 @@ public static void updateTsFileResource(
           }
         }
       }
-      if (cacheLastValue) {
-        deviceLastValues
-            .computeIfAbsent(entry.getKey(), deviceID -> new ArrayList<>())
-            .addAll(seriesLastValues);
+      if (deviceLastValues != null) {
+        lastValueMemCost += entry.getKey().ramBytesUsed();
+        for (Pair<String, TimeValuePair> lastValue : seriesLastValues) {
+          if (lastValue == null) {
+            continue;
+          }
+          // pair
+          lastValueMemCost += RamUsageEstimator.shallowSizeOf(lastValue);
+          // measurement name
+          lastValueMemCost += RamUsageEstimator.sizeOf(lastValue.left);
+          TimeValuePair right = lastValue.getRight();
+          lastValueMemCost +=
+              right != null ? right.getSize() : RamUsageEstimator.NUM_BYTES_OBJECT_REF;
+        }
+        // ArrayList
+        lastValueMemCost +=
+            (long) seriesLastValues.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF
+                + RamUsageEstimator.NUM_BYTES_OBJECT_HEADER
+                + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
+        lastValueMemCost += RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY;
+        if (lastValueMemCost <= config.getCacheLastValuesMemoryBudgetInByte()) {
+          deviceLastValues
+              .computeIfAbsent(entry.getKey(), deviceID -> new ArrayList<>())
+              .addAll(seriesLastValues);
+        } else {
+          deviceLastValues = null;
+        }
       }
     }
     tsFileResource.setTimeIndex(newTimeIndex);

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

@@ -60,6 +60,7 @@
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.read.TimeValuePair;
 import org.apache.tsfile.utils.Pair;
+import org.apache.tsfile.utils.RamUsageEstimator;
 import org.apache.tsfile.utils.TsPrimitiveType;
 import org.apache.tsfile.write.writer.TsFileIOWriter;
 import org.slf4j.Logger;
@@ -78,6 +79,7 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -543,33 +545,60 @@ private void endTsFileResource(
       if (IoTDBDescriptor.getInstance().getConfig().isCacheLastValuesForLoad()) {
         deviceLastValues = new HashMap<>();
       }
+      AtomicLong lastValuesMemCost = new AtomicLong(0);
+
       for (final ChunkGroupMetadata chunkGroupMetadata : writer.getChunkGroupMetadataList()) {
         final IDeviceID device = chunkGroupMetadata.getDevice();
         for (final ChunkMetadata chunkMetadata : chunkGroupMetadata.getChunkMetadataList()) {
           tsFileResource.updateStartTime(device, chunkMetadata.getStartTime());
           tsFileResource.updateEndTime(device, chunkMetadata.getEndTime());
           if (deviceLastValues != null) {
-            deviceLastValues
-                .computeIfAbsent(device, d -> new HashMap<>())
-                .compute(
-                    chunkMetadata.getMeasurementUid(),
-                    (m, oldPair) -> {
-                      if (oldPair != null && oldPair.getTimestamp() > chunkMetadata.getEndTime()) {
-                        return oldPair;
-                      }
-                      TsPrimitiveType lastValue =
-                          chunkMetadata.getStatistics() != null
-                                  && chunkMetadata.getDataType() != TSDataType.BLOB
-                              ? TsPrimitiveType.getByType(
-                                  chunkMetadata.getDataType() == TSDataType.VECTOR
-                                      ? TSDataType.INT64
-                                      : chunkMetadata.getDataType(),
-                                  chunkMetadata.getStatistics().getLastValue())
-                              : null;
-                      return lastValue != null
+            Map<String, TimeValuePair> deviceMap =
+                deviceLastValues.computeIfAbsent(
+                    device,
+                    d -> {
+                      Map<String, TimeValuePair> map = new HashMap<>();
+                      lastValuesMemCost.addAndGet(RamUsageEstimator.shallowSizeOf(map));
+                      lastValuesMemCost.addAndGet(device.ramBytesUsed());
+                      return map;
+                    });
+            int prevSize = deviceMap.size();
+            deviceMap.compute(
+                chunkMetadata.getMeasurementUid(),
+                (m, oldPair) -> {
+                  if (oldPair != null && oldPair.getTimestamp() > chunkMetadata.getEndTime()) {
+                    return oldPair;
+                  }
+                  TsPrimitiveType lastValue =
+                      chunkMetadata.getStatistics() != null
+                              && chunkMetadata.getDataType() != TSDataType.BLOB
+                          ? TsPrimitiveType.getByType(
+                              chunkMetadata.getDataType() == TSDataType.VECTOR
+                                  ? TSDataType.INT64
+                                  : chunkMetadata.getDataType(),
+                              chunkMetadata.getStatistics().getLastValue())
+                          : null;
+                  TimeValuePair timeValuePair =
+                      lastValue != null
                           ? new TimeValuePair(chunkMetadata.getEndTime(), lastValue)
                           : null;
-                    });
+                  if (oldPair != null) {
+                    lastValuesMemCost.addAndGet(-oldPair.getSize());
+                  }
+                  if (timeValuePair != null) {
+                    lastValuesMemCost.addAndGet(timeValuePair.getSize());
+                  }
+                  return timeValuePair;
+                });
+            int afterSize = deviceMap.size();
+            lastValuesMemCost.addAndGet(
+                (afterSize - prevSize) * RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY);
+            if (lastValuesMemCost.get()
+                > IoTDBDescriptor.getInstance()
+                    .getConfig()
+                    .getCacheLastValuesMemoryBudgetInByte()) {
+              deviceLastValues = null;
+            }
           }
         }
       }