Extend region migration speed limit to transfers and deletion

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

@@ -36,6 +36,7 @@
 import org.apache.iotdb.commons.utils.KillPoint.IoTConsensusDeleteLocalPeerKillPoints;
 import org.apache.iotdb.commons.utils.KillPoint.IoTConsensusRemovePeerCoordinatorKillPoints;
 import org.apache.iotdb.commons.utils.KillPoint.KillPoint;
+import org.apache.iotdb.commons.utils.RegionMigrationRateLimiter;
 import org.apache.iotdb.commons.utils.StatusUtils;
 import org.apache.iotdb.consensus.IConsensus;
 import org.apache.iotdb.consensus.IStateMachine;
@@ -378,7 +379,9 @@ public void deleteLocalPeer(ConsensusGroupId groupId) throws ConsensusException
     if (!exist.get()) {
       throw new ConsensusGroupNotExistException(groupId);
     }
-    FileUtils.deleteFileOrDirectory(new File(buildPeerDir(storageDir, groupId)));
+    FileUtils.deleteFileOrDirectoryWithRateLimiter(
+        new File(buildPeerDir(storageDir, groupId)),
+        RegionMigrationRateLimiter.getInstance()::acquire);
     KillPoint.setKillPoint(IoTConsensusDeleteLocalPeerKillPoints.AFTER_DELETE);
   }
 

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/snapshot/IoTConsensusRateLimiter.java

@@ -19,22 +19,16 @@
 
 package org.apache.iotdb.consensus.iot.snapshot;
 
