blockchain: Fix and enhance the blockchain iterator (#2308)

Author: g11tech

packages/blockchain/src/blockchain.ts

@@ -1,13 +1,13 @@
 import { Block, BlockHeader } from '@ethereumjs/block'
 import { Chain, Common, ConsensusAlgorithm, ConsensusType, Hardfork } from '@ethereumjs/common'
+import { Lock } from '@ethereumjs/util'
 import { MemoryLevel } from 'memory-level'
 
 import { CasperConsensus, CliqueConsensus, EthashConsensus } from './consensus'
 import { DBOp, DBSaveLookups, DBSetBlockOrHeader, DBSetHashToNumber, DBSetTD } from './db/helpers'
 import { DBManager } from './db/manager'
 import { DBTarget } from './db/operation'
 import { genesisStateRoot } from './genesisStates'
-import { Lock } from './lock'
 
 import type { Consensus } from './consensus'
 import type { GenesisState } from './genesisStates'
@@ -911,33 +911,47 @@ export class Blockchain implements BlockchainInterface {
    * @param maxBlocks - How many blocks to run. By default, run all unprocessed blocks in the canonical chain.
    * @returns number of blocks actually iterated
    */
-  async iterator(name: string, onBlock: OnBlock, maxBlocks?: number): Promise<number> {
-    return this._iterator(name, onBlock, maxBlocks)
-  }
-
-  /**
-   * @hidden
-   */
-  private async _iterator(name: string, onBlock: OnBlock, maxBlocks?: number): Promise<number> {
+  async iterator(
+    name: string,
+    onBlock: OnBlock,
+    maxBlocks?: number,
+    releaseLockOnCallback?: boolean
+  ): Promise<number> {
     return this.runWithLock<number>(async (): Promise<number> => {
-      const headHash = this._heads[name] ?? this.genesisBlock.hash()
+      let headHash = this._heads[name] ?? this.genesisBlock.hash()
 
       if (typeof maxBlocks === 'number' && maxBlocks < 0) {
         throw 'If maxBlocks is provided, it has to be a non-negative number'
       }
 
-      const headBlockNumber = await this.dbManager.hashToNumber(headHash)
+      let headBlockNumber = await this.dbManager.hashToNumber(headHash)
       let nextBlockNumber = headBlockNumber + BigInt(1)
       let blocksRanCounter = 0
       let lastBlock: Block | undefined
 
       while (maxBlocks !== blocksRanCounter) {
         try {
-          const nextBlock = await this._getBlock(nextBlockNumber)
+          let nextBlock = await this._getBlock(nextBlockNumber)
+          const reorg = lastBlock ? !lastBlock.hash().equals(nextBlock.header.parentHash) : false
+          if (reorg) {
+            // If reorg has happened, the _heads must have been updated so lets reload the counters
+            headHash = this._heads[name] ?? this.genesisBlock.hash()
+            headBlockNumber = await this.dbManager.hashToNumber(headHash)
+            nextBlockNumber = headBlockNumber + BigInt(1)
+            nextBlock = await this._getBlock(nextBlockNumber)
+          }
           this._heads[name] = nextBlock.hash()
-          const reorg = lastBlock ? lastBlock.hash().equals(nextBlock.header.parentHash) : false
           lastBlock = nextBlock
-          await onBlock(nextBlock, reorg)
+          if (releaseLockOnCallback === true) {
+            this._lock.release()
+          }
+          try {
+            await onBlock(nextBlock, reorg)
+          } finally {
+            if (releaseLockOnCallback === true) {
+              await this._lock.acquire()
+            }
+          }
           nextBlockNumber++
           blocksRanCounter++
         } catch (error: any) {

packages/blockchain/src/types.ts

@@ -37,7 +37,12 @@ export interface BlockchainInterface {
    * @param maxBlocks - optional maximum number of blocks to iterate through
    * reorg: boolean)
    */
-  iterator(name: string, onBlock: OnBlock, maxBlocks?: number): Promise<number>
+  iterator(
+    name: string,
+    onBlock: OnBlock,
+    maxBlocks?: number,
+    releaseLockOnCallback?: boolean
+  ): Promise<number>
 
   /**
    * Returns a copy of the blockchain

packages/blockchain/test/index.spec.ts

@@ -368,143 +368,6 @@ tape('blockchain test', (t) => {
     st.end()
   })
 
-  t.test('should iterate through 24 blocks', async (st) => {
-    const { blockchain, blocks, error } = await generateBlockchain(25)
-    st.error(error, 'no error')
-    let i = 0
-    const iterated = await blockchain.iterator('test', (block: Block) => {
-      if (block.hash().equals(blocks[i + 1].hash())) {
-        i++
-      }
-    })
-    st.equal(iterated, 24)
-    st.equal(i, 24)
-    st.end()
-  })
-
-  t.test(
-    'should iterate through maxBlocks blocks if maxBlocks parameter is provided',
-    async (st) => {
-      const { blockchain, blocks, error } = await generateBlockchain(25)
-      st.error(error, 'no error')
-      let i = 0
-      const iterated = await blockchain.iterator(
-        'test',
-        (block: Block) => {
-          if (block.hash().equals(blocks[i + 1].hash())) {
-            i++
-          }
-        },
-        5
-      )
-      st.equal(iterated, 5)
-      st.equal(i, 5)
-      st.end()
-    }
-  )
-
-  t.test(
-    'should iterate through 0 blocks in case 0 maxBlocks parameter is provided',
-    async (st) => {
-      const { blockchain, blocks, error } = await generateBlockchain(25)
-      st.error(error, 'no error')
-      let i = 0
-      const iterated = await blockchain
-        .iterator(
-          'test',
-          (block: Block) => {
-            if (block.hash().equals(blocks[i + 1].hash())) {
-              i++
-            }
-          },
-          0
-        )
-        .catch(() => {
-          st.fail('Promise cannot throw when running 0 blocks')
-        })
-      st.equal(iterated, 0)
-      st.equal(i, 0)
-      st.end()
-    }
-  )
-
-  t.test('should throw on a negative maxBlocks parameter in iterator', async (st) => {
-    const { blockchain, blocks, error } = await generateBlockchain(25)
-    st.error(error, 'no error')
-    let i = 0
-    await blockchain
-      .iterator(
-        'test',
-        (block: Block) => {
-          if (block.hash().equals(blocks[i + 1].hash())) {
-            i++
-          }
-        },
-        -1
-      )
-      .catch(() => {
-        st.end()
-      })
-    // Note: if st.end() is not called (Promise did not throw), then this test fails, as it does not end.
-  })
-
-  t.test('should test setIteratorHead method', async (st) => {
-    const { blockchain, blocks, error } = await generateBlockchain(25)
-    st.error(error, 'no error')
-
-    const headBlockIndex = 5
-
-    const headHash = blocks[headBlockIndex].hash()
-    await blockchain.setIteratorHead('myHead', headHash)
-    const currentHeadBlock = await blockchain.getIteratorHead('myHead')
-
-    st.ok(headHash.equals(currentHeadBlock.hash()), 'head hash equals the provided head hash')
-
-    let i = 0
-    // check that iterator starts from this head block
-    await blockchain.iterator(
-      'myHead',
-      (block: Block) => {
-        if (block.hash().equals(blocks[headBlockIndex + 1].hash())) {
-          i++
-        }
-      },
-      5
-    )
-
-    st.equal(i, 1)
-
-    st.end()
-  })
-
-  t.test('should catch iterator func error', async (st) => {
-    const { blockchain, error } = await generateBlockchain(25)
-    st.error(error, 'no error')
-    try {
-      await blockchain.iterator('error', () => {
-        throw new Error('iterator func error')
-      })
-    } catch (error: any) {
-      st.ok(error)
-      st.equal(error.message, 'iterator func error', 'should return correct error')
-      st.end()
-    }
-  })
-
-  t.test('should not call iterator function in an empty blockchain', async (st) => {
-    const blockchain = await Blockchain.create({
-      validateBlocks: true,
-      validateConsensus: false,
-    })
-
-    await blockchain.iterator('test', () => {
-      st.fail('should not call iterator function')
-    })
-
-    st.pass('should finish iterating')
-    st.end()
-  })
-
   t.test('should add fork header and reset stale heads', async (st) => {
     const { blockchain, blocks, error } = await generateBlockchain(15)
     st.error(error, 'no error')
@@ -615,23 +478,6 @@ tape('blockchain test', (t) => {
     st.end()
   })
 
-  t.test('should get heads', async (st) => {
-    const [db, genesis] = await createTestDB()
-    const blockchain = await Blockchain.create({ db, genesisBlock: genesis })
-    const head = await blockchain.getIteratorHead()
-    if (typeof genesis !== 'undefined') {
-      st.ok(head.hash().equals(genesis.hash()), 'should get head')
-      st.equal(
-        (blockchain as any)._heads['head0'].toString('hex'),
-        'abcd',
-        'should get state root heads'
-      )
-      st.end()
-    } else {
-      st.fail()
-    }
-  })
-
   t.test('should validate', async (st) => {
     const genesisBlock = Block.fromBlockData({ header: { gasLimit: 8000000 } })
     const blockchain = await Blockchain.create({