feat: load blocks from db (#392)

* feat: load blocks from db

* fix load blocks

* fix: make tests independent

* fix: override storage after load blocks

* fix: storage

* fix: load latest block only

* feat: resume using block number hash

packages/chopsticks/src/cli.ts

@@ -44,6 +44,12 @@ const commands = yargs(hideBin(process.argv))
           desc: 'Max memory block count',
           number: true,
         },
+        resume: {
+          desc: `Resume from the specified block hash or block number in db.
+                 If true, it will resume from the latest block in db.
+                 Note this will override the block option`,
+          string: true,
+        },
       }),
     async (argv) => {
       await setupWithServer(argv as Config)

packages/chopsticks/src/context.ts

@@ -1,4 +1,5 @@
 import './utils/tunnel'
+import { BlockEntity } from '@acala-network/chopsticks-core/db/entities'
 import { Config } from './schema'
 import { HexString } from '@polkadot/util/types'
 import { overrideStorage, overrideWasm } from './utils/override'
@@ -19,6 +20,28 @@ export const setupContext = async (argv: Config, overrideParent = false) => {
     maxMemoryBlockCount: argv['max-memory-block-count'],
   })
 
+  // load block from db
+  if (chain.db) {
+    if (argv.resume) {
+      const where: Record<string, string | number> = {}
+      switch (typeof argv.resume) {
+        case 'string':
+          where.hash = argv.resume
+          break
+        case 'number':
+          where.number = argv.resume
+          break
+        default:
+          break
+      }
+      const blockData = await chain.db.getRepository(BlockEntity).findOne({ where, order: { number: 'desc' } })
+      if (blockData) {
+        const block = await chain.loadBlockFromDB(blockData?.number)
+        block && (await chain.setHead(block))
+      }
+    }
+  }
+
   if (argv.timestamp) await timeTravel(chain, argv.timestamp)
 
   let at: HexString | undefined

packages/chopsticks/src/schema/index.ts

@@ -22,6 +22,7 @@ export const configSchema = z
     'registered-types': z.any().optional(),
     'runtime-log-level': z.number().min(0).max(5).optional(),
     'offchain-worker': z.boolean().optional(),
+    resume: z.union([z.string().length(66).startsWith('0x'), z.number(), z.boolean()]).optional(),
   })
   .strict()
 

packages/core/src/blockchain/index.ts

@@ -149,13 +149,18 @@ export class Blockchain {
         .getRepository(BlockEntity)
         .findOne({ where: { [typeof key === 'number' ? 'number' : 'hash']: key } })
       if (blockData) {
-        const { hash, number, header, extrinsics, parentHash } = blockData
-        const parentBlock = parentHash ? this.#blocksByHash[parentHash] : undefined
+        const { hash, number, header, extrinsics } = blockData
+        const parentHash = blockData.parentHash || undefined
+        let parentBlock = parentHash ? this.#blocksByHash[parentHash] : undefined
+        if (!parentBlock) {
+          parentBlock = await this.getBlock(parentHash)
+        }
         const storageDiff = blockData.storageDiff ?? undefined
         const registry = await this.head.registry
         const block = new Block(this, number, hash, parentBlock, {
           header: registry.createType<Header>('Header', header),
           extrinsics,
+          storage: parentBlock?.storage,
           storageDiff,
         })
         this.#registerBlock(block)

packages/e2e/src/blocks-save.test.ts

@@ -1,29 +1,74 @@
-import { afterAll, assert, describe, expect, it } from 'vitest'
+import { assert, describe, expect, it } from 'vitest'
 import { resolve } from 'node:path'
 import { tmpdir } from 'node:os'
 import networks from './networks'
 
 describe('block-save', async () => {
-  const acala = await networks.acala({ db: resolve(tmpdir(), 'db.sqlite') })
-  const { chain, dev } = acala
-
-  afterAll(async () => {
+  const buildBlocks = async () => {
+    // save blocks
+    const acala = await networks.acala({ db: resolve(tmpdir(), 'db.sqlite') })
+    const { chain, dev } = acala
+    await dev.newBlock({ count: 2 })
+    const head = await chain.getBlockAt(chain.head.number)
+    const savedHeadHash = head?.hash
     await acala.teardown()
-  })
+
+    return savedHeadHash
+  }
 
   it('saved blocks data', async () => {
+    const acala = await networks.acala({ db: resolve(tmpdir(), 'db.sqlite') })
+    const { chain, dev } = acala
     await dev.newBlock({ count: 2 })
 
     const numberOfBlocks = await chain.db!.getRepository('Block').count()
     expect(numberOfBlocks).toEqual(2)
 
     const block = await chain.getBlockAt(chain.head.number)
     const blockData = await chain.db!.getRepository('Block').findOne({ where: { number: chain.head.number } })
+
     assert(block && blockData, 'block and blockData should be defined')
     expect(blockData.hash).toEqual(block.hash)
     expect(JSON.stringify(blockData.header)).toEqual(JSON.stringify(block.header))
     expect(blockData.parentHash).toEqual((await block.parentBlock)!.hash)
     expect(JSON.stringify(blockData.extrinsics)).toEqual(JSON.stringify(await block.extrinsics))
     expect(JSON.stringify(blockData.storageDiff)).toEqual(JSON.stringify(await block.storageDiff()))
+
+    await acala.teardown()
+  })
+
+  it('load chain using the latest saved block', async () => {
+    const savedHeadHash = await buildBlocks()
+
+    // load block
+    const newAcala = await networks.acala({ db: resolve(tmpdir(), 'db.sqlite'), resume: true })
+    const newHeadNumber = newAcala.chain.head.number
+    const loadedHead = await newAcala.chain.getBlockAt(newHeadNumber)
+
+    expect(loadedHead?.hash).toEqual(savedHeadHash)
+    await newAcala.teardown()
+  })
+
+  it('load chain using a block number', async () => {
+    await buildBlocks()
+
+    // load blocks
+    const newAcala = await networks.acala({ db: resolve(tmpdir(), 'db.sqlite'), resume: 3000001 })
+    const newHeadNumber = newAcala.chain.head.number
+
+    expect(newHeadNumber).toEqual(3000001)
+    await newAcala.teardown()
+  })
+
+  it('load chain using a block hash', async () => {
+    const savedHeadHash = await buildBlocks()
+
+    // load blocks
+    const newAcala = await networks.acala({ db: resolve(tmpdir(), 'db.sqlite'), resume: savedHeadHash })
+    const newHeadNumber = newAcala.chain.head.number
+    const loadedHead = await newAcala.chain.getBlockAt(newHeadNumber)
+
+    expect(loadedHead?.hash).toEqual(savedHeadHash)
+    await newAcala.teardown()
   })
 })

packages/testing/src/index.ts

@@ -25,6 +25,7 @@ export type SetupOption = {
   timeout?: number
   port?: number
   maxMemoryBlockCount?: number
+  resume?: boolean | HexString | number
 }
 
 export type SetupConfig = Config & {
@@ -40,6 +41,7 @@ export const createConfig = ({
   timeout,
   port,
   maxMemoryBlockCount,
+  resume,
 }: SetupOption): SetupConfig => {
   // random port if not specified
   port = port ?? Math.floor(Math.random() * 10000) + 10000
@@ -53,6 +55,7 @@ export const createConfig = ({
     db,
     'wasm-override': wasmOverride,
     timeout,
+    resume: resume ?? false,
   }
   return config
 }