Storage cache, concurrency, and block-state changes for test-runner consumer

[rockbmb]

This issue collects the internal Chopsticks performance changes originally proposed in open-web3-stack/polkadot-ecosystem-tests#606 by @ggwpez (with Claude). It is filed again here because the changes are entirely chopsticks-internal.

The original analysis and the snippets below are reproduced from #606 with two further changes appended at the end (sections 7 and 8) that are required for the PET shared-client test refactor in open-web3-stack/polkadot-ecosystem-tests#609 to work correctly. Those two were also developed by Claude in collaboration with @rockbmb.

Headline result on the PET suite (~1000 tests, ~70 files, ~22 chains): full run 477s on 8 forks with a warm disk cache.

---

1. chopsticks-db SQLite backend is unusable from worker_threads

The default TypeORM driver type: 'sqlite' (the sqlite3 package) fails to initialize inside worker_threads.Worker with DriverPackageNotInstalledError. Native binding is not initialized in worker contexts.

Patch: swap to type: 'better-sqlite3'. better-sqlite3 works inside workers and additionally enables enableWAL: true for safe multi-process concurrent reads/writes (vitest with --pool=forks --maxWorkers=8 against a shared DB_PATH).

- type: 'sqlite',
+ type: 'better-sqlite3',
+ enableWAL: true,
+ busyErrorRetry: 1000,
+ busyTimeout: 5000,

---

2. Disk cache only covers state_getStorage — getKeysPaged and chain/system metadata are re-fetched every run

chopsticks-db ships with one entity, KeyValuePair. Two other call categories are cold every run:

• state_getKeysPaged results: the runtime's next_key host calls during Core_initialize_block re-query upstream every cold setup.
• chain_getBlockHash(N), chain_getHeader(hash), chain_getBlock(hash), system_chain/properties/name: fetched on every chopsticks bootstrap.

Patch: two new TypeORM entities in chopsticks-db.

// PagedKeys: per-block-hash cache for state_getKeysPaged results
PagedKeys = new EntitySchema({
  columns: {
    blockHash: { primary: true, type: 'varchar' },
    prefix:    { primary: true, type: 'varchar' },
    startKey:  { primary: true, type: 'varchar' },
    keys:      { type: 'simple-array' },
    complete:  { type: 'boolean' },
  },
})

// RpcCall: scoped cache for deterministic chain/system RPC calls
RpcCall = new EntitySchema({
  columns: {
    scope:     { primary: true, type: 'varchar' },  // endpoint URL — DB_PATH is shared across chains
    method:    { primary: true, type: 'varchar' },
    paramsKey: { primary: true, type: 'varchar' },  // JSON.stringify(params)
    value:     { type: 'text' },
  },
})

Wired into RemoteStorageLayer.getKeysPaged (consults db.queryPagedKeys before going upstream) and Api.#cachedSend for the deterministic chain/system getters. chain_getFinalizedHead is intentionally NOT cached (non-deterministic).

scope is needed because in PET's setup DB_PATH is shared across all chains; without it, system_chain for kusama would collide with polkadot's.

---

3. getKeysPaged early-return bypasses both caches when prefix === startKey

In RemoteStorageLayer.getKeysPaged:

// can't handle keyCache without prefix
if (prefix === startKey || prefix.length < minPrefixLen || startKey.length < minPrefixLen) {
    return this.#api.getKeysPaged(prefix, pageSize, startKey, this.#at);
}

The prefix === startKey case fires every time the runtime asks "is there anything starting with this exact key" via next_key (very common during Core_initialize_block). The early-return goes straight to upstream RPC, bypassing both the in-memory KeyCache and the new PagedKeys disk cache. Each call costs one upstream round-trip (~120ms on public RPC).

Fix: drop the prefix === startKey clause. The KeyCache actually handles startKey === prefix correctly once it's been seeded with the empty-suffix marker: feed([startKey, ...batch]) produces a range like ['', 'aaaa', 'bbbb'] and next(prefix) returns 'aaaa'.

