fix subscription

e2e/__snapshots__/storage.test.ts.snap

@@ -455,3 +455,19 @@ exports[`storage > getStorageMulti 2`] = `
   },
 ]
 `;
+
+exports[`storage > subscription 1`] = `
+[
+  [
+    1666661202090,
+  ],
+]
+`;
+
+exports[`storage > subscription 3`] = `
+[
+  [
+    1666828820000,
+  ],
+]
+`;

e2e/storage.test.ts

@@ -1,6 +1,7 @@
+import '@polkadot/api-augment'
 import { describe, expect, it } from 'vitest'
 
-import { api, expectJson } from './helper'
+import { api, delay, dev, expectJson, mockCallback } from './helper'
 
 describe('storage', () => {
   it('getStorage', async () => {
@@ -58,4 +59,33 @@ describe('storage', () => {
     })
     expect(entries2).toMatchSnapshot()
   })
+
+  it('subscription', async () => {
+    const { callback, next } = mockCallback()
+    const unsub = await api.query.timestamp.now(callback)
+
+    await next()
+
+    expect(callback.mock.calls).toMatchSnapshot()
+    callback.mockClear()
+
+    expect(await dev.newBlock()).toMatchInlineSnapshot(
+      '"0x5e29ae2538ffa601a9da913b75de8c95d0ce0bc7458756a094348d7f7e9b146a"'
+    )
+
+    await next()
+
+    expect(callback.mock.calls).toMatchSnapshot()
+    callback.mockClear()
+
+    unsub()
+
+    expect(await dev.newBlock()).toMatchInlineSnapshot(
+      '"0xe300c88d4790076560300b914c7a742929121cb2812fd931f859aa97e38b9393"'
+    )
+
+    await delay(100)
+
+    expect(callback).not.toHaveBeenCalled()
+  })
 })

src/blockchain/block.ts

@@ -93,7 +93,7 @@ class StorageLayer implements StorageLayerProvider {
     return undefined
   }
 
-  async set(key: string, value: StorageValue): Promise<void> {
+  set(key: string, value: StorageValue): void {
     switch (value) {
       case StorageValueKind.Deleted:
         this.#store[key] = value
@@ -110,12 +110,12 @@ class StorageLayer implements StorageLayerProvider {
     }
   }
 
-  async setAll(values: Record<string, StorageValue | null> | [string, StorageValue | null][]) {
+  setAll(values: Record<string, StorageValue | null> | [string, StorageValue | null][]) {
     if (!Array.isArray(values)) {
       values = Object.entries(values)
     }
     for (const [key, value] of values) {
-      await this.set(key, value || undefined)
+      this.set(key, value || undefined)
     }
   }
 
@@ -160,9 +160,9 @@ class StorageLayer implements StorageLayerProvider {
     return res
   }
 
-  mergeInto(into: Record<string, string>) {
+  async mergeInto(into: Record<string, string>) {
     for (const key of this.#keys) {
-      const value = this.#store[key]
+      const value = await this.#store[key]
       if (value === StorageValueKind.Deleted) {
         delete into[key]
       } else {
@@ -262,11 +262,11 @@ export class Block {
     this.#storages.pop()
   }
 
-  storageDiff(): Record<string, string> {
+  async storageDiff(): Promise<Record<string, string>> {
     const storage = {}
 
     for (const layer of this.#storages) {
-      layer.mergeInto(storage)
+      await layer.mergeInto(storage)
     }
 
     return storage

src/blockchain/head-state.ts

@@ -51,14 +51,14 @@ export class HeadState {
     void id
   }
 
-  setHead(head: Block) {
+  async setHead(head: Block) {
     this.#head = head
 
     for (const cb of Object.values(this.#headListeners)) {
       cb(head)
     }
 
-    const diff = this.#head.storageDiff()
+    const diff = await this.#head.storageDiff()
 
     for (const [keys, cb] of Object.values(this.#storageListeners)) {
       const changed = keys.filter((key) => diff[key]).map((key) => [key, diff[key]] as [string, string])

src/blockchain/txpool.ts

@@ -129,9 +129,12 @@ export class TxPool {
     const finalBlock = new Block(this.#api, this.#chain, newBlock.number, blockData.hash.toHex(), head, {
       header,
       extrinsics,
-      storage: newBlock.storage,
+      storage: head.storage,
     })
 
+    const diff = await newBlock.storageDiff()
+    finalBlock.pushStorageLayer().setAll(diff)
+
     this.#chain.unregisterBlock(newBlock)
     this.#chain.setHead(finalBlock)