This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new fd398988672 Bound PartitionInfo snapshot buffers via
config_node_snapshot_buffer_size_max (#18580)
fd398988672 is described below
commit fd398988672d84d12163103b251a0cea6891a07e
Author: Yongzao <[email protected]>
AuthorDate: Fri Sep 4 16:07:05 2026 +0800
Bound PartitionInfo snapshot buffers via
config_node_snapshot_buffer_size_max (#18580)
---
.../iotdb/confignode/conf/ConfigNodeConfig.java | 20 ++
.../confignode/conf/ConfigNodeDescriptor.java | 8 +
.../persistence/partition/PartitionInfo.java | 23 +-
.../confignode/conf/ConfigNodeConfigTest.java | 44 +++
.../confignode/persistence/PartitionInfoTest.java | 65 +++++
.../resources/conf/edge/iotdb-system.properties | 5 +
.../conf/iotdb-system.properties.template | 5 +
.../apache/iotdb/commons/i18n/CommonMessages.java | 3 +
.../apache/iotdb/commons/i18n/CommonMessages.java | 3 +
.../snapshot/ReusableBufferedOutputStream.java | 134 +++++++++
.../commons/snapshot/SnapshotStreamFactory.java | 130 +++++++++
.../snapshot/SnapshotStreamFactoryTest.java | 307 +++++++++++++++++++++
12 files changed, 737 insertions(+), 10 deletions(-)
diff --git
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeConfig.java
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeConfig.java
index f49525ea7da..c1baa98209c 100644
---
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeConfig.java
+++
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeConfig.java
@@ -26,6 +26,7 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import
org.apache.iotdb.commons.client.property.ClientPoolProperty.DefaultProperty;
import org.apache.iotdb.commons.conf.CommonDescriptor;
import org.apache.iotdb.commons.conf.IoTDBConstant;
+import org.apache.iotdb.commons.i18n.CommonMessages;
import org.apache.iotdb.confignode.i18n.ConfigNodeMessages;
import org.apache.iotdb.confignode.manager.load.balancer.RegionBalancer;
import
org.apache.iotdb.confignode.manager.load.balancer.router.leader.AbstractLeaderBalancer;
@@ -245,6 +246,9 @@ public class ConfigNodeConfig {
private long configNodeRatisConsensusLogAppenderBufferSize = 16 * 1024 *
1024L;
private long schemaRegionRatisConsensusLogAppenderBufferSize = 16 * 1024 *
1024L;
+ /** Max size (in bytes) of the in-memory buffer used when taking/loading
ConfigNode snapshots. */
+ private long configNodeSnapshotBufferSizeMax = 4 * 1024 * 1024L;
+
/**
* RatisConsensus protocol, trigger a snapshot when
ratis_snapshot_trigger_threshold logs are
* written.
@@ -973,6 +977,22 @@ public class ConfigNodeConfig {
schemaRegionRatisConsensusLogAppenderBufferSize;
}
+ public long getConfigNodeSnapshotBufferSizeMax() {
+ return configNodeSnapshotBufferSizeMax;
+ }
+
+ public void setConfigNodeSnapshotBufferSizeMax(long
configNodeSnapshotBufferSizeMax) {
+ if (configNodeSnapshotBufferSizeMax > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException(
+ String.format(
+ CommonMessages
+
.EXCEPTION_SNAPSHOT_BUFFER_SIZE_MUST_NOT_EXCEED_ARG_BYTES_BUT_WAS_ARG_D1DA6F7E,
+ Integer.MAX_VALUE,
+ configNodeSnapshotBufferSizeMax));
+ }
+ this.configNodeSnapshotBufferSizeMax = configNodeSnapshotBufferSizeMax;
+ }
+
public long getSchemaRegionRatisSnapshotTriggerThreshold() {
return schemaRegionRatisSnapshotTriggerThreshold;
}
diff --git
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeDescriptor.java
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeDescriptor.java
index 16dfd4150f1..b4e8b6f9d89 100644
---
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeDescriptor.java
+++
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeDescriptor.java
@@ -28,6 +28,7 @@ import org.apache.iotdb.commons.exception.BadNodeUrlException;
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 @@ public class ConfigNodeDescriptor {
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) {
diff --git
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java
index 8b2079b581d..71fef7c21d0 100644
---
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java
+++
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java
@@ -30,6 +30,7 @@ import org.apache.iotdb.commons.partition.DataPartitionTable;
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.apache.tsfile.utils.ReadWriteIOUtils;
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,9 +120,6 @@ 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;
-
/** For Cluster Partition. */
// For allocating Regions
private final AtomicInteger nextRegionGroupId;
@@ -993,9 +991,11 @@ public class PartitionInfo implements SnapshotProcessor {
// snapshot operation.
File tmpFile = new File(snapshotFile.getAbsolutePath() + "-" +
UUID.randomUUID());
+ // The write buffer is bounded 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);
TIOStreamTransport tioStreamTransport = new
TIOStreamTransport(bufferedOutputStream)) {
TProtocol protocol = new TBinaryProtocol(tioStreamTransport);
@@ -1050,9 +1050,12 @@ public class PartitionInfo implements SnapshotProcessor {
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
diff --git
a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/conf/ConfigNodeConfigTest.java
b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/conf/ConfigNodeConfigTest.java
new file mode 100644
index 00000000000..da016c5e163
--- /dev/null
+++
b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/conf/ConfigNodeConfigTest.java
@@ -0,0 +1,44 @@
+/*
+ * 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());
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> configNodeConfig.setConfigNodeSnapshotBufferSizeMax((long)
Integer.MAX_VALUE + 1));
+ Assert.assertEquals(256 * 1024L,
configNodeConfig.getConfigNodeSnapshotBufferSizeMax());
+ }
+}
diff --git
a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/PartitionInfoTest.java
b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/PartitionInfoTest.java
index afccb0c0eba..21744963066 100644
---
a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/PartitionInfoTest.java
+++
b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/PartitionInfoTest.java
@@ -29,6 +29,7 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
import org.apache.iotdb.commons.partition.DataPartitionTable;
import org.apache.iotdb.commons.partition.SchemaPartitionTable;
import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.snapshot.SnapshotStreamFactory;
import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType;
import
org.apache.iotdb.confignode.consensus.request.read.region.GetRegionInfoListPlan;
import
org.apache.iotdb.confignode.consensus.request.write.database.DatabaseSchemaPlan;
@@ -66,6 +67,8 @@ public class PartitionInfoTest {
private static PartitionInfo partitionInfo;
private static final File snapshotDir = new File(BASE_OUTPUT_PATH,
"snapshot");
+ private long originalSnapshotBufferSizeMax;
+
public enum testFlag {
DataPartition(20),
SchemaPartition(30);
@@ -87,6 +90,10 @@ public class PartitionInfoTest {
if (!snapshotDir.exists()) {
snapshotDir.mkdirs();
}
+ // Run the snapshot round-trips of this class with a small buffer cap,
proving that snapshot
+ // correctness does not depend on a large fixed buffer.
+ originalSnapshotBufferSizeMax = SnapshotStreamFactory.getBufferSizeMax();
+ SnapshotStreamFactory.setBufferSizeMax(64 * 1024);
}
@After
@@ -95,6 +102,7 @@ public class PartitionInfoTest {
if (snapshotDir.exists()) {
FileUtils.deleteDirectory(snapshotDir);
}
+ SnapshotStreamFactory.setBufferSizeMax(originalSnapshotBufferSizeMax);
}
@Test
@@ -152,6 +160,63 @@ public class PartitionInfoTest {
Assert.assertEquals(partitionInfo, partitionInfo1);
}
+ @Test
+ public void testSnapshotWithWriteBufferSmallerThanSnapshot() throws
TException, IOException {
+ partitionInfo.generateNextRegionGroupId();
+
+ // Set StorageGroup
+ partitionInfo.createDatabase(
+ new DatabaseSchemaPlan(
+ ConfigPhysicalPlanType.CreateDatabase, new
TDatabaseSchema("root.test")));
+
+ // Create a SchemaRegion
+ CreateRegionGroupsPlan createRegionGroupsReq = new
CreateRegionGroupsPlan();
+ final TRegionReplicaSet schemaRegionReplicaSet =
+ generateTRegionReplicaSet(
+ testFlag.SchemaPartition.getFlag(),
+ generateTConsensusGroupId(
+ testFlag.SchemaPartition.getFlag(),
TConsensusGroupType.SchemaRegion));
+ createRegionGroupsReq.addRegionGroup("root.test", schemaRegionReplicaSet);
+ partitionInfo.createRegionGroups(createRegionGroupsReq);
+
+ // Create a DataRegion
+ createRegionGroupsReq = new CreateRegionGroupsPlan();
+ final TRegionReplicaSet dataRegionReplicaSet =
+ generateTRegionReplicaSet(
+ testFlag.DataPartition.getFlag(),
+ generateTConsensusGroupId(
+ testFlag.DataPartition.getFlag(),
TConsensusGroupType.DataRegion));
+ createRegionGroupsReq.addRegionGroup("root.test", dataRegionReplicaSet);
+ partitionInfo.createRegionGroups(createRegionGroupsReq);
+
+ // Create a data partition table far larger than the 64KB write buffer
configured in setup(),
+ // so the buffered snapshot stream must flush and wrap around many times.
+ final CreateDataPartitionPlan createDataPartitionPlan = new
CreateDataPartitionPlan();
+ final Map<String, DataPartitionTable> dataPartitionMap = new HashMap<>();
+ final Map<TSeriesPartitionSlot, SeriesPartitionTable> slotInfo = new
HashMap<>();
+ final TConsensusGroupId dataRegionId = dataRegionReplicaSet.getRegionId();
+ for (int seriesSlot = 0; seriesSlot < 2000; seriesSlot++) {
+ final Map<TTimePartitionSlot, List<TConsensusGroupId>> relationInfo =
new HashMap<>();
+ for (int timeSlot = 0; timeSlot < 8; timeSlot++) {
+ relationInfo.put(new TTimePartitionSlot(timeSlot),
Collections.singletonList(dataRegionId));
+ }
+ slotInfo.put(new TSeriesPartitionSlot(seriesSlot), new
SeriesPartitionTable(relationInfo));
+ }
+ dataPartitionMap.put("root.test", new DataPartitionTable(slotInfo));
+ createDataPartitionPlan.setAssignedDataPartition(dataPartitionMap);
+ partitionInfo.createDataPartition(createDataPartitionPlan);
+
+ Assert.assertTrue(partitionInfo.processTakeSnapshot(snapshotDir));
+
+ // The snapshot must actually be larger than the 64KB buffer for this test
to be meaningful.
+ final File snapshotFile = new File(snapshotDir, "partition_info.bin");
+ Assert.assertTrue(snapshotFile.length() > 64 * 1024);
+
+ final PartitionInfo partitionInfo1 = new PartitionInfo();
+ partitionInfo1.processLoadSnapshot(snapshotDir);
+ Assert.assertEquals(partitionInfo, partitionInfo1);
+ }
+
@Test
public void testGetRegionType() {
diff --git
a/iotdb-core/node-commons/src/assembly/resources/conf/edge/iotdb-system.properties
b/iotdb-core/node-commons/src/assembly/resources/conf/edge/iotdb-system.properties
index a4686233b71..01cac68af52 100644
---
a/iotdb-core/node-commons/src/assembly/resources/conf/edge/iotdb-system.properties
+++
b/iotdb-core/node-commons/src/assembly/resources/conf/edge/iotdb-system.properties
@@ -169,6 +169,11 @@ schema_region_ratis_preserve_logs_num_when_purge=200
config_node_ratis_periodic_snapshot_interval=1800
schema_region_ratis_periodic_snapshot_interval=1800
+# Cap the snapshot I/O buffer of the ConfigNode at 8KB (stock default: 4MB).
+# Snapshots are tiny on an edge node, so this keeps their transient heap
+# allocation negligible.
+config_node_snapshot_buffer_size_max=8192
+
# ---- realtime pipe sync out of the box ----
# The pipe memory pool is 10% of the heap (~22MB at the default 224M budget),
# while the stock pipe memory estimates are sized for datacenter nodes: each
diff --git
a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
index e325ff42573..d2e00d2181d 100644
---
a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
+++
b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
@@ -2123,6 +2123,11 @@ config_node_ratis_snapshot_trigger_threshold=400000
schema_region_ratis_snapshot_trigger_threshold=400000
data_region_ratis_snapshot_trigger_threshold=400000
+# max size (in byte) of the in-memory buffer used when taking/loading
ConfigNode snapshot files
+# effectiveMode: restart
+# Datatype: long
+config_node_snapshot_buffer_size_max=4194304
+
# allow flushing Raft Log asynchronously
# effectiveMode: restart
# Datatype: Boolean
diff --git
a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/CommonMessages.java
b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/CommonMessages.java
index 89c4600776b..a6b958fd142 100644
---
a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/CommonMessages.java
+++
b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/CommonMessages.java
@@ -342,4 +342,7 @@ public final class CommonMessages {
EXCEPTION_XCORR_REQUIRES_EXACTLY_TWO_CALCULATION_COLUMNS_BUT_FOUND_ARG_2FF8EB0C
=
"XCorr requires exactly two calculation columns, but found %d.";
public static final String EXCEPTION_COLUMN_LACK_OF_NAME = "the column in
table lack of the name";
+ public static final String
+
EXCEPTION_SNAPSHOT_BUFFER_SIZE_MUST_NOT_EXCEED_ARG_BYTES_BUT_WAS_ARG_D1DA6F7E =
+ "Snapshot buffer size must not exceed %d bytes, but was %d.";
}
diff --git
a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/CommonMessages.java
b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/CommonMessages.java
index f854ddb77f3..c3af2dbe7e2 100644
---
a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/CommonMessages.java
+++
b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/CommonMessages.java
@@ -238,4 +238,7 @@ public final class CommonMessages {
EXCEPTION_XCORR_REQUIRES_EXACTLY_TWO_CALCULATION_COLUMNS_BUT_FOUND_ARG_2FF8EB0C
=
"XCorr 要求必须正好有两列计算列，但实际找到 %d 列。";
public static final String EXCEPTION_COLUMN_LACK_OF_NAME = "表参数中列缺少名字";
+ public static final String
+
EXCEPTION_SNAPSHOT_BUFFER_SIZE_MUST_NOT_EXCEED_ARG_BYTES_BUT_WAS_ARG_D1DA6F7E =
+ "快照缓冲区大小不得超过 %d 字节，但实际为 %d。";
}
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/snapshot/ReusableBufferedOutputStream.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/snapshot/ReusableBufferedOutputStream.java
new file mode 100644
index 00000000000..e18922cac06
--- /dev/null
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/snapshot/ReusableBufferedOutputStream.java
@@ -0,0 +1,134 @@
+/*
+ * 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.commons.snapshot;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.Objects;
+
+/**
+ * A {@link java.io.BufferedOutputStream} variant whose backing byte array is
borrowed from and
+ * returned to {@link SnapshotStreamFactory}'s per-thread pool instead of
being freshly allocated
+ * per stream. Behavior follows {@link java.io.BufferedOutputStream}: {@link
#close()} propagates a
+ * flush failure while still closing the underlying stream.
+ *
+ * <p>Writes larger than the buffer bypass it entirely, exactly like the JDK
implementation would
+ * after flushing.
+ */
+final class ReusableBufferedOutputStream extends OutputStream {
+
+ private final OutputStream out;
+ private final int bufferSize;
+
+ /** Lazily borrowed from {@link SnapshotStreamFactory} on the first write. */
+ private byte[] buffer;
+
+ private int count;
+
+ private boolean closed;
+
+ ReusableBufferedOutputStream(final OutputStream out, final int bufferSize) {
+ this.out = Objects.requireNonNull(out);
+ this.bufferSize = bufferSize;
+ }
+
+ @Override
+ public void write(final int b) throws IOException {
+ ensureOpen();
+ ensureBuffer();
+ if (count >= buffer.length) {
+ flushBuffer();
+ }
+ buffer[count++] = (byte) b;
+ }
+
+ @Override
+ public void write(final byte[] b, final int off, final int len) throws
IOException {
+ ensureOpen();
+ Objects.checkFromIndexSize(off, len, b.length);
+ if (len == 0) {
+ return;
+ }
+ if (buffer == null && len >= bufferSize) {
+ // Large write while the buffer has not even been allocated: skip the
buffer entirely
+ // instead of allocating it just to flush it immediately.
+ out.write(b, off, len);
+ return;
+ }
+ ensureBuffer();
+ if (len >= buffer.length) {
+ flushBuffer();
+ out.write(b, off, len);
+ return;
+ }
+ if (len > buffer.length - count) {
+ flushBuffer();
+ }
+ System.arraycopy(b, off, buffer, count, len);
+ count += len;
+ }
+
+ @Override
+ public void flush() throws IOException {
+ ensureOpen();
+ flushBuffer();
+ out.flush();
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (closed) {
+ return;
+ }
+ try (final OutputStream outputStream = out) {
+ flush();
+ } finally {
+ closed = true;
+ releaseBuffer();
+ }
+ }
+
+ private void ensureBuffer() {
+ if (buffer == null) {
+ buffer = SnapshotStreamFactory.acquireBuffer(bufferSize);
+ }
+ }
+
+ private void ensureOpen() throws IOException {
+ if (closed) {
+ throw new IOException();
+ }
+ }
+
+ private void flushBuffer() throws IOException {
+ if (count > 0) {
+ out.write(buffer, 0, count);
+ count = 0;
+ }
+ }
+
+ private void releaseBuffer() {
+ count = 0;
+ if (buffer != null) {
+ SnapshotStreamFactory.releaseBuffer(buffer);
+ buffer = null;
+ }
+ }
+}
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/snapshot/SnapshotStreamFactory.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/snapshot/SnapshotStreamFactory.java
new file mode 100644
index 00000000000..ceb1b008e28
--- /dev/null
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/snapshot/SnapshotStreamFactory.java
@@ -0,0 +1,130 @@
+/*
+ * 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.commons.snapshot;
+
+import org.apache.iotdb.commons.i18n.CommonMessages;
+
+import java.io.BufferedInputStream;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.lang.ref.SoftReference;
+
+/**
+ * Creates the buffered streams used to write and read ConfigNode snapshot
files.
+ *
+ * <p>The buffer size is bounded by {@link #bufferSizeMax}, which is
configurable through {@code
+ * config_node_snapshot_buffer_size_max} (0 disables buffering). This replaces
the previous fixed
+ * 32MB buffer of {@code PartitionInfo}: memory-constrained deployments can
lower the cap, while the
+ * default keeps snapshot I/O fast for large partition tables. Write buffers
are pooled per thread
+ * through {@link SoftReference}s, so consecutive snapshots on the same thread
reuse the array
+ * instead of re-allocating it.
+ */
+public final class SnapshotStreamFactory {
+
+ /** Default upper bound of a snapshot stream buffer, 4MB. */
+ public static final long DEFAULT_BUFFER_SIZE_MAX = 4 * 1024 * 1024L;
+
+ private static volatile long bufferSizeMax = DEFAULT_BUFFER_SIZE_MAX;
+
+ /** Thread-local pool of reusable write buffers, kept only as long as GC
allows. */
+ private static final ThreadLocal<SoftReference<byte[]>> WRITE_BUFFER_POOL =
+ ThreadLocal.withInitial(() -> new SoftReference<>(null));
+
+ private SnapshotStreamFactory() {
+ // Utility class
+ }
+
+ /**
+ * Set the upper bound of snapshot stream buffers, in bytes. Values below or
equal to zero disable
+ * buffering entirely (the raw stream is returned unchanged). Thread-safe;
takes effect on the
+ * next stream creation.
+ */
+ public static void setBufferSizeMax(final long sizeInBytes) {
+ if (sizeInBytes > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException(
+ String.format(
+ CommonMessages
+
.EXCEPTION_SNAPSHOT_BUFFER_SIZE_MUST_NOT_EXCEED_ARG_BYTES_BUT_WAS_ARG_D1DA6F7E,
+ Integer.MAX_VALUE,
+ sizeInBytes));
+ }
+ bufferSizeMax = Math.max(0L, sizeInBytes);
+ }
+
+ public static long getBufferSizeMax() {
+ return bufferSizeMax;
+ }
+
+ /**
+ * Wrap {@code raw} with a buffered output stream whose buffer is at most
{@link #bufferSizeMax}
+ * bytes. The buffer is reusable, so the same thread writing several
snapshots does not repeatedly
+ * allocate it.
+ *
+ * @param raw the raw stream to buffer
+ * @return a buffered stream, or {@code raw} itself if buffering is disabled
+ */
+ public static OutputStream createOutputStream(final OutputStream raw) {
+ final int bufferSize = (int) bufferSizeMax;
+ return bufferSize <= 0 ? raw : new ReusableBufferedOutputStream(raw,
bufferSize);
+ }
+
+ /**
+ * Wrap {@code raw} with a buffered input stream whose buffer is sized to at
most {@code
+ * fileSize}, capped by {@link #bufferSizeMax}. Reading never allocates more
than the configured
+ * cap, no matter how large the snapshot file is.
+ *
+ * @param raw the raw stream to buffer
+ * @param fileSize size of the file being read, in bytes
+ * @return a buffered stream, or {@code raw} itself if buffering is disabled
or the file is empty
+ */
+ public static InputStream createInputStream(final InputStream raw, final
long fileSize) {
+ final long bufferSize = Math.min(fileSize, bufferSizeMax);
+ return bufferSize <= 0 ? raw : new BufferedInputStream(raw, (int)
bufferSize);
+ }
+
+ /**
+ * Borrow a reusable buffer of at least {@code minSize} bytes that does not
exceed the current
+ * cap. If the thread's pool holds a large enough buffer that still fits
within {@link
+ * #bufferSizeMax}, it is handed out (and removed from the pool, so
concurrent borrowers never see
+ * the same array); otherwise a new buffer is allocated.
+ */
+ static byte[] acquireBuffer(final int minSize) {
+ final long cap = bufferSizeMax;
+ final SoftReference<byte[]> reference = WRITE_BUFFER_POOL.get();
+ final byte[] cached = reference == null ? null : reference.get();
+ // Reuse only a buffer that fits the current cap: after the cap has been
lowered, a larger
+ // pooled buffer must not be handed out again, since the backing buffer
may never exceed
+ // bufferSizeMax.
+ if (cached != null && cached.length >= minSize && cached.length <= cap) {
+ // Take ownership of the cached buffer so a nested borrower allocates
its own instead of
+ // silently sharing the array.
+ WRITE_BUFFER_POOL.set(new SoftReference<>(null));
+ return cached;
+ }
+ return new byte[minSize];
+ }
+
+ /** Return a buffer to the thread's pool for reuse by a later snapshot on
the same thread. */
+ static void releaseBuffer(final byte[] buffer) {
+ if (buffer != null) {
+ WRITE_BUFFER_POOL.set(new SoftReference<>(buffer));
+ }
+ }
+}
diff --git
a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/snapshot/SnapshotStreamFactoryTest.java
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/snapshot/SnapshotStreamFactoryTest.java
new file mode 100644
index 00000000000..e4e33e5e83f
--- /dev/null
+++
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/snapshot/SnapshotStreamFactoryTest.java
@@ -0,0 +1,307 @@
+/*
+ * 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.commons.snapshot;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+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.util.Arrays;
+import java.util.Random;
+
+public class SnapshotStreamFactoryTest {
+
+ @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ private long originalBufferSizeMax;
+
+ @Before
+ public void setUp() {
+ originalBufferSizeMax = SnapshotStreamFactory.getBufferSizeMax();
+ }
+
+ @After
+ public void tearDown() {
+ // The buffer size cap is process-global; restore it so tests do not leak
state into each
+ // other.
+ SnapshotStreamFactory.setBufferSizeMax(originalBufferSizeMax);
+ }
+
+ @Test
+ public void testDisabledBufferingReturnsRawStreams() {
+ SnapshotStreamFactory.setBufferSizeMax(0);
+
+ final OutputStream rawOut = new NullOutputStream();
+ final InputStream rawIn = new NullInputStream();
+ Assert.assertSame(rawOut,
SnapshotStreamFactory.createOutputStream(rawOut));
+ Assert.assertSame(rawIn, SnapshotStreamFactory.createInputStream(rawIn,
1024));
+ }
+
+ @Test
+ public void testBufferSizeMaximumBoundary() {
+ SnapshotStreamFactory.setBufferSizeMax(Integer.MAX_VALUE);
+ Assert.assertEquals(Integer.MAX_VALUE,
SnapshotStreamFactory.getBufferSizeMax());
+
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> SnapshotStreamFactory.setBufferSizeMax((long) Integer.MAX_VALUE
+ 1));
+ Assert.assertEquals(Integer.MAX_VALUE,
SnapshotStreamFactory.getBufferSizeMax());
+ }
+
+ @Test
+ public void testRoundTripThroughBufferedStreams() throws IOException {
+ // 100KB of random data with a 64KB buffer: writes must wrap the buffer
several times.
+ SnapshotStreamFactory.setBufferSizeMax(64 * 1024);
+ final byte[] data = new byte[100 * 1024];
+ new Random(42).nextBytes(data);
+
+ final File file = temporaryFolder.newFile();
+ try (FileOutputStream fileOutputStream = new FileOutputStream(file);
+ OutputStream outputStream =
SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
+ // Write in 4KB chunks so the 64KB buffer is actually filled, flushed
and wrapped around.
+ for (int offset = 0; offset < data.length; offset += 4096) {
+ outputStream.write(data, offset, Math.min(4096, data.length - offset));
+ }
+ }
+
+ final byte[] readBack = new byte[data.length];
+ try (FileInputStream fileInputStream = new FileInputStream(file);
+ InputStream inputStream =
+ SnapshotStreamFactory.createInputStream(fileInputStream,
file.length())) {
+ int offset = 0;
+ while (offset < readBack.length) {
+ final int read = inputStream.read(readBack, offset, readBack.length -
offset);
+ if (read < 0) {
+ break;
+ }
+ offset += read;
+ }
+ Assert.assertEquals(readBack.length, offset);
+ }
+ Assert.assertArrayEquals(data, readBack);
+ }
+
+ @Test
+ public void testLargeWriteBypassesBuffer() throws IOException {
+ // A single write larger than the buffer must bypass it and still land
correctly on the
+ // underlying stream.
+ SnapshotStreamFactory.setBufferSizeMax(64 * 1024);
+ final byte[] data = new byte[1024 * 1024];
+ new Random(7).nextBytes(data);
+
+ final File file = temporaryFolder.newFile();
+ try (FileOutputStream fileOutputStream = new FileOutputStream(file);
+ OutputStream outputStream =
SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
+ outputStream.write(data);
+ }
+ Assert.assertEquals(data.length, file.length());
+
+ final byte[] readBack = new byte[data.length];
+ try (InputStream inputStream = new FileInputStream(file)) {
+ int offset = 0;
+ while (offset < readBack.length) {
+ final int read = inputStream.read(readBack, offset, readBack.length -
offset);
+ if (read < 0) {
+ break;
+ }
+ offset += read;
+ }
+ Assert.assertEquals(readBack.length, offset);
+ }
+ Assert.assertTrue(Arrays.equals(data, readBack));
+ }
+
+ @Test
+ public void testWriteBufferReuse() {
+ // Use an explicit cap so the pool semantics do not depend on state left
by other tests.
+ SnapshotStreamFactory.setBufferSizeMax(128 * 1024);
+ final byte[] first = SnapshotStreamFactory.acquireBuffer(64 * 1024);
+ SnapshotStreamFactory.releaseBuffer(first);
+
+ // The released buffer is handed out again for a request that fits.
+ final byte[] second = SnapshotStreamFactory.acquireBuffer(64 * 1024);
+ Assert.assertSame(first, second);
+ SnapshotStreamFactory.releaseBuffer(second);
+
+ // A borrowed buffer is removed from the pool: a concurrent borrow must
not share it.
+ final byte[] borrowed = SnapshotStreamFactory.acquireBuffer(64 * 1024);
+ final byte[] other = SnapshotStreamFactory.acquireBuffer(64 * 1024);
+ Assert.assertNotSame(borrowed, other);
+ SnapshotStreamFactory.releaseBuffer(other);
+
+ // A smaller request reuses the larger cached buffer while it still fits
the cap.
+ final byte[] smaller = SnapshotStreamFactory.acquireBuffer(1024);
+ Assert.assertSame(other, smaller);
+ SnapshotStreamFactory.releaseBuffer(smaller);
+
+ // A request larger than the cached buffer allocates a fresh one.
+ final byte[] tiny = new byte[32 * 1024];
+ SnapshotStreamFactory.releaseBuffer(tiny);
+ final byte[] larger = SnapshotStreamFactory.acquireBuffer(64 * 1024);
+ Assert.assertNotSame(tiny, larger);
+ SnapshotStreamFactory.releaseBuffer(borrowed);
+ SnapshotStreamFactory.releaseBuffer(larger);
+ }
+
+ @Test
+ public void testBufferNotReusedAfterCapDecrease() {
+ // Prime the pool with a large write buffer under a large cap.
+ SnapshotStreamFactory.setBufferSizeMax(128 * 1024);
+ final byte[] large = SnapshotStreamFactory.acquireBuffer(128 * 1024);
+ SnapshotStreamFactory.releaseBuffer(large);
+
+ // Lower the cap: the pooled buffer now exceeds it and must not be handed
out again.
+ SnapshotStreamFactory.setBufferSizeMax(8192);
+ final byte[] borrowed = SnapshotStreamFactory.acquireBuffer(8192);
+ Assert.assertNotSame(large, borrowed);
+ Assert.assertTrue(borrowed.length <= 8192);
+ SnapshotStreamFactory.releaseBuffer(borrowed);
+
+ // The buffer allocated under the new cap is still reused for requests
that fit it.
+ final byte[] again = SnapshotStreamFactory.acquireBuffer(8192);
+ Assert.assertSame(borrowed, again);
+ SnapshotStreamFactory.releaseBuffer(again);
+ }
+
+ @Test
+ public void testCloseIsIdempotentAndWriteAfterCloseFailsOnFlush() throws
IOException {
+ final File file = temporaryFolder.newFile();
+ final OutputStream outputStream =
+ SnapshotStreamFactory.createOutputStream(new FileOutputStream(file));
+ outputStream.write(1);
+ outputStream.close();
+ // Double close must not throw.
+ outputStream.close();
+
+ // Writes and flushes after close must fail instead of silently losing
data.
+ Assert.assertThrows(IOException.class, () -> outputStream.write(2));
+ Assert.assertThrows(IOException.class, outputStream::flush);
+ // Closing an already closed stream remains idempotent.
+ outputStream.close();
+ }
+
+ @Test
+ public void testClosePropagatesFlushFailureAndClosesUnderlyingStream()
throws IOException {
+ SnapshotStreamFactory.setBufferSizeMax(64);
+ final FailingOutputStream rawOut = new FailingOutputStream(false);
+ final OutputStream outputStream =
SnapshotStreamFactory.createOutputStream(rawOut);
+ outputStream.write(1);
+
+ final IOException exception = Assert.assertThrows(IOException.class,
outputStream::close);
+ Assert.assertEquals("write failure", exception.getMessage());
+ Assert.assertTrue(rawOut.closed);
+ }
+
+ @Test
+ public void testCloseSuppressesUnderlyingCloseFailureAfterFlushFailure()
throws IOException {
+ SnapshotStreamFactory.setBufferSizeMax(64);
+ final FailingOutputStream rawOut = new FailingOutputStream(true);
+ final OutputStream outputStream =
SnapshotStreamFactory.createOutputStream(rawOut);
+ outputStream.write(1);
+
+ final IOException exception = Assert.assertThrows(IOException.class,
outputStream::close);
+ Assert.assertEquals("write failure", exception.getMessage());
+ Assert.assertEquals(1, exception.getSuppressed().length);
+ Assert.assertEquals("close failure",
exception.getSuppressed()[0].getMessage());
+ Assert.assertTrue(rawOut.closed);
+ }
+
+ @Test
+ public void testInputBufferNeverExceedsFileSizeOrCap() throws IOException {
+ // The read buffer of a stream created for a small file must be capped,
and reading must
+ // still see the whole content.
+ final byte[] data = new byte[100];
+ new Random(3).nextBytes(data);
+ final File file = temporaryFolder.newFile();
+ try (FileOutputStream fileOutputStream = new FileOutputStream(file)) {
+ fileOutputStream.write(data);
+ }
+
+ SnapshotStreamFactory.setBufferSizeMax(64);
+ final byte[] readBack = new byte[data.length];
+ try (FileInputStream fileInputStream = new FileInputStream(file);
+ InputStream inputStream =
+ SnapshotStreamFactory.createInputStream(fileInputStream,
file.length())) {
+ int offset = 0;
+ while (offset < readBack.length) {
+ final int read = inputStream.read(readBack, offset, readBack.length -
offset);
+ if (read < 0) {
+ break;
+ }
+ offset += read;
+ }
+ Assert.assertEquals(readBack.length, offset);
+ }
+ Assert.assertArrayEquals(data, readBack);
+ }
+
+ /** OutputStream that discards everything, used to test the
disabled-buffering fast path. */
+ private static final class NullOutputStream extends OutputStream {
+ @Override
+ public void write(final int b) {
+ // discard
+ }
+ }
+
+ /** InputStream that is always at EOF, used to test the disabled-buffering
fast path. */
+ private static final class NullInputStream extends InputStream {
+ @Override
+ public int read() {
+ return -1;
+ }
+ }
+
+ private static final class FailingOutputStream extends OutputStream {
+
+ private final boolean failOnClose;
+ private boolean closed;
+
+ private FailingOutputStream(final boolean failOnClose) {
+ this.failOnClose = failOnClose;
+ }
+
+ @Override
+ public void write(final int b) throws IOException {
+ throw new IOException("write failure");
+ }
+
+ @Override
+ public void write(final byte[] b, final int off, final int len) throws
IOException {
+ throw new IOException("write failure");
+ }
+
+ @Override
+ public void close() throws IOException {
+ closed = true;
+ if (failOnClose) {
+ throw new IOException("close failure");
+ }
+ }
+ }
+}
