[ETCM-377] Fix consistency issue between cache and persisted block number on rollbacks

Author: Nicolas Tallar

src/main/scala/io/iohk/ethereum/domain/Blockchain.scala

@@ -224,7 +224,12 @@ class BlockchainImpl(
   // There is always only one writer thread (ensured by actor), but can by many readers (api calls)
   // to ensure visibility of writes, needs to be volatile or atomic ref
   private val bestKnownBlockAndLatestCheckpoint: AtomicReference[BestBlockLatestCheckpointNumbers] =
-    new AtomicReference(BestBlockLatestCheckpointNumbers(BigInt(0), BigInt(0)))
+    new AtomicReference(
+      BestBlockLatestCheckpointNumbers(
+        appStateStorage.getBestBlockNumber(),
+        appStateStorage.getLatestCheckpointBlockNumber()
+      )
+    )
 
   override def getBlockHeaderByHash(hash: ByteString): Option[BlockHeader] =
     blockHeadersStorage.get(hash)
@@ -246,22 +251,15 @@ class BlockchainImpl(
       bestSavedBlockNumber,
       bestKnownBlockNumber
     )
-    if (bestKnownBlockNumber > bestSavedBlockNumber)
-      bestKnownBlockNumber
-    else
-      bestSavedBlockNumber
-  }
 
-  override def getLatestCheckpointBlockNumber(): BigInt = {
-    val latestCheckpointNumberInStorage = appStateStorage.getLatestCheckpointBlockNumber()
-    // The latest checkpoint number is firstly saved in memory and then persisted to the storage only when it's time to persist cache.
-    // The latest checkpoint number in memory can be bigger than the number in storage because the cache wasn't persisted yet
-    if (bestKnownBlockAndLatestCheckpoint.get().latestCheckpointNumber > latestCheckpointNumberInStorage)
-      bestKnownBlockAndLatestCheckpoint.get().latestCheckpointNumber
-    else
-      latestCheckpointNumberInStorage
+    // The cached best block number should always be more up-to-date than the one on disk, we are keeping access to disk
+    // above only for logging purposes
+    bestKnownBlockNumber
   }
 
