Bound ConfigNode snapshot buffers and avoid whole-file heap reads

- Add SnapshotStreamFactory providing adaptive, capped and reusable
  buffered streams for ConfigNode snapshot files, configured via the new
  config_node_snapshot_buffer_size_max parameter (default 4MB)
- PartitionInfo: replace the fixed 32MB snapshot buffer with the adaptive
  one, sized from the previous snapshot file size
- TemplateTable / CNPhysicalPlanGenerator: memory-map template snapshot
  files instead of copying them into the heap
- Edge distribution: cap snapshot buffers at 256KB

Author: CRZbulabula

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeConfig.java

@@ -245,6 +245,12 @@ public class ConfigNodeConfig {
   private long configNodeRatisConsensusLogAppenderBufferSize = 16 * 1024 * 1024L;
   private long schemaRegionRatisConsensusLogAppenderBufferSize = 16 * 1024 * 1024L;
 
+  /**
+   * Upper bound (in bytes) of the in-memory buffer used when taking/loading ConfigNode snapshot
+   * files. Snapshot streams size their buffers adaptively up to this cap; 0 disables buffering.
+   */
+  private long configNodeSnapshotBufferSizeMax = 4 * 1024 * 1024L;
+
   /**
    * RatisConsensus protocol, trigger a snapshot when ratis_snapshot_trigger_threshold logs are
    * written.
@@ -973,6 +979,14 @@ public void setSchemaRegionRatisConsensusLogAppenderBufferSize(
         schemaRegionRatisConsensusLogAppenderBufferSize;
   }
 
+  public long getConfigNodeSnapshotBufferSizeMax() {
+    return configNodeSnapshotBufferSizeMax;
+  }
+
+  public void setConfigNodeSnapshotBufferSizeMax(long configNodeSnapshotBufferSizeMax) {
+    this.configNodeSnapshotBufferSizeMax = configNodeSnapshotBufferSizeMax;
+  }
+
   public long getSchemaRegionRatisSnapshotTriggerThreshold() {
     return schemaRegionRatisSnapshotTriggerThreshold;
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeDescriptor.java

@@ -28,6 +28,7 @@
 import org.apache.iotdb.commons.log.LoggerPeriodicalLogReducer;
 import org.apache.iotdb.commons.pipe.config.PipeDescriptor;
 import org.apache.iotdb.commons.schema.SchemaConstant;
+import org.apache.iotdb.commons.snapshot.SnapshotStreamFactory;
 import org.apache.iotdb.commons.utils.NodeUrlUtils;
 import org.apache.iotdb.confignode.i18n.ConfigNodeMessages;
 import org.apache.iotdb.confignode.manager.load.balancer.RegionBalancer;
@@ -403,6 +404,13 @@ private void loadProperties(TrimProperties properties) throws BadNodeUrlExceptio
 
     loadRatisConsensusConfig(properties);
     loadCQConfig(properties);
+
+    conf.setConfigNodeSnapshotBufferSizeMax(
+        Long.parseLong(
+            properties.getProperty(
+                "config_node_snapshot_buffer_size_max",
+                String.valueOf(conf.getConfigNodeSnapshotBufferSizeMax()))));
+    SnapshotStreamFactory.setBufferSizeMax(conf.getConfigNodeSnapshotBufferSizeMax());
   }
 
   private void loadRatisConsensusConfig(TrimProperties properties) {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java

@@ -30,6 +30,7 @@
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
 import org.apache.iotdb.commons.schema.table.Audit;
 import org.apache.iotdb.commons.snapshot.SnapshotProcessor;
+import org.apache.iotdb.commons.snapshot.SnapshotStreamFactory;
 import org.apache.iotdb.commons.utils.PathUtils;
 import org.apache.iotdb.confignode.consensus.request.read.partition.CountTimeSlotListPlan;
 import org.apache.iotdb.confignode.consensus.request.read.partition.GetDataPartitionPlan;
@@ -80,11 +81,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.BufferedInputStream;
-import java.io.BufferedOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.BitSet;
@@ -119,8 +120,8 @@ public class PartitionInfo implements SnapshotProcessor {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(PartitionInfo.class);
 
-  // Allocate 8MB buffer for load snapshot of PartitionInfo
-  private static final int PARTITION_TABLE_BUFFER_SIZE = 32 * 1024 * 1024;
+  /** Size of the last written snapshot file, used to size the next snapshot buffer adaptively. */
+  private volatile long lastSnapshotSize = -1L;
 
   /** For Cluster Partition. */
   // For allocating Regions
@@ -993,9 +994,13 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
     // snapshot operation.
     File tmpFile = new File(snapshotFile.getAbsolutePath() + "-" + UUID.randomUUID());
 
+    // The write buffer is sized adaptively from the previous snapshot size and capped by
+    // config_node_snapshot_buffer_size_max, so a small partition table no longer allocates a
+    // fixed 32MB buffer per snapshot.
     try (FileOutputStream fileOutputStream = new FileOutputStream(tmpFile);
-        BufferedOutputStream bufferedOutputStream =
-            new BufferedOutputStream(fileOutputStream, PARTITION_TABLE_BUFFER_SIZE);
+        OutputStream bufferedOutputStream =
+            SnapshotStreamFactory.createOutputStream(
+                fileOutputStream, lastSnapshotSize > 0 ? lastSnapshotSize : 0);
         TIOStreamTransport tioStreamTransport = new TIOStreamTransport(bufferedOutputStream)) {
       TProtocol protocol = new TBinaryProtocol(tioStreamTransport);
 
@@ -1024,7 +1029,11 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
       tioStreamTransport.close();
 
       // rename file
-      return tmpFile.renameTo(snapshotFile);
+      final boolean renamed = tmpFile.renameTo(snapshotFile);
+      if (renamed) {
+        lastSnapshotSize = snapshotFile.length();
+      }
+      return renamed;
     } finally {
       // with or without success, delete temporary files anyway
       for (int retry = 0; retry < 5; retry++) {
@@ -1050,9 +1059,12 @@ public void processLoadSnapshot(final File snapshotDir) throws TException, IOExc
       return;
     }
 
-    try (final BufferedInputStream fileInputStream =
-            new BufferedInputStream(
-                Files.newInputStream(snapshotFile.toPath()), PARTITION_TABLE_BUFFER_SIZE);
+    // The read buffer is sized from the file size and capped by
+    // config_node_snapshot_buffer_size_max,
+    // so loading a snapshot never allocates more than the configured cap.
+    try (final InputStream fileInputStream =
+            SnapshotStreamFactory.createInputStream(
+                Files.newInputStream(snapshotFile.toPath()), snapshotFile.length());
         final TIOStreamTransport tioStreamTransport = new TIOStreamTransport(fileInputStream)) {
       final TProtocol protocol = new TBinaryProtocol(tioStreamTransport);
       // before restoring a snapshot, clear all old data
@@ -1082,6 +1094,8 @@ public void processLoadSnapshot(final File snapshotDir) throws TException, IOExc
         regionMaintainTaskList.add(task);
       }
     }
+    // Remember the loaded file size as the estimate for the next snapshot's write buffer.
+    lastSnapshotSize = snapshotFile.length();
   }
 
   /**

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/CNPhysicalPlanGenerator.java

@@ -47,7 +47,6 @@
 import org.apache.iotdb.confignode.persistence.schema.mnode.impl.ConfigTableNode;
 import org.apache.iotdb.confignode.rpc.thrift.TDatabaseSchema;
 
-import org.apache.tsfile.external.commons.io.IOUtils;
 import org.apache.tsfile.utils.Pair;
 import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.slf4j.Logger;
@@ -60,8 +59,10 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -89,7 +90,9 @@ public class CNPhysicalPlanGenerator
 
   // File input stream.
   private InputStream inputStream = null;
-  private InputStream templateInputStream = null;
+
+  // Template snapshot file, memory-mapped on demand instead of being read into the heap.
+  private Path templateFile = null;
 
   private static final String STRING_ENCODING = "utf-8";
 
@@ -126,7 +129,7 @@ public CNPhysicalPlanGenerator(final Path schemaInfoFile, final Path templateFil
     inputStream = Files.newInputStream(schemaInfoFile);
     // Template file is null for table mTree
     if (Objects.nonNull(templateFile)) {
-      templateInputStream = Files.newInputStream(templateFile);
+      this.templateFile = templateFile;
     }
     snapshotFileType = CNSnapshotFileType.SCHEMA;
   }
@@ -153,7 +156,7 @@ public boolean hasNext() {
     } else if (snapshotFileType == CNSnapshotFileType.TTL) {
       generateSetTTLPlan();
     } else if (snapshotFileType == CNSnapshotFileType.SCHEMA) {
-      if (Objects.nonNull(templateInputStream)) {
+      if (Objects.nonNull(templateFile)) {
         generateTemplatePlan();
       }
       if (latestException != null) {
@@ -171,10 +174,6 @@ public boolean hasNext() {
         inputStream.close();
         inputStream = null;
       }
-      if (templateInputStream != null) {
-        templateInputStream.close();
-        templateInputStream = null;
-      }
     } catch (final IOException ioException) {
       latestException = ioException;
     }
@@ -513,9 +512,14 @@ private void generateDatabasePhysicalPlan() {
   }
 
   private void generateTemplatePlan() {
-    try (final BufferedInputStream bufferedInputStream =
-        new BufferedInputStream(templateInputStream)) {
-      final ByteBuffer byteBuffer = ByteBuffer.wrap(IOUtils.toByteArray(bufferedInputStream));
+    // The template snapshot file is memory-mapped instead of being read into the heap, so
+    // generating plans from a large template file no longer temporarily doubles its memory
+    // footprint.
+    try (final FileChannel channel = FileChannel.open(templateFile, StandardOpenOption.READ)) {
+      if (channel.size() == 0) {
+        return;
+      }
+      final ByteBuffer byteBuffer = channel.map(FileChannel.MapMode.READ_ONLY, 0, channel.size());
       // Skip id
       ReadWriteIOUtils.readInt(byteBuffer);
       int size = ReadWriteIOUtils.readInt(byteBuffer);

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplateTable.java

@@ -28,23 +28,21 @@
 
 import org.apache.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.enums.TSDataType;
-import org.apache.tsfile.external.commons.io.IOUtils;
 import org.apache.tsfile.file.metadata.enums.CompressionType;
 import org.apache.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -194,8 +192,7 @@ private void serializeTemplate(Template template, OutputStream outputStream) thr
     template.serialize(outputStream);
   }
 
-  private void deserialize(InputStream inputStream) throws IOException {
-    ByteBuffer byteBuffer = ByteBuffer.wrap(IOUtils.toByteArray(inputStream));
+  private void deserialize(ByteBuffer byteBuffer) {
     templateIdGenerator.set(ReadWriteIOUtils.readInt(byteBuffer));
     int size = ReadWriteIOUtils.readInt(byteBuffer);
     while (size > 0) {
@@ -259,12 +256,14 @@ public void processLoadSnapshot(File snapshotDir) throws IOException {
       return;
     }
     templateReadWriteLock.writeLock().lock();
-    try (FileInputStream fileInputStream = new FileInputStream(snapshotFile);
-        BufferedInputStream bufferedInputStream = new BufferedInputStream(fileInputStream)) {
+    // The snapshot file is memory-mapped instead of copied into the heap, so loading a large
+    // template file no longer temporarily doubles its memory footprint.
+    try (final FileChannel channel =
+        FileChannel.open(snapshotFile.toPath(), StandardOpenOption.READ)) {
       // Load snapshot of template
       this.templateMap.clear();
       this.templateIdMap.clear();
-      deserialize(bufferedInputStream);
+      deserialize(channel.map(FileChannel.MapMode.READ_ONLY, 0, channel.size()));
     } finally {
       templateReadWriteLock.writeLock().unlock();
     }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/conf/ConfigNodeConfigTest.java

@@ -0,0 +1,40 @@
+/*
+ * 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.confignode.conf;
+
+import org.apache.iotdb.commons.snapshot.SnapshotStreamFactory;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class ConfigNodeConfigTest {
+
+  @Test
+  public void testSnapshotBufferSizeMaxDefault() {
+    final ConfigNodeConfig configNodeConfig = new ConfigNodeConfig();
+    // The code default must stay aligned with SnapshotStreamFactory's default cap.
+    Assert.assertEquals(
+        SnapshotStreamFactory.DEFAULT_BUFFER_SIZE_MAX,
+        configNodeConfig.getConfigNodeSnapshotBufferSizeMax());
+
+    configNodeConfig.setConfigNodeSnapshotBufferSizeMax(256 * 1024L);
+    Assert.assertEquals(256 * 1024L, configNodeConfig.getConfigNodeSnapshotBufferSizeMax());
+  }
+}

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/CNPhysicalPlanGeneratorTest.java

@@ -391,6 +391,37 @@ public void databaseWithoutTemplateGeneratorTest() throws Exception {
     Assert.assertEquals(11, count);
   }
 
+  @Test
+  public void emptyTemplateFileGeneratorTest() throws Exception {
+    setupClusterSchemaInfo();
+    final TDatabaseSchema tDatabaseSchema = new TDatabaseSchema();
+    tDatabaseSchema.setName("root.sg");
+    final DatabaseSchemaPlan databaseSchemaPlan =
+        new DatabaseSchemaPlan(ConfigPhysicalPlanType.CreateDatabase, tDatabaseSchema);
+    clusterSchemaInfo.createDatabase(databaseSchemaPlan);
+    Assert.assertTrue(clusterSchemaInfo.processTakeSnapshot(snapshotDir));
+
+    final File schemaInfo =
+        SystemFileFactory.INSTANCE.getFile(snapshotDir + File.separator + SCHEMA_INFO_FILE_NAME);
+    // A zero-byte template file must be tolerated: no template plans are generated and no
+    // exception is reported.
+    final File emptyTemplate =
+        SystemFileFactory.INSTANCE.getFile(
+            snapshotDir + File.separator + "empty_" + TEMPLATE_INFO_FILE_NAME);
+    Assert.assertTrue(emptyTemplate.createNewFile());
+
+    final CNPhysicalPlanGenerator planGenerator =
+        new CNPhysicalPlanGenerator(schemaInfo.toPath(), emptyTemplate.toPath());
+    int templateCount = 0;
+    for (ConfigPhysicalPlan plan : planGenerator) {
+      if (plan.getType() == ConfigPhysicalPlanType.CreateSchemaTemplate) {
+        templateCount++;
+      }
+    }
+    planGenerator.checkException();
+    Assert.assertEquals(0, templateCount);
+  }
+
   @Test
   public void templateGeneratorTest() throws Exception {
     setupClusterSchemaInfo();