[To dev/1.3] Cherry-pick some IoTConsensus related commits to 1.3

integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java

@@ -100,4 +100,16 @@ public DataNodeConfig setCacheLastValuesForLoad(boolean cacheLastValuesForLoad)
  setProperty("cache_last_values_for_load", String.valueOf(cacheLastValuesForLoad));
  return this;
  }
+
+ @Override
+ public DataNodeConfig setWalThrottleSize(long walThrottleSize) {
+ setProperty("wal_throttle_threshold_in_byte", String.valueOf(walThrottleSize));
+ return this;
+ }
+
+ @Override
+ public DataNodeConfig setDeleteWalFilesPeriodInMs(long deleteWalFilesPeriodInMs) {
+ setProperty("delete_wal_files_period_in_ms", String.valueOf(deleteWalFilesPeriodInMs));
+ return this;
+ }
 }

integration-test/src/main/java/org/apache/iotdb/it/env/cluster/node/AbstractNodeWrapper.java

@@ -758,4 +758,8 @@ public long getPid() {
  return -1;
  }
  }
+
+ public Process getInstance() {
+ return instance;
+ }
 }

integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java

@@ -63,4 +63,14 @@ public DataNodeConfig setLoadLastCacheStrategy(String strategyName) {
  public DataNodeConfig setCacheLastValuesForLoad(boolean cacheLastValuesForLoad) {
  return this;
  }
+
+ @Override
+ public DataNodeConfig setWalThrottleSize(long walThrottleSize) {
+ return this;
+ }
+
+ @Override
+ public DataNodeConfig setDeleteWalFilesPeriodInMs(long deleteWalFilesPeriodInMs) {
+ return this;
+ }
 }

integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java

@@ -39,4 +39,8 @@ DataNodeConfig setLoadTsFileAnalyzeSchemaMemorySizeInBytes(
  DataNodeConfig setLoadLastCacheStrategy(String strategyName);
 
  DataNodeConfig setCacheLastValuesForLoad(boolean cacheLastValuesForLoad);
+
+ DataNodeConfig setWalThrottleSize(long walThrottleSize);
+
+ DataNodeConfig setDeleteWalFilesPeriodInMs(long deleteWalFilesPeriodInMs);
 }

integration-test/src/test/java/org/apache/iotdb/confignode/it/regionmigration/IoTDBRegionOperationReliabilityITFramework.java

@@ -44,6 +44,7 @@
 
 import org.apache.thrift.TException;
 import org.apache.tsfile.read.common.Field;
+import org.apache.tsfile.utils.Pair;
 import org.awaitility.Awaitility;
 import org.awaitility.core.ConditionTimeoutException;
 import org.junit.After;
@@ -307,7 +308,7 @@ public void generalTestWithAllOptions(
  }
  }
 