-import com.google.common.util.concurrent.RateLimiter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import org.apache.iotdb.commons.utils.RegionMigrationRateLimiter;
 
 public class IoTConsensusRateLimiter {
-  private static final Logger logger = LoggerFactory.getLogger(IoTConsensusRateLimiter.class);
 
-  private final RateLimiter rateLimiter = RateLimiter.create(Double.MAX_VALUE);
+  private final RegionMigrationRateLimiter rateLimiter = RegionMigrationRateLimiter.getInstance();
 
   private IoTConsensusRateLimiter() {}
 
   public void init(long regionMigrationSpeedLimitBytesPerSecond) {
-    rateLimiter.setRate(
-        regionMigrationSpeedLimitBytesPerSecond <= 0
-            ? Double.MAX_VALUE
-            : regionMigrationSpeedLimitBytesPerSecond);
+    rateLimiter.init(regionMigrationSpeedLimitBytesPerSecond);
   }
 
   /**
@@ -43,15 +37,7 @@ public void init(long regionMigrationSpeedLimitBytesPerSecond) {
    * @param transitDataSize the size of the data to be sent
    */
   public void acquireTransitDataSizeWithRateLimiter(long transitDataSize) {
-    while (transitDataSize > 0) {
-      if (transitDataSize > Integer.MAX_VALUE) {
-        rateLimiter.acquire(Integer.MAX_VALUE);
-        transitDataSize -= Integer.MAX_VALUE;
-      } else {
-        rateLimiter.acquire((int) transitDataSize);
-        return;
-      }
-    }
+    rateLimiter.acquire(transitDataSize);
   }
 
   private static final IoTConsensusRateLimiter INSTANCE = new IoTConsensusRateLimiter();

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/pipe/IoTConsensusV2.java

@@ -34,6 +34,7 @@
 import org.apache.iotdb.commons.utils.KillPoint.IoTConsensusDeleteLocalPeerKillPoints;
 import org.apache.iotdb.commons.utils.KillPoint.IoTConsensusRemovePeerCoordinatorKillPoints;
 import org.apache.iotdb.commons.utils.KillPoint.KillPoint;
+import org.apache.iotdb.commons.utils.RegionMigrationRateLimiter;
 import org.apache.iotdb.commons.utils.StatusUtils;
 import org.apache.iotdb.consensus.IConsensus;
 import org.apache.iotdb.consensus.IStateMachine;
@@ -323,7 +324,8 @@ public void deleteLocalPeer(ConsensusGroupId groupId) throws ConsensusException
         consensus.clear();
         stateMachineMap.remove(groupId);
 
-        FileUtils.deleteFileOrDirectory(new File(getPeerDir(groupId)));
+        FileUtils.deleteFileOrDirectoryWithRateLimiter(
+            new File(getPeerDir(groupId)), RegionMigrationRateLimiter.getInstance()::acquire);
         KillPoint.setKillPoint(IoTConsensusDeleteLocalPeerKillPoints.AFTER_DELETE);
         LOGGER.info(IoTConsensusV2Messages.FINISH_DELETE_LOCAL_PEER, CLASS_NAME, groupId);
       } finally {

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/ratis/RateLimitedGrpcFactory.java

@@ -0,0 +1,40 @@
+/*
+ * 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.consensus.ratis;
+
+import org.apache.ratis.conf.Parameters;
+import org.apache.ratis.grpc.GrpcFactory;
+import org.apache.ratis.server.RaftServer;
+import org.apache.ratis.server.leader.FollowerInfo;
+import org.apache.ratis.server.leader.LeaderState;
+import org.apache.ratis.server.leader.LogAppender;
+
+class RateLimitedGrpcFactory extends GrpcFactory {
+
+  RateLimitedGrpcFactory(Parameters parameters) {
+    super(parameters);
+  }
+
+  @Override
+  public LogAppender newLogAppender(
+      RaftServer.Division server, LeaderState leaderState, FollowerInfo follower) {
+    return new RateLimitedGrpcLogAppender(server, leaderState, follower);
+  }
+}

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/ratis/RateLimitedGrpcLogAppender.java

@@ -0,0 +1,76 @@
+/*
+ * 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.consensus.ratis;
+
+import org.apache.iotdb.commons.utils.RegionMigrationRateLimiter;
+
+import org.apache.ratis.grpc.server.GrpcLogAppender;
+import org.apache.ratis.proto.RaftProtos.FileChunkProto;
+import org.apache.ratis.proto.RaftProtos.InstallSnapshotRequestProto;
+import org.apache.ratis.server.RaftServer;
+import org.apache.ratis.server.leader.FollowerInfo;
+import org.apache.ratis.server.leader.LeaderState;
+import org.apache.ratis.statemachine.SnapshotInfo;
+
+import java.util.Iterator;
+
+class RateLimitedGrpcLogAppender extends GrpcLogAppender {
+
+  private final RegionMigrationRateLimiter rateLimiter = RegionMigrationRateLimiter.getInstance();
+
+  RateLimitedGrpcLogAppender(
+      RaftServer.Division server, LeaderState leaderState, FollowerInfo follower) {
+    super(server, leaderState, follower);
+  }
+
+  @Override
+  public Iterable<InstallSnapshotRequestProto> newInstallSnapshotRequests(
+      String requestId, SnapshotInfo snapshot) {
+    final Iterable<InstallSnapshotRequestProto> requests =
+        super.newInstallSnapshotRequests(requestId, snapshot);
+    return () -> {
+      final Iterator<InstallSnapshotRequestProto> iterator = requests.iterator();
+      return new Iterator<InstallSnapshotRequestProto>() {
+        @Override
+        public boolean hasNext() {
+          return iterator.hasNext();
+        }
+
+        @Override
+        public InstallSnapshotRequestProto next() {
+          final InstallSnapshotRequestProto request = iterator.next();
+          rateLimiter.acquire(getSnapshotChunkDataSize(request));
+          return request;
+        }
+      };
+    };
+  }
+
+  static long getSnapshotChunkDataSize(InstallSnapshotRequestProto request) {
+    if (!request.hasSnapshotChunk()) {
+      return 0;
+    }
+
+    return request.getSnapshotChunk().getFileChunksList().stream()
+        .map(FileChunkProto::getData)
+        .mapToLong(data -> data == null ? 0 : data.size())
+        .sum();
+  }
+}

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/ratis/RateLimitedGrpcRpcType.java

@@ -0,0 +1,37 @@
+/*
+ * 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.consensus.ratis;
+
+import org.apache.ratis.conf.Parameters;
+import org.apache.ratis.rpc.RpcFactory;
+import org.apache.ratis.rpc.RpcType;
+
+public class RateLimitedGrpcRpcType implements RpcType {
+
+  @Override
+  public String name() {
+    return RateLimitedGrpcRpcType.class.getName();
+  }
+
+  @Override
+  public RpcFactory newFactory(Parameters parameters) {
+    return new RateLimitedGrpcFactory(parameters);
+  }
+}

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/ratis/RatisConsensus.java

@@ -56,6 +56,7 @@
 import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
+import org.apache.ratis.RaftConfigKeys;
 import org.apache.ratis.client.RaftClientRpc;
 import org.apache.ratis.conf.Parameters;
 import org.apache.ratis.conf.RaftProperties;
@@ -159,6 +160,7 @@ public RatisConsensus(ConsensusConfig config, IStateMachine.Registry registry) {
     this.storageDir = new File(config.getStorageDir());
 
     RaftServerConfigKeys.setStorageDir(properties, Collections.singletonList(storageDir));
+    RaftConfigKeys.Rpc.setType(properties, new RateLimitedGrpcRpcType());
     GrpcConfigKeys.Server.setHost(properties, config.getThisNodeEndPoint().getIp());
     GrpcConfigKeys.Server.setPort(properties, config.getThisNodeEndPoint().getPort());
 

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/simple/SimpleConsensus.java

@@ -26,6 +26,7 @@
 import org.apache.iotdb.commons.request.IConsensusRequest;
 import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics;
 import org.apache.iotdb.commons.utils.FileUtils;
+import org.apache.iotdb.commons.utils.RegionMigrationRateLimiter;
 import org.apache.iotdb.commons.utils.StatusUtils;
 import org.apache.iotdb.consensus.IConsensus;
 import org.apache.iotdb.consensus.IStateMachine;
@@ -197,7 +198,8 @@ public void deleteLocalPeer(ConsensusGroupId groupId) throws ConsensusException
         (k, v) -> {
           exist.set(true);
           v.stop();
-          FileUtils.deleteFileOrDirectory(new File(buildPeerDir(groupId)));
+          FileUtils.deleteFileOrDirectoryWithRateLimiter(
+              new File(buildPeerDir(groupId)), RegionMigrationRateLimiter.getInstance()::acquire);
           return null;
         });
     if (!exist.get()) {

iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/ratis/RateLimitedGrpcLogAppenderTest.java

@@ -0,0 +1,66 @@
+/*
+ * 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.consensus.ratis;
+
+import org.apache.ratis.RaftConfigKeys;
+import org.apache.ratis.conf.Parameters;
+import org.apache.ratis.conf.RaftProperties;
+import org.apache.ratis.proto.RaftProtos.FileChunkProto;
+import org.apache.ratis.proto.RaftProtos.InstallSnapshotRequestProto;
+import org.apache.ratis.rpc.RpcType;
+import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class RateLimitedGrpcLogAppenderTest {
+
+  @Test
+  public void testGetSnapshotChunkDataSize() {
+    final InstallSnapshotRequestProto request =
+        InstallSnapshotRequestProto.newBuilder()
+            .setSnapshotChunk(
+                InstallSnapshotRequestProto.SnapshotChunkProto.newBuilder()
+                    .addFileChunks(
+                        FileChunkProto.newBuilder()
+                            .setData(ByteString.copyFrom(new byte[3]))
+                            .build())
+                    .addFileChunks(
+                        FileChunkProto.newBuilder()
+                            .setData(ByteString.copyFrom(new byte[5]))
+                            .build()))
+            .build();
+
+    Assert.assertEquals(8, RateLimitedGrpcLogAppender.getSnapshotChunkDataSize(request));
+    Assert.assertEquals(
+        0,
+        RateLimitedGrpcLogAppender.getSnapshotChunkDataSize(
+            InstallSnapshotRequestProto.newBuilder().buildPartial()));
+  }
+
+  @Test
+  public void testRateLimitedGrpcRpcTypeIsResolvedByRatis() {
+    final RaftProperties properties = new RaftProperties();
+
+    RaftConfigKeys.Rpc.setType(properties, new RateLimitedGrpcRpcType());
+    final RpcType rpcType = RaftConfigKeys.Rpc.type(properties, ignored -> {});
+
+    Assert.assertTrue(rpcType instanceof RateLimitedGrpcRpcType);
+    Assert.assertTrue(rpcType.newFactory(new Parameters()) instanceof RateLimitedGrpcFactory);
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java

@@ -33,6 +33,7 @@
 import org.apache.iotdb.commons.service.metric.MetricService;
 import org.apache.iotdb.commons.utils.JVMCommonUtils;
 import org.apache.iotdb.commons.utils.NodeUrlUtils;
+import org.apache.iotdb.commons.utils.RegionMigrationRateLimiter;
 import org.apache.iotdb.confignode.rpc.thrift.TCQConfig;
 import org.apache.iotdb.confignode.rpc.thrift.TGlobalConfig;
 import org.apache.iotdb.confignode.rpc.thrift.TRatisConfig;
@@ -1290,6 +1291,8 @@ private void loadIoTConsensusProps(TrimProperties properties) throws IOException
                 "region_migration_speed_limit_bytes_per_second",
                 ConfigurationFileUtils.getConfigurationDefaultValue(
                     "region_migration_speed_limit_bytes_per_second"))));
+    RegionMigrationRateLimiter.getInstance()
+        .init(conf.getRegionMigrationSpeedLimitBytesPerSecond());
     conf.setDataRegionIotSnapshotTransmissionProgressLogIntervalMs(
         Long.parseLong(
             properties.getProperty(