Load: 内存控制 ver0.1

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

@@ -174,8 +174,6 @@ public class IoTDBConfig {
   /** The proportion of write memory for loading TsFile */
   private double loadTsFileProportion = 0.125;
 
-  private int maxLoadingTimeseriesNumber = 2000;
-
   /**
    * If memory cost of data region increased more than proportion of {@linkplain
    * IoTDBConfig#getAllocateMemoryForStorageEngine()}*{@linkplain
@@ -1089,6 +1087,16 @@ public class IoTDBConfig {
   private int maxPendingBatchesNum = 5;
   private double maxMemoryRatioForQueue = 0.6;
 
+  /** Load related */
+  private int maxLoadingTimeseriesNumber = 2000;
+
+  // TODO: remove loadMemoryTotalSizeFromQueryInBytes after introducing queryEngine memory manager
+  private long loadMemoryTotalSizeFromQueryInBytes = 200 * 1024 * 1024L;
+  private long initLoadMemoryTotalSizeInBytes = 100 * 1024 * 1024L; // 100MB
+  private long loadMemoryAllocateRetryIntervalMs = 1000;
+  private int loadMemoryAllocateMaxRetries = 10;
+  private long loadMemoryAllocateMinSizeInBytes = 32;
+
   /** Pipe related */
   /** initialized as empty, updated based on the latest `systemDir` during querying */
   private String[] pipeReceiverFileDirs = new String[0];
@@ -3269,14 +3277,6 @@ public double getLoadTsFileProportion() {
     return loadTsFileProportion;
   }
 
-  public int getMaxLoadingTimeseriesNumber() {
-    return maxLoadingTimeseriesNumber;
-  }
-
-  public void setMaxLoadingTimeseriesNumber(int maxLoadingTimeseriesNumber) {
-    this.maxLoadingTimeseriesNumber = maxLoadingTimeseriesNumber;
-  }
-
   public static String getEnvironmentVariables() {
     return "\n\t"
         + IoTDBConstant.IOTDB_HOME
@@ -3731,6 +3731,54 @@ public int getModeMapSizeThreshold() {
     return modeMapSizeThreshold;
   }
 
+  public long getLoadMemoryTotalSizeFromQueryInBytes() {
+    return loadMemoryTotalSizeFromQueryInBytes;
+  }
+
+  public void setLoadMemoryTotalSizeFromQueryInBytes(long loadMemoryTotalSizeFromQueryInBytes) {
+    this.loadMemoryTotalSizeFromQueryInBytes = loadMemoryTotalSizeFromQueryInBytes;
+  }
+
+  public long getInitLoadMemoryTotalSizeInBytes() {
+    return initLoadMemoryTotalSizeInBytes;
+  }
+
+  public void setInitLoadMemoryTotalSizeInBytes(long initLoadMemoryTotalSizeInBytes) {
+    this.initLoadMemoryTotalSizeInBytes = initLoadMemoryTotalSizeInBytes;
+  }
+
+  public int getMaxLoadingTimeseriesNumber() {
+    return maxLoadingTimeseriesNumber;
+  }
+
+  public void setMaxLoadingTimeseriesNumber(int maxLoadingTimeseriesNumber) {
+    this.maxLoadingTimeseriesNumber = maxLoadingTimeseriesNumber;
+  }
+
+  public long getLoadMemoryAllocateRetryIntervalMs() {
+    return loadMemoryAllocateRetryIntervalMs;
+  }
+
+  public void setLoadMemoryAllocateRetryIntervalMs(long loadMemoryAllocateRetryIntervalMs) {
+    this.loadMemoryAllocateRetryIntervalMs = loadMemoryAllocateRetryIntervalMs;
+  }
+
+  public int getLoadMemoryAllocateMaxRetries() {
+    return loadMemoryAllocateMaxRetries;
+  }
+
+  public void setLoadMemoryAllocateMaxRetries(int loadMemoryAllocateMaxRetries) {
+    this.loadMemoryAllocateMaxRetries = loadMemoryAllocateMaxRetries;
+  }
+
+  public long getLoadMemoryAllocateMinSizeInBytes() {
+    return loadMemoryAllocateMinSizeInBytes;
+  }
+
+  public void setLoadMemoryAllocateMinSizeInBytes(long loadMemoryAllocateMinSizeInBytes) {
+    this.loadMemoryAllocateMinSizeInBytes = loadMemoryAllocateMinSizeInBytes;
+  }
+
   public void setPipeReceiverFileDirs(String[] pipeReceiverFileDirs) {
     this.pipeReceiverFileDirs = pipeReceiverFileDirs;
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/load/analyze/LoadTsfileAnalyzer.java

@@ -45,6 +45,7 @@
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.common.schematree.DeviceSchemaInfo;
 import org.apache.iotdb.db.queryengine.common.schematree.ISchemaTree;
+import org.apache.iotdb.db.queryengine.load.memory.DeviceToTimeseriesSchemasMap;
 import org.apache.iotdb.db.queryengine.plan.Coordinator;
 import org.apache.iotdb.db.queryengine.plan.analyze.Analysis;
 import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
@@ -177,6 +178,7 @@ public Analysis analyzeFileByFile() {
     }
 
     LOGGER.info("Load - Analysis Stage: all tsfiles have been analyzed.");
+    schemaAutoCreatorAndVerifier.close();
 
     // data partition will be queried in the scheduler
     final Analysis analysis = new Analysis();
@@ -252,8 +254,8 @@ private Analysis createFailAnalysisForAuthException(AuthException e) {
   private final class SchemaAutoCreatorAndVerifier {
 
     private final Map<String, Boolean> tsfileDevice2IsAligned = new HashMap<>();
-    private final Map<String, Set<MeasurementSchema>> currentBatchDevice2TimeseriesSchemas =
-        new HashMap<>();
+    private final DeviceToTimeseriesSchemasMap currentBatchDevice2TimeseriesSchemas =
+        new DeviceToTimeseriesSchemasMap();
 
     private final Set<PartialPath> alreadySetDatabases = new HashSet<>();
 
@@ -263,6 +265,7 @@ public void autoCreateAndVerify(
         TsFileSequenceReader reader,
         Map<String, List<TimeseriesMetadata>> device2TimeseriesMetadataList)
         throws IOException, AuthException {
+      boolean isReachedMaxMemorySize = false;
       for (final Map.Entry<String, List<TimeseriesMetadata>> entry :
           device2TimeseriesMetadataList.entrySet()) {
         final String device = entry.getKey();
@@ -304,9 +307,9 @@ public void autoCreateAndVerify(
             }
             final Pair<CompressionType, TSEncoding> compressionEncodingPair =
                 reader.readTimeseriesCompressionTypeAndEncoding(timeseriesMetadata);
-            currentBatchDevice2TimeseriesSchemas
-                .computeIfAbsent(device, o -> new HashSet<>())
-                .add(
+            isReachedMaxMemorySize =
+                currentBatchDevice2TimeseriesSchemas.add(
+                    device,
                     new MeasurementSchema(
                         timeseriesMetadata.getMeasurementId(),
                         dataType,
@@ -315,6 +318,11 @@ public void autoCreateAndVerify(
 
             tsfileDevice2IsAligned.putIfAbsent(device, false);
           }
+
+          if (isReachedMaxMemorySize) {
+            flush();
+            isReachedMaxMemorySize = false;
+          }
         }
       }
 
@@ -616,5 +624,11 @@ public void clear() {
       currentBatchDevice2TimeseriesSchemas.clear();
       alreadySetDatabases.clear();
     }
+
+    public void close() {
+      tsfileDevice2IsAligned.clear();
+      currentBatchDevice2TimeseriesSchemas.close();
+      alreadySetDatabases.clear();
+    }
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/load/exception/LoadRuntimeOutOfMemoryException.java

@@ -0,0 +1,26 @@
+/*
+ * 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.queryengine.load.exception;
+
+public class LoadRuntimeOutOfMemoryException extends RuntimeException {
+  public LoadRuntimeOutOfMemoryException(String message) {
+    super(message);
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/load/memory/DeviceToTimeseriesSchemasMap.java

@@ -0,0 +1,89 @@
+/*
+ * 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.queryengine.load.memory;
+
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.queryengine.load.memory.block.LoadMemoryBlock;
+import org.apache.iotdb.tsfile.write.schema.MeasurementSchema;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public class DeviceToTimeseriesSchemasMap {
+  private final LoadTsFileMemoryManager loadTsFileMemoryManager =
+      LoadTsFileMemoryManager.getInstance();
+
+  private final LoadMemoryBlock block;
+
+  private long totalMemorySizeInBytes =
+      IoTDBDescriptor.getInstance().getConfig().getInitLoadMemoryTotalSizeInBytes() / 10;
+  private long usedMemorySizeInBytes = 0;
+
+  private final Map<String, Set<MeasurementSchema>> currentBatchDevice2TimeseriesSchemas =
+      new HashMap<>();
+
+  public DeviceToTimeseriesSchemasMap() {
+    block = loadTsFileMemoryManager.forceAllocate(totalMemorySizeInBytes);
+  }
+
+  // false if the memory is still enough and the schema is added successfully, true when it is out
+  // of memory.
+  public boolean add(String device, MeasurementSchema measurementSchema) {
+    boolean isNewDevice = !currentBatchDevice2TimeseriesSchemas.containsKey(device);
+
+    currentBatchDevice2TimeseriesSchemas
+        .computeIfAbsent(device, k -> new HashSet<>())
+        .add(measurementSchema);
+
+    if (isNewDevice) {
+      // estimate the memory size of a new hashmap node and a hash set
+      usedMemorySizeInBytes += (100);
+    }
+    usedMemorySizeInBytes += measurementSchema.serializedSize();
+
+    return usedMemorySizeInBytes > totalMemorySizeInBytes;
+  }
+
+  public void clear() {
+    currentBatchDevice2TimeseriesSchemas.clear();
+    usedMemorySizeInBytes = 0;
+  }
+
+  public boolean isEmpty() {
+    return currentBatchDevice2TimeseriesSchemas.isEmpty();
+  }
+
+  public Set<Map.Entry<String, Set<MeasurementSchema>>> entrySet() {
+    return currentBatchDevice2TimeseriesSchemas.entrySet();
+  }
+
+  public Set<String> keySet() {
+    return currentBatchDevice2TimeseriesSchemas.keySet();
+  }
+
+  public void close() {
+    currentBatchDevice2TimeseriesSchemas.clear();
+    totalMemorySizeInBytes = 0;
+    usedMemorySizeInBytes = 0;
+    loadTsFileMemoryManager.release(block);
+  }
+}