+  override def getLatestCheckpointBlockNumber(): BigInt =
+    bestKnownBlockAndLatestCheckpoint.get().latestCheckpointNumber
+
   override def getBestBlock(): Block = {
     val bestBlockNumber = getBestBlockNumber()
     log.debug("Trying to get best block with number {}", bestBlockNumber)
@@ -313,10 +311,6 @@ class BlockchainImpl(
       .and(storeChainWeight(block.header.hash, weight))
       .commit()
 
-    // not transactional part
-    // the best blocks data will be persisted only when the cache will be persisted
-    stateStorage.onBlockSave(block.header.number, appStateStorage.getBestBlockNumber())(persistBestBlocksData)
-
     if (saveAsBestBlock && block.hasCheckpoint) {
       log.debug(
         "New best known block block number - {}, new best checkpoint number - {}",
@@ -331,6 +325,10 @@ class BlockchainImpl(
       )
       saveBestKnownBlock(block.header.number)
     }
+
+    // not transactional part
+    // the best blocks data will be persisted only when the cache will be persisted
+    stateStorage.onBlockSave(block.header.number, appStateStorage.getBestBlockNumber())(persistBestBlocksData)
   }
 
   override def storeBlockHeader(blockHeader: BlockHeader): DataSourceBatchUpdate = {
@@ -388,87 +386,77 @@ class BlockchainImpl(
     blockNumberMappingStorage.remove(number)
   }
 
-// scalastyle:off method.length
   override def removeBlock(blockHash: ByteString, withState: Boolean): Unit = {
-    val maybeBlockHeader = getBlockHeaderByHash(blockHash)
-
-    log.debug(
-      "Trying to remove block with hash {} and number {}",
-      ByteStringUtils.hash2string(blockHash),
-      maybeBlockHeader.map(_.number)
-    )
+    val maybeBlock = getBlockByHash(blockHash)
 
-    val maybeTxList = getBlockBodyByHash(blockHash).map(_.transactionList)
-    val bestBlocks = bestKnownBlockAndLatestCheckpoint.get()
-    // as we are decreasing block numbers in memory more often than in storage,
-    // we can't use here getBestBlockNumber / getLatestCheckpointBlockNumber
-    val bestBlockNumber =
-      if (bestBlocks.bestBlockNumber != 0) bestBlocks.bestBlockNumber else appStateStorage.getBestBlockNumber()
-    val latestCheckpointNumber = {
-      if (bestBlocks.latestCheckpointNumber != 0) bestBlocks.latestCheckpointNumber
-      else appStateStorage.getLatestCheckpointBlockNumber()
+    maybeBlock match {
+      case Some(block) => removeBlock(block, withState)
+      case None =>
+        log.warn(s"Attempted removing block with hash $blockHash that we don't have")
     }
+  }
 
-    val blockNumberMappingUpdates = {
-      maybeBlockHeader.fold(blockNumberMappingStorage.emptyBatchUpdate)(h =>
-        if (getHashByBlockNumber(h.number).contains(blockHash))
-          removeBlockNumberMapping(h.number)
-        else blockNumberMappingStorage.emptyBatchUpdate
-      )
-    }
+  // scalastyle:off method.length
+  private def removeBlock(block: Block, withState: Boolean): Unit = {
+    val blockHash = block.hash
 
-    val (checkpointUpdates, prevCheckpointNumber): (DataSourceBatchUpdate, Option[BigInt]) = maybeBlockHeader match {
-      case Some(header) =>
-        if (header.hasCheckpoint && header.number == latestCheckpointNumber) {
-          val prev = findPreviousCheckpointBlockNumber(header.number, header.number)
-          prev
-            .map { num =>
-              (appStateStorage.putLatestCheckpointBlockNumber(num), Some(num))
-            }
-            .getOrElse {
-              (appStateStorage.removeLatestCheckpointBlockNumber(), Some(0))
-            }
-        } else (appStateStorage.emptyBatchUpdate, None)
-      case None =>
-        (appStateStorage.emptyBatchUpdate, None)
-    }
+    log.debug(s"Trying to remove block block ${block.idTag}")
 
-    val newBestBlockNumber: BigInt = if (bestBlockNumber >= 1) bestBlockNumber - 1 else 0
+    val txList = block.body.transactionList
+    val bestBlockNumber = getBestBlockNumber()
+    val latestCheckpointNumber = getLatestCheckpointBlockNumber()
+
+    val blockNumberMappingUpdates =
+      if (getHashByBlockNumber(block.number).contains(blockHash))
+        removeBlockNumberMapping(block.number)
+      else blockNumberMappingStorage.emptyBatchUpdate
+
+    val newBestBlockNumber: BigInt = (bestBlockNumber - 1).max(0)
+    val newLatestCheckpointNumber: BigInt =
+      if (block.hasCheckpoint && block.number == latestCheckpointNumber) {
+        findPreviousCheckpointBlockNumber(block.number, block.number)
+      } else latestCheckpointNumber
+
+    // Block number updates are only done if the persisted value is larger, as we won't have the associated mpt nodes if not
+    val bestBlockNumberUpdates =
+      if (appStateStorage.getBestBlockNumber() > newBestBlockNumber)
+        appStateStorage.putBestBlockNumber(newBestBlockNumber)
+      else appStateStorage.emptyBatchUpdate
+
+    // Checkpoint number updates are only done if the persisted value is larger, as we won't have the associated mpt nodes if not
+    val latestCheckpointNumberUpdates =
+      if (appStateStorage.getLatestCheckpointBlockNumber() > newLatestCheckpointNumber)
+        appStateStorage.putLatestCheckpointBlockNumber(newLatestCheckpointNumber)
+      else appStateStorage.emptyBatchUpdate
+
+    log.debug(
+      "Persisting block info data into database. Persisted block number is {}. Persisted checkpoint number is {}",
+      newBestBlockNumber,
+      newLatestCheckpointNumber
+    )
 
     blockHeadersStorage
       .remove(blockHash)
       .and(blockBodiesStorage.remove(blockHash))
       .and(chainWeightStorage.remove(blockHash))
       .and(receiptStorage.remove(blockHash))
-      .and(maybeTxList.fold(transactionMappingStorage.emptyBatchUpdate)(removeTxsLocations))
+      .and(removeTxsLocations(txList))
       .and(blockNumberMappingUpdates)
+      .and(bestBlockNumberUpdates)
+      .and(latestCheckpointNumberUpdates)
       .commit()
 
     // not transactional part
-    saveBestKnownBlocks(newBestBlockNumber, prevCheckpointNumber)
+    saveBestKnownBlocks(newBestBlockNumber, Some(newLatestCheckpointNumber))
     log.debug(
       "Removed block with hash {}. New best block number - {}, new best checkpoint block number - {}",
       ByteStringUtils.hash2string(blockHash),
       newBestBlockNumber,
-      prevCheckpointNumber
+      newLatestCheckpointNumber
     )
 
-    maybeBlockHeader.foreach { h =>
-      if (withState) {
-        val bestBlocksUpdates = appStateStorage
-          .putBestBlockNumber(newBestBlockNumber)
-          .and(checkpointUpdates)
-        stateStorage.onBlockRollback(h.number, bestBlockNumber) { () =>
-          log.debug(
-            "Persisting block info data into database. Persisted block number is {}. " +
-              "Persisted checkpoint number is {}",
-            newBestBlockNumber,
-            prevCheckpointNumber
-          )
-          bestBlocksUpdates.commit()
-        }
-      }
-    }
+    if (withState)
+      stateStorage.onBlockRollback(block.number, bestBlockNumber) { () => persistBestBlocksData() }
   }
   // scalastyle:on method.length
 
@@ -485,7 +473,7 @@ class BlockchainImpl(
   private def findPreviousCheckpointBlockNumber(
       blockNumberToCheck: BigInt,
       latestCheckpointBlockNumber: BigInt
-  ): Option[BigInt] = {
+  ): BigInt = {
     if (blockNumberToCheck > 0) {
       val maybePreviousCheckpointBlockNumber = for {
         currentBlock <- getBlockByNumber(blockNumberToCheck)
@@ -494,10 +482,10 @@ class BlockchainImpl(
       } yield currentBlock.number
 
       maybePreviousCheckpointBlockNumber match {
-        case Some(_) => maybePreviousCheckpointBlockNumber
+        case Some(previousCheckpointBlockNumber) => previousCheckpointBlockNumber
         case None => findPreviousCheckpointBlockNumber(blockNumberToCheck - 1, latestCheckpointBlockNumber)
       }
-    } else None
+    } else 0
   }
 
   private def saveTxsLocations(blockHash: ByteString, blockBody: BlockBody): DataSourceBatchUpdate =

src/test/scala/io/iohk/ethereum/domain/BlockchainSpec.scala

@@ -7,11 +7,15 @@ import io.iohk.ethereum.db.dataSource.EphemDataSource
 import io.iohk.ethereum.db.storage.StateStorage
 import io.iohk.ethereum.domain.BlockHeader.HeaderExtraFields.HefPostEcip1097
 import io.iohk.ethereum.mpt.MerklePatriciaTrie
-import io.iohk.ethereum.{Fixtures, ObjectGenerators}
+import io.iohk.ethereum.{BlockHelpers, Fixtures, ObjectGenerators}
+import io.iohk.ethereum.ObjectGenerators._
+import org.scalacheck.Gen
+import org.scalamock.scalatest.MockFactory
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
+import org.scalatestplus.scalacheck.ScalaCheckPropertyChecks
 
-class BlockchainSpec extends AnyFlatSpec with Matchers {
+class BlockchainSpec extends AnyFlatSpec with Matchers with ScalaCheckPropertyChecks {
 
   val checkpoint = ObjectGenerators.fakeCheckpointGen(2, 5).sample.get
   val checkpointBlockGenerator = new CheckpointBlockGenerator
@@ -127,4 +131,108 @@ class BlockchainSpec extends AnyFlatSpec with Matchers {
     val retrievedAccount = blockchain.getAccount(address, headerWithAcc.number)
     retrievedAccount shouldEqual Some(account)
   }
+
+  it should "return correct best block number after applying and rollbacking blocks" in new TestSetup {
+    forAll(intGen(min = 1: Int, max = maxNumberBlocksToImport)) { numberBlocksToImport =>
+      val testSetup = newSetup()
+      import testSetup._
+
+      // Import blocks
+      val blocksToImport = BlockHelpers.generateChain(numberBlocksToImport, Fixtures.Blocks.Genesis.block)
+
+      // Randomly select the block import to persist (empty means no persistance)
+      val blockImportToPersist = Gen.option(Gen.oneOf(blocksToImport)).sample.get
+      (stubStateStorage
+        .onBlockSave(_: BigInt, _: BigInt)(_: () => Unit))
+        .when(*, *, *)
+        .onCall((bn, _, persistFn) => {
+          if (blockImportToPersist.exists(_.number == bn)) persistFn()
+        })
+
+      blocksToImport.foreach { block =>
+        blockchainWithStubPersisting.save(block, Nil, ChainWeight.zero, true)
+      }
+
+      blockchainWithStubPersisting.getBestBlockNumber() shouldBe blocksToImport.last.number
+      blockchainStoragesWithStubPersisting.appStateStorage.getBestBlockNumber() shouldBe blockImportToPersist.fold(0: BigInt)(_.number)
+
+
+      // Rollback blocks
+      val numberBlocksToRollback = intGen(0, numberBlocksToImport).sample.get
+      val (blocksNotRollbacked, blocksToRollback) = blocksToImport.splitAt(numberBlocksToRollback)
+
+      // Randomly select the block rollback to persist (empty means no persistance)
+      val blockRollbackToPersist = if (blocksToRollback.isEmpty) None else Gen.option(Gen.oneOf(blocksToRollback)).sample.get
+      (stubStateStorage
+        .onBlockRollback(_: BigInt, _: BigInt)(_: () => Unit))
+        .when(*, *, *)
+        .onCall((bn, _, persistFn) => {
+          if (blockRollbackToPersist.exists(_.number == bn)) persistFn()
+        })
+
+      blocksToRollback.reverse.foreach { block =>
+        blockchainWithStubPersisting.removeBlock(block.hash, true)
+      }
+
+      val expectedMemoryBestBlock = blocksNotRollbacked.lastOption.fold(0: BigInt)(_.number)
+      val expectedPersistedBestBlock = calculatePersistedBestBlock(
+        blockImportToPersist.map(_.number),
+        blockRollbackToPersist.map(_.number),
+        blocksToRollback.map(_.number)
+      )
+      blockchainWithStubPersisting.getBestBlockNumber() shouldBe expectedMemoryBestBlock
+      blockchainStoragesWithStubPersisting.appStateStorage.getBestBlockNumber() shouldBe expectedPersistedBestBlock
+    }
+  }
+
+  trait TestSetup extends MockFactory {
+    val maxNumberBlocksToImport: Int = 30
+
+    def calculatePersistedBestBlock(blockImportPersisted: Option[BigInt], blockRollbackPersisted: Option[BigInt], blocksRollbacked: Seq[BigInt]): BigInt = {
+      (blocksRollbacked, blockImportPersisted) match {
+        case (Nil, Some(bi)) =>
+          // No blocks rollbacked, last persist was the persist during import
+          bi
+        case (nonEmptyRollbackedBlocks, Some(bi)) =>
+          // Last forced persist during apply/rollback
+          val maxForcedPersist = blockRollbackPersisted.fold(bi){ br => (br - 1).max(bi)}
+
+          // The above number would have been decreased by any rollbacked blocks
+          (nonEmptyRollbackedBlocks.head - 1).min(maxForcedPersist)
+        case (_, None) =>
+          // If persisted rollback, then it  was decreased by the future rollbacks, if not no persistance was ever done
+          blockRollbackPersisted.fold(0: BigInt)(_ => blocksRollbacked.head - 1)
+      }
+    }
+
+    trait StubPersistingBlockchainSetup {
+      def stubStateStorage: StateStorage
+      def blockchainStoragesWithStubPersisting: BlockchainStorages
+      def blockchainWithStubPersisting: BlockchainImpl
+    }
+
+    def newSetup(): StubPersistingBlockchainSetup = {
+      new StubPersistingBlockchainSetup with EphemBlockchainTestSetup {
+        override val stubStateStorage = stub[StateStorage]
+        override val blockchainStoragesWithStubPersisting = new BlockchainStorages {
+          val blockHeadersStorage = storagesInstance.storages.blockHeadersStorage
+          val blockBodiesStorage = storagesInstance.storages.blockBodiesStorage
+          val blockNumberMappingStorage = storagesInstance.storages.blockNumberMappingStorage
+          val receiptStorage = storagesInstance.storages.receiptStorage
+          val evmCodeStorage = storagesInstance.storages.evmCodeStorage
+          val chainWeightStorage = storagesInstance.storages.chainWeightStorage
+          val transactionMappingStorage = storagesInstance.storages.transactionMappingStorage
+          val nodeStorage = storagesInstance.storages.nodeStorage
+          val pruningMode = storagesInstance.storages.pruningMode
+          val appStateStorage = storagesInstance.storages.appStateStorage
+          val cachedNodeStorage = storagesInstance.storages.cachedNodeStorage
+          val stateStorage = stubStateStorage
+        }
+        override val blockchainWithStubPersisting = BlockchainImpl(blockchainStoragesWithStubPersisting)
+
+        blockchainWithStubPersisting.storeBlock(Fixtures.Blocks.Genesis.block)
+      }
+    }
+
+  }
 }