add unittest 1. doNotSwith when no common block between the current main chain with the fork chain, 2. switchBack when switch fork failed

Author: Yrp

src/main/java/org/tron/core/db/Manager.java

@@ -618,6 +618,8 @@ private void switchFork(BlockCapsule newHead)
           throw e;
         } finally {
           if (exception != null) {
+            logger.warn("switch back because exception thrown while switching forks. " + exception.getMessage(),
+                exception);
             first.forEach(khaosBlock -> khaosDb.removeBlk(khaosBlock.getBlk().getBlockId()));
             khaosDb.setHead(binaryTree.getValue().peekFirst());
 

src/test/java/org/tron/core/db/ManagerTest.java

@@ -7,8 +7,10 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import lombok.extern.slf4j.Slf4j;
+import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
@@ -53,13 +55,10 @@ public class ManagerTest {
   private static BlockCapsule blockCapsule2;
   private static String dbPath = "output_manager_test";
 
-  static {
+  @Before
+  public void init() {
     Args.setParam(new String[]{"-d", dbPath, "-w"}, Constant.TEST_CONF);
     context = new AnnotationConfigApplicationContext(DefaultConfig.class);
-  }
-
-  @BeforeClass
-  public static void init() {
 
     dbManager = context.getBean(Manager.class);
 
@@ -80,8 +79,8 @@ public static void init() {
         ByteArray.fromHexString(Args.getInstance().getLocalWitnesses().getPrivateKey()));
   }
 
-  @AfterClass
-  public static void removeDb() {
+  @After
+  public void removeDb() {
     Args.clearParam();
     FileUtil.deleteDir(new File(dbPath));
     context.destroy();
@@ -275,6 +274,104 @@ public void fork()
         dbManager.getDynamicPropertiesStore().getLatestBlockHeaderHash());
   }
 
+  @Test
+  public void doNotSwitch()
+      throws ValidateSignatureException, ContractValidateException, ContractExeException,
+      UnLinkedBlockException, ValidateScheduleException, BadItemException,
+      ItemNotFoundException, HeaderNotFound, AccountResourceInsufficientException,
+      TransactionExpirationException, TooBigTransactionException,
+      DupTransactionException, BadBlockException,
+      TaposException, BadNumberBlockException, NonCommonBlockException {
+    Args.setParam(new String[]{"--witness"}, Constant.TEST_CONF);
+    long size = dbManager.getBlockStore().dbSource.allKeys().size();
+    System.out.print("block store size:" + size + "\n");
+    String key = "f31db24bfbd1a2ef19beddca0a0fa37632eded9ac666a05d3bd925f01dde1f62";
+    byte[] privateKey = ByteArray.fromHexString(key);
+    final ECKey ecKey = ECKey.fromPrivate(privateKey);
+    byte[] address = ecKey.getAddress();
+    WitnessCapsule witnessCapsule = new WitnessCapsule(ByteString.copyFrom(address));
+    dbManager.addWitness(ByteString.copyFrom(address));
+    dbManager.generateBlock(witnessCapsule, System.currentTimeMillis(), privateKey);
+
+    Map<ByteString, String> addressToProvateKeys = addTestWitnessAndAccount();
+
+    long num = dbManager.getDynamicPropertiesStore().getLatestBlockHeaderNumber();
+    BlockCapsule blockCapsule0 =
+        createTestBlockCapsule(
+            num + 1,
+            dbManager.getDynamicPropertiesStore().getLatestBlockHeaderHash().getByteString(),
+            addressToProvateKeys);
+
+    BlockCapsule blockCapsule1 =
+        createTestBlockCapsule(
+            num + 1,
+            dbManager.getDynamicPropertiesStore().getLatestBlockHeaderHash().getByteString(),
+            addressToProvateKeys);
+
+    BlockCapsule blockCapsule2 =
+        createTestBlockCapsule(
+            num + 2, blockCapsule1.getBlockId().getByteString(), addressToProvateKeys);
+
+    dbManager.pushBlock(blockCapsule0);
+    dbManager.pushBlock(blockCapsule1);
+    context.getBean(KhaosDatabase.class).removeBlk(dbManager.getBlockIdByNum(num));
+    try {
+      dbManager.pushBlock(blockCapsule2);
+    } catch (NonCommonBlockException e) {
+      logger.info("do not switch fork");
+      Assert.assertNotNull(dbManager.getBlockStore().get(blockCapsule0.getBlockId().getBytes()));
+    }
+  }
+
+  @Test
+  public void switchBack()
+      throws ValidateSignatureException, ContractValidateException, ContractExeException,
+      UnLinkedBlockException, ValidateScheduleException, BadItemException,
+      ItemNotFoundException, HeaderNotFound, AccountResourceInsufficientException,
+      TransactionExpirationException, TooBigTransactionException,
+      DupTransactionException, BadBlockException,
+      TaposException, BadNumberBlockException, NonCommonBlockException {
+    Args.setParam(new String[]{"--witness"}, Constant.TEST_CONF);
+    long size = dbManager.getBlockStore().dbSource.allKeys().size();
+    System.out.print("block store size:" + size + "\n");
+    String key = "f31db24bfbd1a2ef19beddca0a0fa37632eded9ac666a05d3bd925f01dde1f62";
+    byte[] privateKey = ByteArray.fromHexString(key);
+    final ECKey ecKey = ECKey.fromPrivate(privateKey);
+    byte[] address = ecKey.getAddress();
+    WitnessCapsule witnessCapsule = new WitnessCapsule(ByteString.copyFrom(address));
+    dbManager.addWitness(ByteString.copyFrom(address));
+    dbManager.generateBlock(witnessCapsule, System.currentTimeMillis(), privateKey);
+
+    Map<ByteString, String> addressToProvateKeys = addTestWitnessAndAccount();
+
+    long num = dbManager.getDynamicPropertiesStore().getLatestBlockHeaderNumber();
+    BlockCapsule blockCapsule0 =
+        createTestBlockCapsule(
+            num + 1,
+            dbManager.getDynamicPropertiesStore().getLatestBlockHeaderHash().getByteString(),
+            addressToProvateKeys);
+
+    BlockCapsule blockCapsule1 =
+        createTestBlockCapsule(
+            num + 1,
+            dbManager.getDynamicPropertiesStore().getLatestBlockHeaderHash().getByteString(),
+            addressToProvateKeys);
+
+    BlockCapsule blockCapsule2 =
+        createTestBlockCapsuleError(
+            num + 2, blockCapsule1.getBlockId().getByteString(), addressToProvateKeys);
+
+    dbManager.pushBlock(blockCapsule0);
+    dbManager.pushBlock(blockCapsule1);
+    try {
+      dbManager.pushBlock(blockCapsule2);
+    } catch (Exception e) {
+      logger.info("the fork chain has error block");
+    }
+
+    Assert.assertNotNull(dbManager.getBlockStore().get(blockCapsule0.getBlockId().getBytes()));
+  }
+
   private Map<ByteString, String> addTestWitnessAndAccount() {
     dbManager.getWitnesses().clear();
     return IntStream.range(0, 2)
@@ -309,4 +406,17 @@ private BlockCapsule createTestBlockCapsule(
     blockCapsule.sign(ByteArray.fromHexString(addressToProvateKeys.get(witnessAddress)));
     return blockCapsule;
   }
+
+  private BlockCapsule createTestBlockCapsuleError(
+      long number, ByteString hash, Map<ByteString, String> addressToProvateKeys) {
+    long time = System.currentTimeMillis();
+    WitnessController witnessController = dbManager.getWitnessController();
+    ByteString witnessAddress =
+        witnessController.getScheduledWitness(witnessController.getSlotAtTime(time));
+    BlockCapsule blockCapsule = new BlockCapsule(number, Sha256Hash.wrap(hash), time, ByteString.copyFromUtf8("onlyTest"));
+    blockCapsule.generatedByMyself = true;
+    blockCapsule.setMerkleRoot();
+    blockCapsule.sign(ByteArray.fromHexString(addressToProvateKeys.get(witnessAddress)));
+    return blockCapsule;
+  }
 }