feat(load): route decoded tsfile pieces through consensus

integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java

@@ -0,0 +1,231 @@
+/*
+ * 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.it;
+
+import org.apache.iotdb.consensus.ConsensusFactory;
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.it.utils.TsFileGenerator;
+import org.apache.iotdb.itbase.category.ClusterIT;
+import org.apache.iotdb.jdbc.Config;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.write.schema.IMeasurementSchema;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.List;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({ClusterIT.class})
+public class IoTDBLoadTsFileClusterIT {
+
+  private static final long PARTITION_INTERVAL = 10_000L;
+  private static final String DATABASE = "root.load_cluster";
+  private static final List<String> DEVICES =
+      Arrays.asList(
+          DATABASE + ".d1", DATABASE + ".d2", DATABASE + ".d3", DATABASE + ".d4", DATABASE + ".d5");
+  private static final List<String> ALIGNED_DEVICES =
+      Arrays.asList(DATABASE + ".a1", DATABASE + ".a2");
+  private static final List<String> MEASUREMENTS = Arrays.asList("s1", "s2", "s3", "s4");
+  private static final int POINT_COUNT_PER_DEVICE = 10_000;
+  private static final int REPLICATION_FACTOR = 3;
+  private static final int CONFIG_NODE_NUM = 3;
+  private static final int DATA_NODE_NUM = 3;
+
+  private File tmpDir;
+
+  @Before
+  public void setUp() throws Exception {
+    tmpDir = new File(Files.createTempDirectory("load-cluster-it").toUri());
+    EnvFactory.getEnv().getConfig().getCommonConfig().setTimePartitionInterval(PARTITION_INTERVAL);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setPipeMemoryManagementEnabled(false);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setDatanodeMemoryProportion("1:10:1:1:1:0");
+    EnvFactory.getEnv().getConfig().getCommonConfig().setTargetChunkPointNum(1_000);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setMaxNumberOfPointsInPage(500);
+    EnvFactory.getEnv()
+        .getConfig()
+        .getCommonConfig()
+        .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
+        .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
+        .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS)
+        .setSchemaReplicationFactor(REPLICATION_FACTOR)
+        .setDataReplicationFactor(REPLICATION_FACTOR);
+    EnvFactory.getEnv().initClusterEnvironment(CONFIG_NODE_NUM, DATA_NODE_NUM);
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute("delete database " + DATABASE);
+    } catch (final Exception ignored) {
+    }
+
+    EnvFactory.getEnv().cleanClusterEnvironment();
+
+    final File[] files = tmpDir.listFiles();
+    if (files != null) {
+      for (final File file : files) {
+        Assert.assertTrue(file.delete());
+      }
+    }
+    Assert.assertTrue(tmpDir.delete());
+  }
+
+  @Test
+  public void testLoadTsFileInCluster() throws Exception {
+    final long writtenPointCount;
+    try (final TsFileGenerator generator =
+        new TsFileGenerator(new File(tmpDir, "load-cluster-1-0-0-0.tsfile"))) {
+      final List<IMeasurementSchema> schemas =
+          Arrays.asList(
+              new MeasurementSchema(MEASUREMENTS.get(0), TSDataType.INT64, TSEncoding.PLAIN),
+              new MeasurementSchema(MEASUREMENTS.get(1), TSDataType.INT64, TSEncoding.PLAIN),
+              new MeasurementSchema(MEASUREMENTS.get(2), TSDataType.INT64, TSEncoding.PLAIN),
+              new MeasurementSchema(MEASUREMENTS.get(3), TSDataType.INT64, TSEncoding.PLAIN));
+      for (final String device : DEVICES) {
+        generator.registerTimeseries(device, schemas);
+        generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, false);
+      }
+      for (final String device : ALIGNED_DEVICES) {
+        generator.registerAlignedTimeseries(device, schemas);
+        generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, true, 1_000_000L);
+      }
+      writtenPointCount = generator.getTotalNumber();
+    }
+
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute("create database " + DATABASE);
+      for (final String device : DEVICES) {
+        for (final String measurement : MEASUREMENTS) {
+          statement.execute(
+              "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name());
+        }
+      }
+      for (final String device : ALIGNED_DEVICES) {
+        statement.execute(
+            "create aligned timeseries " + device + "(s1 INT64, s2 INT64, s3 INT64, s4 INT64)");
+      }
+      statement.execute("load \"" + tmpDir.getAbsolutePath() + "\"");
+
+      // queryPointCount also covers the aligned devices, which writtenPointCount includes.
+      Assert.assertEquals(writtenPointCount, queryPointCount(statement));
+    }
+  }
+
+  @Test
+  public void testLoadAfterOneDataNodeDown() throws Exception {
+    final long firstPointCount = generateTsFile("load-first.tsfile", 5_000, 0L);
+
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute("create database " + DATABASE);
+      for (final String device : DEVICES) {
+        for (final String measurement : MEASUREMENTS) {
+          statement.execute(
+              "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name());
+        }
+      }
+      for (final String device : ALIGNED_DEVICES) {
+        statement.execute(
+            "create aligned timeseries " + device + "(s1 INT64, s2 INT64, s3 INT64, s4 INT64)");
+      }
+
+      statement.execute("load \"" + tmpDir.getAbsolutePath() + "\"");
+    }
+
+    EnvFactory.getEnv().getDataNodeWrapper(0).stop();
+    Thread.sleep(1_000L);
+
+    final long secondPointCount = generateTsFile("load-second.tsfile", 3_000, 100_000L);
+    try (final Connection connection =
+            DriverManager.getConnection(
+                Config.IOTDB_URL_PREFIX
+                    + EnvFactory.getEnv().getDataNodeWrapper(1).getIpAndPortString(),
+                "root",
+                "root");
+        final Statement statement = connection.createStatement()) {
+      statement.execute("load \"" + tmpDir.getAbsolutePath() + "\"");
+
+      final long expectedPointCount = firstPointCount + secondPointCount;
+      long actualPointCount = queryPointCount(statement);
+      for (int retry = 0; actualPointCount < expectedPointCount && retry < 30; retry++) {
+        Thread.sleep(1_000L);
+        actualPointCount = queryPointCount(statement);
+      }
+      Assert.assertEquals(expectedPointCount, actualPointCount);
+    }
+  }
+
+  private long generateTsFile(
+      final String fileName, final int pointCountPerDevice, final long startTimestamp)
+      throws Exception {
+    final List<IMeasurementSchema> schemas =
+        Arrays.asList(
+            new MeasurementSchema(MEASUREMENTS.get(0), TSDataType.INT64, TSEncoding.PLAIN),
+            new MeasurementSchema(MEASUREMENTS.get(1), TSDataType.INT64, TSEncoding.PLAIN),
+            new MeasurementSchema(MEASUREMENTS.get(2), TSDataType.INT64, TSEncoding.PLAIN),
+            new MeasurementSchema(MEASUREMENTS.get(3), TSDataType.INT64, TSEncoding.PLAIN));
+    try (final TsFileGenerator generator = new TsFileGenerator(new File(tmpDir, fileName))) {
+      for (final String device : DEVICES) {
+        generator.registerTimeseries(device, schemas);
+        generator.generateData(device, pointCountPerDevice, 1L, false, startTimestamp);
+      }
+      final long alignedStartTimestamp = startTimestamp + 1_000_000L;
+      for (final String device : ALIGNED_DEVICES) {
+        generator.registerAlignedTimeseries(device, schemas);
+        generator.generateData(device, pointCountPerDevice, 1L, true, alignedStartTimestamp);
+      }
+      return generator.getTotalNumber();
+    }
+  }
+
+  private long queryPointCount(final Statement statement) throws Exception {
+    long actualPointCount = 0;
+    final List<String> allDevices = new java.util.ArrayList<>(DEVICES);
+    allDevices.addAll(ALIGNED_DEVICES);
+    for (final String device : allDevices) {
+      for (final String measurement : MEASUREMENTS) {
+        try (final ResultSet resultSet =
+            statement.executeQuery("select count(" + measurement + ") from " + device)) {
+          Assert.assertTrue(resultSet.next());
+          actualPointCount += resultSet.getLong(1);
+        }
+      }
+    }
+    return actualPointCount;
+  }
+}

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/common/request/IndexedConsensusRequest.java

@@ -49,6 +49,7 @@ public class IndexedConsensusRequest implements IConsensusRequest {
   private long memorySize = 0;
   private long retainedMemorySize = 0;
   private boolean serializedRequestsBuilt = false;
+  private long queueReservedMemorySize = -1L;
   private boolean containsUserData = false;
   private final AtomicLong referenceCnt = new AtomicLong();
 
@@ -86,11 +87,19 @@ public ByteBuffer serializeToByteBuffer() {
     throw new UnsupportedOperationException();
   }
 
+  /** Returns whether any request of this entry materializes its bytes only when it is sent. */
+  public synchronized boolean hasDeferredRequests() {
+    return requests.stream().anyMatch(IConsensusRequest::isSerializationDeferred);
+  }
+
   public List<IConsensusRequest> getRequests() {
     return requests;
   }
 
