Fix ConfigNode snapshot buffering and recovery

Author: CRZbulabula

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

@@ -26,6 +26,7 @@
 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;
@@ -981,6 +982,14 @@ public long getConfigNodeSnapshotBufferSizeMax() {
   }
 
   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;
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/externalservice/ExternalServiceInfo.java

@@ -277,6 +277,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws IOException {
       serializeInfos(bufferedOutputStream);
 
       // fsync
+      bufferedOutputStream.flush();
       fileOutputStream.getFD().sync();
 
       return true;

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

@@ -235,6 +235,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
     try (FileOutputStream fileOutputStream = new FileOutputStream(snapshotFile);
         OutputStream outputStream = SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
       ttlCache.serialize(outputStream);
+      outputStream.flush();
       fileOutputStream.getFD().sync();
     } finally {
       lock.writeLock().unlock();

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

@@ -263,6 +263,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
       triggerTable.serializeTriggerTable(bufferedOutputStream);
 
       // fsync
+      bufferedOutputStream.flush();
       fileOutputStream.getFD().sync();
       return true;
     }

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

@@ -230,6 +230,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws IOException {
       udfTable.serializeUDFTable(bufferedOutputStream);
 
       // fsync
+      bufferedOutputStream.flush();
       fileOutputStream.getFD().sync();
 
       return true;

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/cq/CQInfo.java

@@ -246,6 +246,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
             SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
 
       serialize(bufferedOutputStream);
+      bufferedOutputStream.flush();
       fileOutputStream.getFD().sync();
       return true;
     } finally {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/node/NodeInfo.java

@@ -26,7 +26,6 @@
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.cluster.NodeStatus;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
-import org.apache.iotdb.commons.snapshot.ByteBufferInputStream;
 import org.apache.iotdb.commons.snapshot.SnapshotProcessor;
 import org.apache.iotdb.commons.snapshot.SnapshotStreamFactory;
 import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
@@ -61,7 +60,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
 import java.nio.channels.FileChannel;
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
@@ -745,11 +744,9 @@ public void processLoadSnapshot(File snapshotDir) throws IOException, TException
     aiNodeInfoReadWriteLock.writeLock().lock();
     versionInfoReadWriteLock.writeLock().lock();
 
-    // The snapshot file is read into a direct buffer instead of the heap, so loading node_info.bin
-    // no longer allocates a heap array as large as the whole file.
     try (final FileChannel channel =
             FileChannel.open(snapshotFile.toPath(), StandardOpenOption.READ);
-        final InputStream inputStream = readSnapshotIntoByteBufferInputStream(channel);
+        final InputStream inputStream = Channels.newInputStream(channel);
         TIOStreamTransport tioStreamTransport = new TIOStreamTransport(inputStream)) {
       TProtocol protocol = new TBinaryProtocol(tioStreamTransport);
 
@@ -763,7 +760,7 @@ public void processLoadSnapshot(File snapshotDir) throws IOException, TException
 
       // TODO: Compatibility design. Should replace this function to actual deserialization method
       // in IoTDB 2.2 / 1.5
-      tryDeserializeRegisteredAINode(inputStream, protocol);
+      tryDeserializeRegisteredAINode(snapshotFile, channel);
 
       deserializeBuildInfo(inputStream);
 
@@ -775,22 +772,6 @@ public void processLoadSnapshot(File snapshotDir) throws IOException, TException
     }
   }
 
-  private static ByteBufferInputStream readSnapshotIntoByteBufferInputStream(
-      final FileChannel channel) throws IOException {
-    final long fileSize = channel.size();
-    if (fileSize == 0) {
-      return new ByteBufferInputStream(ByteBuffer.allocate(0));
-    }
-    final ByteBuffer byteBuffer = ByteBuffer.allocateDirect((int) fileSize);
-    while (byteBuffer.hasRemaining()) {
-      if (channel.read(byteBuffer) < 0) {
-        break;
-      }
-    }
-    byteBuffer.flip();
-    return new ByteBufferInputStream(byteBuffer);
-  }
-
   private void deserializeRegisteredConfigNode(InputStream inputStream, TProtocol protocol)
       throws IOException, TException {
     int size = ReadWriteIOUtils.readInt(inputStream);
@@ -815,15 +796,20 @@ private void deserializeRegisteredDataNode(InputStream inputStream, TProtocol pr
     }
   }
 
-  private void tryDeserializeRegisteredAINode(InputStream inputStream, TProtocol protocol)
+  private void tryDeserializeRegisteredAINode(File snapshotFile, FileChannel channel)
       throws IOException {
-    try {
-      // 0 has no meaning here
-      inputStream.mark(0);
-      deserializeRegisteredAINode(inputStream, protocol);
+    final long position = channel.position();
+    try (final FileChannel probeChannel =
+            FileChannel.open(snapshotFile.toPath(), StandardOpenOption.READ);
+        final InputStream probeInputStream = Channels.newInputStream(probeChannel);
+        final TIOStreamTransport probeTransport = new TIOStreamTransport(probeInputStream)) {
+      probeChannel.position(position);
+      deserializeRegisteredAINode(probeInputStream, new TBinaryProtocol(probeTransport));
+      channel.position(probeChannel.position());
     } catch (IOException | TException ignore) {
       // Exception happens here means that the data is upgraded from the old version
-      inputStream.reset();
+      registeredAINodes.clear();
+      channel.position(position);
     }
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipePluginInfo.java

@@ -416,6 +416,7 @@ public boolean processTakeSnapshot(final File snapshotDir) throws IOException {
           final OutputStream bufferedOutputStream =
               SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
         pipePluginMetaKeeper.processTakeSnapshot(bufferedOutputStream);
+        bufferedOutputStream.flush();
         fileOutputStream.getFD().sync();
       }
       return true;

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipeTaskInfo.java

@@ -1328,6 +1328,7 @@ public boolean processTakeSnapshot(final File snapshotDir) throws IOException {
           final OutputStream bufferedOutputStream =
               SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
         pipeMetaKeeper.processTakeSnapshot(bufferedOutputStream);
+        bufferedOutputStream.flush();
         fileOutputStream.getFD().sync();
       }
       return true;

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java

@@ -164,6 +164,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept
             SnapshotStreamFactory.createOutputStream(fileOutputStream)) {
       serializeSpaceQuotaLimit(bufferedOutputStream);
       serializeThrottleQuotaLimit(bufferedOutputStream);
+      bufferedOutputStream.flush();
       fileOutputStream.getFD().sync();
     } finally {
       spaceQuotaReadWriteLock.writeLock().unlock();