feat(net): optimize block processing logic (#5754)

feat(net): optimize block processing logic

Author: fyyhtx

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

@@ -125,6 +125,14 @@ private void check(PeerConnection peer, BlockMessage msg) throws P2pException {
 
   private void processBlock(PeerConnection peer, BlockCapsule block) throws P2pException {
     BlockId blockId = block.getBlockId();
+    boolean flag = tronNetDelegate.validBlock(block);
+    if (!flag) {
+      logger.warn("Receive a bad block from {}, {}, {}",
+          peer.getInetSocketAddress(), blockId.getString(),
+          Hex.toHexString(block.getWitnessAddress().toByteArray()));
+      return;
+    }
+
     if (!tronNetDelegate.containBlock(block.getParentBlockId())) {
       logger.warn("Get unlink block {} from {}, head is {}", blockId.getString(),
               peer.getInetAddress(), tronNetDelegate.getHeadBlockId().getString());
@@ -138,16 +146,10 @@ private void processBlock(PeerConnection peer, BlockCapsule block) throws P2pExc
       return;
     }
 
-    boolean flag = tronNetDelegate.validBlock(block);
-    if (flag) {
-      broadcast(new BlockMessage(block));
-    }
+    broadcast(new BlockMessage(block));
 
     try {
       tronNetDelegate.processBlock(block, false);
-      if (!flag) {
-        broadcast(new BlockMessage(block));
-      }
 
       witnessProductBlockService.validWitnessProductTwoBlock(block);
 

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

@@ -1,5 +1,7 @@
 package org.tron.core.net.messagehandler;
 
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import com.google.common.collect.Sets;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
@@ -8,12 +10,11 @@
 import java.security.SignatureException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
@@ -37,7 +38,9 @@
 @Service
 public class PbftDataSyncHandler implements TronMsgHandler, Closeable {
 
-  private Map<Long, PbftCommitMessage> pbftCommitMessageCache = new ConcurrentHashMap<>();
+  private static final Cache<Long, PbftCommitMessage> pbftCommitMessageCache =
+      CacheBuilder.newBuilder().initialCapacity(100).maximumSize(200)
+          .expireAfterWrite(10, TimeUnit.MINUTES).build();
 
   private final String esName = "valid-header-pbft-sign";
 
@@ -51,6 +54,9 @@ public class PbftDataSyncHandler implements TronMsgHandler, Closeable {
   public void processMessage(PeerConnection peer, TronMessage msg) throws P2pException {
     PbftCommitMessage pbftCommitMessage = (PbftCommitMessage) msg;
     try {
+      if (!chainBaseManager.getDynamicPropertiesStore().allowPBFT()) {
+        return;
+      }
       Raw raw = Raw.parseFrom(pbftCommitMessage.getPBFTCommitResult().getData());
       pbftCommitMessageCache.put(raw.getViewN(), pbftCommitMessage);
     } catch (InvalidProtocolBufferException e) {
@@ -64,7 +70,8 @@ public void processPBFTCommitData(BlockCapsule block) {
         return;
       }
       long epoch = 0;
-      PbftCommitMessage pbftCommitMessage = pbftCommitMessageCache.remove(block.getNum());
+      PbftCommitMessage pbftCommitMessage = pbftCommitMessageCache.getIfPresent(block.getNum());
+      pbftCommitMessageCache.invalidate(block.getNum());
       long maintenanceTimeInterval = chainBaseManager.getDynamicPropertiesStore()
           .getMaintenanceTimeInterval();
       if (pbftCommitMessage == null) {
@@ -75,7 +82,8 @@ public void processPBFTCommitData(BlockCapsule block) {
         Raw raw = Raw.parseFrom(pbftCommitMessage.getPBFTCommitResult().getData());
         epoch = raw.getEpoch();
       }
-      pbftCommitMessage = pbftCommitMessageCache.remove(epoch);
+      pbftCommitMessage = pbftCommitMessageCache.getIfPresent(epoch);
+      pbftCommitMessageCache.invalidate(epoch);
       if (pbftCommitMessage != null) {
         processPBFTCommitMessage(pbftCommitMessage);
       }

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

@@ -33,6 +33,9 @@ public class PbftMsgHandler {
   private TronNetDelegate tronNetDelegate;
 
   public void processMessage(PeerConnection peer, PbftMessage msg) throws Exception {
+    if (!tronNetDelegate.allowPBFT()) {
+      return;
+    }
     if (Param.getInstance().getPbftInterface().isSyncing()) {
       return;
     }

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

@@ -32,9 +32,6 @@ public void testProcessMessage() throws Exception {
     Protocol.PBFTMessage.Raw raw = rawBuilder.build();
     PbftSignCapsule pbftSignCapsule = new PbftSignCapsule(raw.toByteString(), new ArrayList<>());
     PbftCommitMessage pbftCommitMessage = new PbftCommitMessage(pbftSignCapsule);
-    pbftDataSyncHandler.processMessage(null, pbftCommitMessage);
-    Assert.assertEquals(Protocol.PBFTMessage.Raw.parseFrom(
-        pbftCommitMessage.getPBFTCommitResult().getData()).getViewN(), 1);
 
     DynamicPropertiesStore dynamicPropertiesStore = Mockito.mock(DynamicPropertiesStore.class);
     PbftSignDataStore pbftSignDataStore = Mockito.mock(PbftSignDataStore.class);
@@ -48,6 +45,10 @@ public void testProcessMessage() throws Exception {
     field.setAccessible(true);
     field.set(pbftDataSyncHandler, chainBaseManager);
 
+    pbftDataSyncHandler.processMessage(null, pbftCommitMessage);
+    Assert.assertEquals(Protocol.PBFTMessage.Raw.parseFrom(
+        pbftCommitMessage.getPBFTCommitResult().getData()).getViewN(), 1);
+
     pbftDataSyncHandler.processPBFTCommitData(blockCapsule);
     Field field1 = PbftDataSyncHandler.class.getDeclaredField("pbftCommitMessageCache");
     field1.setAccessible(true);

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

@@ -34,6 +34,7 @@
 import org.tron.core.net.message.MessageTypes;
 import org.tron.core.net.peer.PeerConnection;
 import org.tron.core.net.peer.PeerManager;
+import org.tron.core.store.DynamicPropertiesStore;
 import org.tron.p2p.P2pConfig;
 import org.tron.p2p.base.Parameter;
 import org.tron.p2p.connection.Channel;
@@ -116,6 +117,14 @@ public void testPbft() throws Exception {
       Assert.assertEquals(P2pException.TypeEnum.BAD_MESSAGE, e.getType());
     }
 
+    DynamicPropertiesStore dynamicPropertiesStore = context.getBean(DynamicPropertiesStore.class);
+    dynamicPropertiesStore.saveAllowPBFT(1);
+    try {
+      context.getBean(PbftMsgHandler.class).processMessage(peer, pbftMessage);
+    } catch (P2pException e) {
+      Assert.assertEquals(P2pException.TypeEnum.BAD_MESSAGE, e.getType());
+    }
+
     Assert.assertEquals(1, PeerManager.getPeers().size());
   }
 }