Size PartitionInfo snapshot buffer from the file size instead of a fixed 32MB

PartitionInfo always allocated a fixed 32MB buffer for snapshot load and
save, regardless of the actual partition table size. On edge deployments
the default heap is only 224MB and snapshots are taken frequently, so a
small partition table can still trigger large, short-lived allocations.

Size the write buffer from the last snapshot file size and the read
buffer from the current file size, still bounded at 32MB so large
partition tables keep the previous behavior.

Author: CRZbulabula

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

@@ -119,8 +119,17 @@ 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;
+  /** Upper bound of a snapshot stream buffer, preserving the previous 32MB cap for large tables. */
+  private static final int PARTITION_TABLE_BUFFER_SIZE_MAX = 32 * 1024 * 1024;
+
+  /** Fallback buffer size used when no snapshot size estimate is available yet. */
+  private static final int DEFAULT_SNAPSHOT_BUFFER_SIZE = 8 * 1024;
+
+  /**
+   * Size of the last written or loaded snapshot file, used to size the buffer of the next snapshot
+   * adaptively instead of always allocating a fixed 32MB buffer.
+   */
+  private volatile long lastSnapshotSize = -1L;
 
   /** For Cluster Partition. */
   // For allocating Regions
@@ -978,6 +987,19 @@ public Map<TSeriesPartitionSlot, TConsensusGroupId> getLastDataAllotTable(String
     return Collections.emptyMap();
   }
 
+  /**
+   * Compute the buffer size for a snapshot file of {@code fileSize} bytes: the actual file size
+   * when known, bounded by {@link #PARTITION_TABLE_BUFFER_SIZE_MAX} so that snapshot I/O never
+   * allocates more than 32MB, falling back to {@link #DEFAULT_SNAPSHOT_BUFFER_SIZE} when no size
+   * estimate is available.
+   */
+  static int getSnapshotBufferSize(final long fileSize) {
+    if (fileSize <= 0) {
+      return DEFAULT_SNAPSHOT_BUFFER_SIZE;
+    }
+    return (int) Math.min(fileSize, PARTITION_TABLE_BUFFER_SIZE_MAX);
+  }
+
   @Override
   public boolean processTakeSnapshot(File snapshotDir) throws TException, IOException {
 
@@ -993,9 +1015,11 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
     // snapshot operation.
     File tmpFile = new File(snapshotFile.getAbsolutePath() + "-" + UUID.randomUUID());
 
+    // Size the write buffer from the last snapshot size, 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);
+            new BufferedOutputStream(fileOutputStream, getSnapshotBufferSize(lastSnapshotSize));
         TIOStreamTransport tioStreamTransport = new TIOStreamTransport(bufferedOutputStream)) {
       TProtocol protocol = new TBinaryProtocol(tioStreamTransport);
 
@@ -1024,7 +1048,12 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
       tioStreamTransport.close();
 
       // rename file
-      return tmpFile.renameTo(snapshotFile);
+      final boolean renamed = tmpFile.renameTo(snapshotFile);
+      if (renamed) {
+        // Remember the file size to size the buffer of the next snapshot.
+        lastSnapshotSize = snapshotFile.length();
+      }
+      return renamed;
     } finally {
       // with or without success, delete temporary files anyway
       for (int retry = 0; retry < 5; retry++) {
@@ -1050,9 +1079,12 @@ public void processLoadSnapshot(final File snapshotDir) throws TException, IOExc
       return;
     }
 
+    // Size the read buffer from the file size, so loading a small snapshot no longer allocates a
+    // fixed 32MB buffer.
     try (final BufferedInputStream fileInputStream =
             new BufferedInputStream(
-                Files.newInputStream(snapshotFile.toPath()), PARTITION_TABLE_BUFFER_SIZE);
+                Files.newInputStream(snapshotFile.toPath()),
+                getSnapshotBufferSize(snapshotFile.length()));
         final TIOStreamTransport tioStreamTransport = new TIOStreamTransport(fileInputStream)) {
       final TProtocol protocol = new TBinaryProtocol(tioStreamTransport);
       // before restoring a snapshot, clear all old data
@@ -1082,6 +1114,9 @@ public void processLoadSnapshot(final File snapshotDir) throws TException, IOExc
         regionMaintainTaskList.add(task);
       }
     }
+
+    // Remember the loaded file size to size the buffer of the next snapshot.
+    lastSnapshotSize = snapshotFile.length();
   }
 
   /**

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

@@ -0,0 +1,41 @@
+/*
+ * 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.persistence.partition;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class PartitionInfoBufferSizeTest {
+
+  @Test
+  public void testGetSnapshotBufferSize() {
+    // No size estimate yet: fall back to the default 8KB buffer.
+    Assert.assertEquals(PartitionInfo.getSnapshotBufferSize(-1L), 8 * 1024);
+    Assert.assertEquals(PartitionInfo.getSnapshotBufferSize(0L), 8 * 1024);
+
+    // A known file size is used directly while it stays below the 32MB cap.
+    Assert.assertEquals(PartitionInfo.getSnapshotBufferSize(1L), 1);
+    Assert.assertEquals(PartitionInfo.getSnapshotBufferSize(4 * 1024L), 4 * 1024);
+
+    // At and above the cap the buffer is bounded at 32MB.
+    Assert.assertEquals(PartitionInfo.getSnapshotBufferSize(32L * 1024 * 1024), 32 * 1024 * 1024);
+    Assert.assertEquals(PartitionInfo.getSnapshotBufferSize(64L * 1024 * 1024), 32 * 1024 * 1024);
+  }
+}