-protected Set<Integer> getAllDataNodes(Statement statement) throws Exception {
+public static Set<Integer> getAllDataNodes(Statement statement) throws Exception {
  ResultSet result = statement.executeQuery(SHOW_DATANODES);
  Set<Integer> allDataNodeId = new HashSet<>();
  while (result.next()) {
@@ -444,6 +445,26 @@ public static Map<Integer, Set<Integer>> getDataRegionMap(Statement statement) t
  return regionMap;
  }
 
+ public static Map<Integer, Pair<Integer, Set<Integer>>> getDataRegionMapWithLeader(
+ Statement statement) throws Exception {
+ ResultSet showRegionsResult = statement.executeQuery(SHOW_REGIONS);
+ Map<Integer, Pair<Integer, Set<Integer>>> regionMap = new HashMap<>();
+ while (showRegionsResult.next()) {
+ if (String.valueOf(TConsensusGroupType.DataRegion)
+ .equals(showRegionsResult.getString(ColumnHeaderConstant.TYPE))) {
+ int regionId = showRegionsResult.getInt(ColumnHeaderConstant.REGION_ID);
+ int dataNodeId = showRegionsResult.getInt(ColumnHeaderConstant.DATA_NODE_ID);
+ Pair<Integer, Set<Integer>> leaderNodesPair =
+ regionMap.computeIfAbsent(regionId, id -> new Pair<>(-1, new HashSet<>()));
+ leaderNodesPair.getRight().add(dataNodeId);
+ if (showRegionsResult.getString(ColumnHeaderConstant.ROLE).equals("Leader")) {
+ leaderNodesPair.setLeft(dataNodeId);
+ }
+ }
+ }
+ return regionMap;
+ }
+
  public static Map<Integer, Set<Integer>> getAllRegionMap(Statement statement) throws Exception {
  ResultSet showRegionsResult = statement.executeQuery(SHOW_REGIONS);
  Map<Integer, Set<Integer>> regionMap = new HashMap<>();

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

@@ -0,0 +1,148 @@
+/*
+ * 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.it.env.cluster.env.SimpleEnv;
+import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.ClusterIT;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.StatementExecutionException;
+import org.apache.iotdb.session.Session;
+
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Date;
+
+import static org.junit.Assert.fail;
+
+/** Tests that may not be satisfied with the default cluster settings. */
+@RunWith(IoTDBTestRunner.class)
+@Category({ClusterIT.class})
+public class IoTDBCustomizedClusterIT {
+
+ private final Logger logger = LoggerFactory.getLogger(IoTDBCustomizedClusterIT.class);
+
+ /**
+ * When the wal size exceeds `walThrottleSize` * 0.8, the timed wal-delete-thread will try
+ * deleting wal forever, which will block the DataNode from exiting, because task of deleting wal
+ * submitted by the ShutdownHook cannot be executed. This test ensures that this blocking is
+ * fixed.
+ */
+ @Test
+ public void testWalThrottleStuck()
+ throws SQLException,
+ IoTDBConnectionException,
+ StatementExecutionException,
+ InterruptedException {
+ SimpleEnv simpleEnv = new SimpleEnv();
+ simpleEnv
+ .getConfig()
+ .getDataNodeConfig()
+ .setWalThrottleSize(1)
+ .setDeleteWalFilesPeriodInMs(100);
+ simpleEnv
+ .getConfig()
+ .getCommonConfig()
+ .setDataReplicationFactor(3)
+ .setSchemaReplicationFactor(3)
+ .setSchemaRegionConsensusProtocolClass("org.apache.iotdb.consensus.ratis.RatisConsensus")
+ .setDataRegionConsensusProtocolClass("org.apache.iotdb.consensus.iot.IoTConsensus");
+ try {
+ simpleEnv.initClusterEnvironment(1, 3);
+
+ int leaderIndex = -1;
+ try (Connection connection = simpleEnv.getConnection();
+ Statement statement = connection.createStatement()) {
+ // write the first data
+ statement.execute("INSERT INTO root.db1.d1 (time, s1) values (1,1)");
+ // find the leader of the data region
+ int port = -1;
+
+ ResultSet resultSet = statement.executeQuery("SHOW REGIONS");
+ while (resultSet.next()) {
+ String regionType = resultSet.getString("Type");
+ if (regionType.equals("DataRegion")) {
+ String role = resultSet.getString("Role");
+ if (role.equals("Leader")) {
+ port = resultSet.getInt("RpcPort");
+ break;
+ }
+ }
+ }
+
+ if (port == -1) {
+ fail("Leader not found");
+ }
+
+ for (int i = 0; i < simpleEnv.getDataNodeWrapperList().size(); i++) {
+ if (simpleEnv.getDataNodeWrapperList().get(i).getPort() == port) {
+ leaderIndex = i;
+ break;
+ }
+ }
+ }
+
+ // stop a follower
+ int followerIndex = (leaderIndex + 1) % simpleEnv.getDataNodeWrapperList().size();
+ simpleEnv.getDataNodeWrapperList().get(followerIndex).stop();
+ System.out.println(
+ new Date()
+ + ":Stopping data node "
+ + simpleEnv.getDataNodeWrapperList().get(followerIndex).getIpAndPortString());
+
+ DataNodeWrapper leader = simpleEnv.getDataNodeWrapperList().get(leaderIndex);
+ // write to the leader to generate wal that cannot be synced
+ try (Session session = new Session(leader.getIp(), leader.getPort())) {
+ session.open();
+
+ session.executeNonQueryStatement("INSERT INTO root.db1.d1 (time, s1) values (1,1)");
+ session.executeNonQueryStatement("INSERT INTO root.db1.d1 (time, s1) values (1,1)");
+ session.executeNonQueryStatement("INSERT INTO root.db1.d1 (time, s1) values (1,1)");
+ session.executeNonQueryStatement("INSERT INTO root.db1.d1 (time, s1) values (1,1)");
+ session.executeNonQueryStatement("INSERT INTO root.db1.d1 (time, s1) values (1,1)");
+ }
+
+ // wait for wal-delete thread to be scheduled
+ Thread.sleep(1000);
+
+ // stop the leader
+ leader.getInstance().destroy();
+ System.out.println(new Date() + ":Stopping data node " + leader.getIpAndPortString());
+ // confirm the death of the leader
+ long startTime = System.currentTimeMillis();
+ while (leader.isAlive()) {
+ if (System.currentTimeMillis() - startTime > 30000) {
+ fail("Leader does not exit after 30s");
+ }
+ }
+ } finally {
+ simpleEnv.cleanClusterEnvironment();
+ }
+ }
+}

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

@@ -29,6 +29,7 @@ public class DeserializedBatchIndexedConsensusRequest
  private final long startSyncIndex;
  private final long endSyncIndex;
  private final List<IConsensusRequest> insertNodes;
+ private long memorySize;
 
  public DeserializedBatchIndexedConsensusRequest(
  long startSyncIndex, long endSyncIndex, int size) {
@@ -52,6 +53,7 @@ public List<IConsensusRequest> getInsertNodes() {
 
  public void add(IConsensusRequest insertNode) {
  this.insertNodes.add(insertNode);
+ this.memorySize += insertNode.getMemorySize();
  }
 
  @Override
@@ -82,4 +84,9 @@ public int hashCode() {
  public ByteBuffer serializeToByteBuffer() {
  return null;
  }
+
+ @Override
+ public long getMemorySize() {
+ return memorySize;
+ }
 }

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

@@ -23,6 +23,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicLong;
 
 /** only used for iot consensus. */
 public class IndexedConsensusRequest implements IConsensusRequest {
@@ -34,6 +35,7 @@ public class IndexedConsensusRequest implements IConsensusRequest {
  private final List<IConsensusRequest> requests;
  private final List<ByteBuffer> serializedRequests;
  private long memorySize = 0;
+ private AtomicLong referenceCnt = new AtomicLong();
 
  public IndexedConsensusRequest(long searchIndex, List<IConsensusRequest> requests) {
  this.searchIndex = searchIndex;
@@ -100,4 +102,12 @@ public boolean equals(Object o) {
  public int hashCode() {
  return Objects.hash(searchIndex, requests);
  }
+
+ public long incRef() {
+ return referenceCnt.getAndIncrement();
+ }
+
+ public long decRef() {
+ return referenceCnt.getAndDecrement();
+ }
 }

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/client/DispatchLogHandler.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.consensus.iot.logdispatcher.Batch;
+import org.apache.iotdb.consensus.iot.logdispatcher.LogDispatcher;
 import org.apache.iotdb.consensus.iot.logdispatcher.LogDispatcher.LogDispatcherThread;
 import org.apache.iotdb.consensus.iot.logdispatcher.LogDispatcherThreadMetrics;
 import org.apache.iotdb.consensus.iot.thrift.TSyncLogEntriesRes;
@@ -89,6 +90,10 @@ public void onComplete(TSyncLogEntriesRes response) {
  }
  completeBatch(batch);
  }
+ if (response.isSetReceiverMemSize()) {
+ LogDispatcher.getReceiverMemSizeSum().addAndGet(response.getReceiverMemSize());
+ LogDispatcher.getSenderMemSizeSum().addAndGet(batch.getMemorySize());
+ }
  logDispatcherThreadMetrics.recordSyncLogTimePerRequest(System.nanoTime() - createTime);
  }
 

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

@@ -22,7 +22,6 @@
 import org.apache.iotdb.consensus.config.IoTConsensusConfig;
 import org.apache.iotdb.consensus.iot.thrift.TLogEntry;
 
-import java.nio.Buffer;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -37,7 +36,7 @@ public class Batch {
 
  private long logEntriesNumFromWAL = 0L;
 
- private long serializedSize;
+ private long memorySize;
  // indicates whether this batch has been successfully synchronized to another node
  private boolean synced;
 
@@ -60,14 +59,20 @@ public void addTLogEntry(TLogEntry entry) {
  if (entry.fromWAL) {
  logEntriesNumFromWAL++;
  }
- // TODO Maybe we need to add in additional fields for more accurate calculations
- serializedSize +=
- entry.getData() == null ? 0 : entry.getData().stream().mapToInt(Buffer::capacity).sum();
+ memorySize += entry.getMemorySize();
  }
 
  public boolean canAccumulate() {
+ // When reading entries from the WAL, the memory size is calculated based on the serialized
+ // size, which can be significantly smaller than the actual size.
+ // Thus, we add a multiplier to sender's memory size to estimate the receiver's memory cost.
+ // The multiplier is calculated based on the receiver's feedback.
+ long receiverMemSize = LogDispatcher.getReceiverMemSizeSum().get();
+ long senderMemSize = LogDispatcher.getSenderMemSizeSum().get();
+ double multiplier = senderMemSize > 0 ? (double) receiverMemSize / senderMemSize : 1.0;
+ multiplier = Math.max(multiplier, 1.0);
  return logEntries.size() < config.getReplication().getMaxLogEntriesNumPerBatch()
- && serializedSize < config.getReplication().getMaxSizePerBatch();
+ && ((long) (memorySize * multiplier)) < config.getReplication().getMaxSizePerBatch();
  }
 
  public long getStartIndex() {
@@ -94,8 +99,8 @@ public boolean isEmpty() {
  return logEntries.isEmpty();
  }
 
- public long getSerializedSize() {
- return serializedSize;
+ public long getMemorySize() {
+ return memorySize;
  }
 
  public long getLogEntriesNumFromWAL() {
@@ -111,8 +116,8 @@ public String toString() {
  + endIndex
  + ", size="
  + logEntries.size()
- + ", serializedSize="
- + serializedSize
+ + ", memorySize="
+ + memorySize
  + '}';
  }
 }