-  public List<ByteBuffer> getSerializedRequests() {
+  public synchronized List<ByteBuffer> getSerializedRequests() {
+    // Requests that deferred their serialization materialize their bytes here, that is, when a
+    // batch is actually built for sending.
+    buildSerializedRequests();
     return serializedRequests;
   }
 
@@ -112,6 +121,20 @@ public synchronized long getRetainedMemorySize() {
     return requests.stream().mapToLong(IConsensusRequest::getMemorySize).sum();
   }
 
+  /**
+   * Returns the amount to reserve for this entry while it waits in the replication queues.
+   *
+   * <p>The amount is captured on the first call, because a request that deferred its serialization
+   * materializes its bytes only when it is sent and its serialized buffers are accounted for by the
+   * batch that carries them. Releasing an entry therefore has to return exactly what was reserved.
+   */
+  public synchronized long getQueueReservedMemorySize() {
+    if (queueReservedMemorySize < 0) {
+      queueReservedMemorySize = getRetainedMemorySize();
+    }
+    return queueReservedMemorySize;
+  }
+
   /**
    * Releases the original request objects after their serialized buffers have been materialized.
    * Replication only needs the serialized buffers, while subscription delivery still needs the

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/log/ConsensusReqReader.java

@@ -38,6 +38,14 @@ public interface ConsensusReqReader {
 
   void setSafelyDeletedSearchIndex(long safelyDeletedSearchIndex);
 
+  /**
+   * @return the search index before which insert nodes can be deleted safely, or {@link
+   *     #DEFAULT_SAFELY_DELETED_SEARCH_INDEX} if the consensus layer never reported a watermark.
+   */
+  default long getSafelyDeletedSearchIndex() {
+    return DEFAULT_SAFELY_DELETED_SEARCH_INDEX;
+  }
+
   /**
    * Gets the consensus requests iterator from the specified start position.
    *

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/IoTConsensusMemoryManager.java

@@ -47,7 +47,9 @@ public boolean reserve(IndexedConsensusRequest request) {
     synchronized (request) {
       long prevRef = request.incRef();
       if (prevRef == 0) {
-        final long retainedMemorySize = request.getRetainedMemorySize();
+        // The amount is captured once, so that an entry that materializes its bytes after it was
+        // queued returns exactly what it reserved when it is released.
+        final long retainedMemorySize = request.getQueueReservedMemorySize();
         boolean reserved = reserve(retainedMemorySize, true);
         if (reserved) {
           if (logger.isDebugEnabled()) {
@@ -114,7 +116,7 @@ public void free(IndexedConsensusRequest request) {
     synchronized (request) {
       long prevRef = request.decRef();
       if (prevRef == 1) {
-        final long retainedMemorySize = request.getRetainedMemorySize();
+        final long retainedMemorySize = request.getQueueReservedMemorySize();
         free(retainedMemorySize, true);
         if (logger.isDebugEnabled()) {
           logger.debug(

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java

@@ -190,8 +190,14 @@ public void offer(IndexedConsensusRequest request) {
   public void offer(IndexedConsensusRequest request, boolean keepRequestsForSubscription) {
     // we don't need to serialize and offer request when replicaNum is 1.
     if (!threads.isEmpty()) {
-      request.buildSerializedRequests();
-      if (!keepRequestsForSubscription) {
+      // A request that deferred its serialization keeps only its source objects, which stay small,
+      // so that the entries waiting in the queues to be replicated do not retain their payloads.
+      // Such requests materialize their bytes when a batch is built instead.
+      final boolean serializationDeferred = request.hasDeferredRequests();
+      if (!serializationDeferred) {
+        request.buildSerializedRequests();
+      }
+      if (!keepRequestsForSubscription && !serializationDeferred) {
         request.clearRequests();
       }
       synchronized (this) {

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/subscription/SubscriptionQueueRegistry.java

@@ -115,8 +115,12 @@ public synchronized boolean offer(final IndexedConsensusRequest indexedConsensus
       return false;
     }
 
-    // Subscription queues reserve request memory, including the serialized buffers.
-    indexedConsensusRequest.buildSerializedRequests();
+    // Subscription queues reserve request memory, including the serialized buffers. A request that
+    // deferred its serialization keeps only its source objects until it is sent, and subscription
+    // delivery reads those objects, so nothing is materialized here.
+    if (!indexedConsensusRequest.hasDeferredRequests()) {
+      indexedConsensusRequest.buildSerializedRequests();
+    }
 
     if (LOGGER.isDebugEnabled()) {
       LOGGER.debug(