-if (prefix === startKey || prefix.length < minPrefixLen || startKey.length < minPrefixLen) {
+if (prefix.length < minPrefixLen || startKey.length < minPrefixLen) {
   return this.#api.getKeysPaged(prefix, pageSize, startKey, this.#at);
}

---

4. Post-cache-hit prefetch loop adds 400ms per getKeysPaged cache hit

After a successful getKeysPaged fetch, the layer runs:

if (this.#db) {
  const newBatch = [];
  for (const key of batch) {                              // <-- sequential await
    const res = await this.#db.queryStorage(this.#at, key);
    if (res) continue;
    newBatch.push(key);
  }
  if (newBatch.length > 0) {
    this.#api.getStorageBatch(prefix, newBatch, this.#at).then(/* save to db */);
  }
}

For BATCH_SIZE = 1000, that's 1000 sequential SQLite point lookups (~0.4ms each on warm OS cache) = 400-470ms per call measured. And the fire-and-forget getStorageBatch silently hits upstream too.

This filter+prefetch is only useful on cache miss, to warm the value cache for keys just discovered. On a cache hit (now possible thanks to the PagedKeys patch above), it's pure overhead.

Fix: one-line guard.

-if (this.#db) {
+if (this.#db && !cached) {
   const newBatch = [];
   ...
}

Measured impact: getKeysPaged warm hit went from 401ms → 5ms.

---

5. No in-flight request dedup in RemoteStorageLayer.get

Two concurrent get(key) calls for the same (blockHash, key) pair both fire upstream. There is no equivalent dedup in Api.send or ChopsticksProvider.send either.

Patch: a per-layer Map<key, Promise> to coalesce in-flight reads.

class RemoteStorageLayer {
  #inFlightGet = new Map();

  async get(key, _cache) {
    if (this.#db) {
      const res = await this.#db.queryStorage(this.#at, key);
      if (res) return res.value ?? undefined;
    }
    const inflight = this.#inFlightGet.get(key);
    if (inflight) return inflight;

    const promise = (async () => {
      try {
        const data = await this.#api.getStorage(key, this.#at);
        this.#db?.saveStorage(this.#at, key, data);
        return data ?? undefined;
      } finally {
        this.#inFlightGet.delete(key);
      }
    })();
    this.#inFlightGet.set(key, promise);
    return promise;
  }
}

Roughly halved cold-setup time on the PET suite (86s → 46s) by eliminating duplicate fetches when the runtime asks for the same key from parallel codepaths.

This is only per-process. Cross-process dedup (multiple vitest workers fetching the same uncached key) is still uncoordinated; SQLite serves as the eventual consistency layer.

---

6. New dryRunExtrinsicsAmortized for tests that need per-extrinsic event isolation

A common pattern in the PET suite: submit N independent extrinsics and snapshot each one's events independently (e.g. proxy-call filtering tests that run ~30 different proxy.proxy(call) variants and check allowed calls succeed, forbidden ones get ProxyExecuted{ ok: false, ... }).

The straightforward implementation calls client.dev.newBlock() per tx; each call pays a full Core_initialize_block (250ms WASM) + inherent (110ms) + BlockBuilder_finalize_block (110ms). For 30 txs that's ~15s of redundant WASM.

Patch: new exported helper in block-builder.js that does init + inherents once, then loops pushStorageLayer → BlockBuilder_apply_extrinsic → snapshot diff → popStorageLayer per tx.

export const dryRunExtrinsicsAmortized = async (head, inherentProviders, extrinsics, params) => {
  const registry = await head.registry;
  const header = await newHeader(head);
  const { block: newBlock } = await initNewBlock(head, header, inherentProviders, params);
  const results = [];
  for (const extrinsic of extrinsics) {
    newBlock.pushStorageLayer();
    const resp = await newBlock.call('BlockBuilder_apply_extrinsic', [extrinsic]);
    const outcome = registry.createType('ApplyExtrinsicResult', resp.result);
    results.push({ outcome, storageDiff: resp.storageDiff });
    newBlock.popStorageLayer();
  }
  return results;
};

Plus a thin chain.dryRunExtrinsicsAmortized(extrinsics, at) wrapper on Blockchain. Caller reads each tx's events from its returned storageDiff (decode Vec<FrameSystemEventRecord> at the canonical events key).

Caveat: because popStorageLayer rolls back account state including nonce, all txs must be pre-signed with the same baseNonce, not sequential nonces.

Measured impact on the proxy-filtering test: ~10× speedup on the inner loop, ~2.7× on the whole test.

---

7. Block.resetStorageLayers and storageLayerCount (required for PET #609)

A consumer that has captured a reference to a Block at a known good state and subsequently mutated it via dev.setStorage needs a way to revert those mutations without losing the block itself.

popStorageLayer is per-call. resetStorageLayers takes a target depth so the caller can record storageLayerCount at capture time and truncate back to it on restore.

class Block {
  // ...
  resetStorageLayers(targetCount: number): void {
    while (this.#storages.length > targetCount) this.#storages.pop()
  }

  get storageLayerCount(): number {
    return this.#storages.length
  }
}

PET's captureSnapshot uses these to clear accumulated dev.setStorage layers between tests in a shared-client suite. Without this, state from one test bleeds into the next.

---

8. HeadState.setHead notifies storage subscribers for keys reverted by a backward setHead (required for PET #609)

HeadState.setHead only fired storage-listener callbacks for keys present in the new head's storageDiff. When setHead is used to roll the chain back to an earlier block (as happens during snapshot restore), keys that were touched by the more recent block are absent from the older block's diff and so listeners were never told their prior values had been reverted.

The polkadot.js WsProvider, which relies on chain_subscribeStorage notifications to track the latest known value, would silently retain the post-rollback storage value and route subsequent state_queryStorageAt calls against the wrong block hash.

async setHead(head: Block) {
  this.#head = head
  // ...

  const diff = await this.#head.storageDiff()
  const oldValues = this.#oldValues

  for (const [keys, cb] of Object.values(this.#storageListeners)) {
    const relevantKeys = keys.filter((key) => diff[key] !== undefined || oldValues[key] !== undefined)
    if (relevantKeys.length > 0) {
      const pairs = await Promise.all(
        relevantKeys.map(async (key): Promise<[string, string | null]> => [
          key,
          diff[key] !== undefined ? (diff[key] ?? null) : ((await head.get(key)) ?? null),
        ]),
      )
      try {
        await cb(head, pairs)
      } catch (error) {
        logger.error(error, 'setHead storage diff callback error')
      }
    }
  }

  this.#oldValues = { ...diff }
}

relevantKeys now also includes any key that had a value in the previous head's diff (tracked in #oldValues). For those keys the fresh value is read from the new head; subscribers always see a correct, current value after each setHead.

---

Benchmarks

Per @ggwpez:

On the PET suite (M-series Mac, ~1000 e2e tests across 22 chains, ~70 files):

Hot run (warm disk cache, 3 forks):

 Test Files  65 passed | 5 skipped (70)
      Tests  1031 passed | 20 skipped (1051)
   Duration  1208.00s

Hot run (warm disk cache, 8 forks):

 Test Files  65 passed | 5 skipped (70)
      Tests  1031 passed | 20 skipped (1051)
   Duration  476.91s

Cold run (empty disk cache, 3 forks):

 Test Files  65 passed | 5 skipped (70)
      Tests  1031 passed | 20 skipped (1051)
   Duration  1256.48s

Multi-fork parallelism only works once getKeysPaged and chain/system calls also hit the disk cache (sections 1, 2, 4 combined).