Max memory block count (#353)

* fix typos

* feat: add maxMemoryBlockCount option

* tests: add max-memory-block-count test

* remove extra re-insert block

packages/chopsticks/src/blockchain/block-builder.ts

@@ -354,9 +354,9 @@ export const dryRunInherents = async (
 ): Promise<[HexString, HexString | null][]> => {
   const header = await newHeader(head)
   const { layers } = await initNewBlock(head, header, inherents)
-  const stoarge = {}
+  const storage = {}
   for (const layer of layers) {
-    await layer.mergeInto(stoarge)
+    await layer.mergeInto(storage)
   }
-  return Object.entries(stoarge) as [HexString, HexString | null][]
+  return Object.entries(storage) as [HexString, HexString | null][]
 }

packages/chopsticks/src/blockchain/head-state.ts

@@ -77,7 +77,7 @@ export class HeadState {
         try {
           await cb(head, changed)
         } catch (error) {
-          logger.error(error, 'setHead storaeg diff callback error')
+          logger.error(error, 'setHead storage diff callback error')
         }
       }
     }

packages/chopsticks/src/blockchain/index.ts

@@ -30,6 +30,7 @@ export interface Options {
   runtimeLogLevel?: number
   registeredTypes: RegisteredTypes
   offchainWorker?: boolean
+  maxMemoryBlockCount?: number
 }
 
 export class Blockchain {
@@ -45,13 +46,14 @@ export class Blockchain {
   readonly #inherentProvider: InherentProvider
 
   #head: Block
-  readonly #blocksByNumber: Block[] = []
+  readonly #blocksByNumber: Map<number, Block> = new Map()
   readonly #blocksByHash: Record<string, Block> = {}
   readonly #loadingBlocks: Record<string, Promise<void>> = {}
 
   readonly headState: HeadState
 
   readonly offchainWorker: OffchainWorker | undefined
+  readonly #maxMemoryBlockCount: number
 
   constructor({
     api,
@@ -64,6 +66,7 @@ export class Blockchain {
     runtimeLogLevel = 0,
     registeredTypes = {},
     offchainWorker = false,
+    maxMemoryBlockCount = 2000,
   }: Options) {
     this.api = api
     this.db = db
@@ -83,10 +86,17 @@ export class Blockchain {
     if (offchainWorker) {
       this.offchainWorker = new OffchainWorker()
     }
+
+    this.#maxMemoryBlockCount = maxMemoryBlockCount
   }
 
   #registerBlock(block: Block) {
-    this.#blocksByNumber[block.number] = block
+    // if exceed max memory block count, delete the oldest block
+    if (this.#blocksByNumber.size === this.#maxMemoryBlockCount) {
+      const firstKey = this.#blocksByNumber.keys().next().value
+      this.#blocksByNumber.delete(firstKey)
+    }
+    this.#blocksByNumber.set(block.number, block)
     this.#blocksByHash[block.hash] = block
   }
 
@@ -105,12 +115,12 @@ export class Blockchain {
     if (number > this.#head.number) {
       return undefined
     }
-    if (!this.#blocksByNumber[number]) {
+    if (!this.#blocksByNumber.has(number)) {
       const hash = await this.api.getBlockHash(number)
       const block = new Block(this, number, hash)
       this.#registerBlock(block)
     }
-    return this.#blocksByNumber[number]
+    return this.#blocksByNumber.get(number)
   }
 
   async getBlock(hash?: HexString): Promise<Block | undefined> {
@@ -140,12 +150,16 @@ export class Blockchain {
     return this.#blocksByHash[hash]
   }
 
+  blocksInMemory(): Block[] {
+    return Array.from(this.#blocksByNumber.values())
+  }
+
   unregisterBlock(block: Block): void {
     if (block.hash === this.head.hash) {
       throw new Error('Cannot unregister head block')
     }
-    if (this.#blocksByNumber[block.number]?.hash === block.hash) {
-      delete this.#blocksByNumber[block.number]
+    if (this.#blocksByNumber.get(block.number)?.hash === block.hash) {
+      this.#blocksByNumber.delete(block.number)
     }
     delete this.#blocksByHash[block.hash]
   }
@@ -280,7 +294,7 @@ export class Blockchain {
 
     const needsDispatch = meta.registry.createType('Vec<u32>', Object.keys(ump))
 
-    const stroageValues: [string, StorageValue | null][] = [
+    const storageValues: [string, StorageValue | null][] = [
       [compactHex(meta.query.ump.needsDispatch()), needsDispatch.toHex()],
     ]
 
@@ -293,11 +307,11 @@ export class Blockchain {
         upwardMessages.map((x) => x.byteLength).reduce((s, i) => s + i, 0),
       ])
 
-      stroageValues.push([compactHex(meta.query.ump.relayDispatchQueues(paraId)), upwardMessages.toHex()])
-      stroageValues.push([compactHex(meta.query.ump.relayDispatchQueueSize(paraId)), queueSize.toHex()])
+      storageValues.push([compactHex(meta.query.ump.relayDispatchQueues(paraId)), upwardMessages.toHex()])
+      storageValues.push([compactHex(meta.query.ump.relayDispatchQueueSize(paraId)), queueSize.toHex()])
     }
 
-    head.pushStorageLayer().setAll(stroageValues)
+    head.pushStorageLayer().setAll(storageValues)
     const inherents = await this.#inherentProvider.createInherents(head, {
       transactions: [],
       downwardMessages: [],

packages/chopsticks/src/cli.ts

@@ -39,6 +39,10 @@ const commands = yargs(hideBin(process.argv))
           desc: 'Allow wasm unresolved imports',
           boolean: true,
         },
+        'max-memory-block-count': {
+          desc: 'Max memory block count',
+          number: true,
+        },
       }),
     async (argv) => {
       await setupWithServer(argv as Config)

packages/chopsticks/src/schema/index.ts

@@ -30,6 +30,7 @@ export const configSchema = z
     'build-block-mode': z.nativeEnum(BuildBlockMode).optional(),
     'import-storage': z.any().optional(),
     'mock-signature-host': z.boolean().optional(),
+    'max-memory-block-count': z.number().optional(),
     db: z.string().optional(),
     'wasm-override': z.string().optional(),
     genesis: z.union([z.string(), genesisSchema]).optional(),

packages/chopsticks/src/setup.ts

@@ -78,6 +78,7 @@ export const setup = async (argv: Config, runBlock = false) => {
     runtimeLogLevel: argv['runtime-log-level'],
     registeredTypes: argv['registered-types'],
     offchainWorker: argv['offchain-worker'],
+    maxMemoryBlockCount: argv['max-memory-block-count'],
   })
 
   if (argv.timestamp) await timeTravel(chain, argv.timestamp)

packages/e2e/src/max-memory-block-count.test.ts

@@ -0,0 +1,20 @@
+import { afterAll, describe, expect, it } from 'vitest'
+import networks from './networks'
+
+describe('max-memory-block-count', async () => {
+  const acala = await networks.acala({ maxMemoryBlockCount: 2 })
+  const { chain } = acala
+
+  afterAll(async () => {
+    await acala.teardown()
+  })
+
+  it('removes the oldest block when exceed', async () => {
+    // ensure the first block is registered
+    const firstBlock = await chain.getBlockAt(chain.head.number)
+    await acala.dev.newBlock({ count: 2 })
+    const blocksInMemory = chain.blocksInMemory()
+    expect(blocksInMemory[0].number).toEqual((firstBlock?.number || 0) + 1)
+    expect(blocksInMemory.length).toEqual(2)
+  })
+})

packages/testing/src/index.ts

@@ -24,6 +24,7 @@ export type SetupOption = {
   db?: string
   timeout?: number
   port?: number
+  maxMemoryBlockCount?: number
 }
 
 export type SetupConfig = Config & {
@@ -38,6 +39,7 @@ export const createConfig = ({
   db,
   timeout,
   port,
+  maxMemoryBlockCount,
 }: SetupOption): SetupConfig => {
   // random port if not specified
   port = port ?? Math.floor(Math.random() * 10000) + 10000
@@ -47,6 +49,7 @@ export const createConfig = ({
     block: blockNumber || blockHash,
     mockSignatureHost: true,
     'build-block-mode': BuildBlockMode.Manual,
+    'max-memory-block-count': maxMemoryBlockCount ?? 100,
     db,
     'wasm-override': wasmOverride,
     timeout,