feat(net):Optimize rate limiting logic and add tests

Author: jwrct

framework/src/main/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandler.java

@@ -41,7 +41,6 @@ public class FetchInvDataMsgHandler implements TronMsgHandler {
       .maximumSize(1000).expireAfterWrite(1, TimeUnit.HOURS).build();
 
   private static final int MAX_SIZE = 1_000_000;
-  private static final int MAX_FETCH_SIZE = 100;
   @Autowired
   private TronNetDelegate tronNetDelegate;
   @Autowired
@@ -56,13 +55,6 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep
 
     FetchInvDataMessage fetchInvDataMsg = (FetchInvDataMessage) msg;
 
-    if (peer.isNeedSyncFromUs() && !peer.getP2pRateLimiter().tryAcquire(msg.getType().asByte())) {
-      // Discard messages that exceed the rate limit
-      logger.warn("{} message from peer {} exceeds the rate limit",
-          msg.getType(), peer.getInetSocketAddress());
-      return;
-    }
-
     check(peer, fetchInvDataMsg);
 
     InventoryType type = fetchInvDataMsg.getInventoryType();
@@ -164,7 +156,11 @@ private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg) thr
         if (!peer.isNeedSyncFromUs()) {
           throw new P2pException(TypeEnum.BAD_MESSAGE, "no need sync");
         }
-        if (fetchInvDataMsg.getHashList().size() > MAX_FETCH_SIZE) {
+        if (!peer.getP2pRateLimiter().tryAcquire(fetchInvDataMsg.getType().asByte())) {
+          throw new P2pException(TypeEnum.BAD_MESSAGE, fetchInvDataMsg.getType()
+              + " message exceeds the rate limit");
+        }
+        if (fetchInvDataMsg.getHashList().size() > NetConstants.MAX_BLOCK_FETCH_PER_PEER) {
           throw new P2pException(TypeEnum.BAD_MESSAGE, "fetch too more blocks, size:"
               + fetchInvDataMsg.getHashList().size());
         }

framework/src/main/java/org/tron/core/net/peer/PeerConnection.java

@@ -59,6 +59,10 @@ public class PeerConnection {
 
   private static List<InetSocketAddress> relayNodes = Args.getInstance().getFastForwardNodes();
 
+  private static final double SYNC_BLOCK_CHAIN_RATE = 3.0;
+  private static final double FETCH_INV_DATA_RATE = 3.0;
+  private static final double P2P_DISCONNECT_RATE = 1.0;
+
   @Getter
   private PeerStatistics peerStatistics = new PeerStatistics();
 
@@ -171,9 +175,9 @@ public void setChannel(Channel channel) {
     }
     this.nodeStatistics = TronStatsManager.getNodeStatistics(channel.getInetAddress());
     lastInteractiveTime = System.currentTimeMillis();
-    p2pRateLimiter.register(SYNC_BLOCK_CHAIN.asByte(), 2);
-    p2pRateLimiter.register(FETCH_INV_DATA.asByte(), 2);
-    p2pRateLimiter.register(P2P_DISCONNECT.asByte(), 1);
+    p2pRateLimiter.register(SYNC_BLOCK_CHAIN.asByte(), SYNC_BLOCK_CHAIN_RATE);
+    p2pRateLimiter.register(FETCH_INV_DATA.asByte(), FETCH_INV_DATA_RATE);
+    p2pRateLimiter.register(P2P_DISCONNECT.asByte(), P2P_DISCONNECT_RATE);
   }
 
   public void setBlockBothHave(BlockId blockId) {

framework/src/test/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandlerTest.java

@@ -58,6 +58,9 @@ public void testProcessMessage() throws Exception {
     Mockito.when(advService.getMessage(new Item(blockId, Protocol.Inventory.InventoryType.BLOCK)))
         .thenReturn(new BlockMessage(blockCapsule));
     ReflectUtils.setFieldValue(fetchInvDataMsgHandler, "advService", advService);
+    P2pRateLimiter p2pRateLimiter = new P2pRateLimiter();
+    p2pRateLimiter.register(FETCH_INV_DATA.asByte(), 2);
+    Mockito.when(peer.getP2pRateLimiter()).thenReturn(p2pRateLimiter);
 
     fetchInvDataMsgHandler.processMessage(peer,
         new FetchInvDataMessage(blockIds, Protocol.Inventory.InventoryType.BLOCK));
@@ -77,6 +80,9 @@ public void testSyncFetchCheck() {
     Cache<Item, Long> advInvSpread = CacheBuilder.newBuilder().maximumSize(100)
         .expireAfterWrite(1, TimeUnit.HOURS).recordStats().build();
     Mockito.when(peer.getAdvInvSpread()).thenReturn(advInvSpread);
+    P2pRateLimiter p2pRateLimiter = new P2pRateLimiter();
+    p2pRateLimiter.register(FETCH_INV_DATA.asByte(), 2);
+    Mockito.when(peer.getP2pRateLimiter()).thenReturn(p2pRateLimiter);
 
     FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler();
 
@@ -113,6 +119,7 @@ public void testRateLimiter() {
     Mockito.when(peer.getAdvInvSpread()).thenReturn(advInvSpread);
     P2pRateLimiter p2pRateLimiter = new P2pRateLimiter();
     p2pRateLimiter.register(FETCH_INV_DATA.asByte(), 1);
+    p2pRateLimiter.acquire(FETCH_INV_DATA.asByte());
     Mockito.when(peer.getP2pRateLimiter()).thenReturn(p2pRateLimiter);
     FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler();
 

framework/src/test/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandlerTest.java

@@ -55,6 +55,7 @@ public void init() throws Exception {
   @Test
   public void testProcessMessage() throws Exception {
     try {
+      peer.setRemainNum(1);
       handler.processMessage(peer, new SyncBlockChainMessage(new ArrayList<>()));
     } catch (P2pException e) {
       Assert.assertEquals("SyncBlockChain blockIds is empty", e.getMessage());
@@ -71,6 +72,10 @@ public void testProcessMessage() throws Exception {
     Assert.assertNotNull(message.toString());
     Assert.assertNotNull(((BlockInventoryMessage) message).getAnswerMessage());
     Assert.assertFalse(f);
+    method.invoke(handler, peer, message);
+    method.invoke(handler, peer, message);
+    f = (boolean)method.invoke(handler, peer, message);
+    Assert.assertFalse(f);
 
     Method method1 = handler.getClass().getDeclaredMethod(
         "getLostBlockIds", List.class, BlockId.class);