From bf0bc9a7c462f07dd2758a6af5f316535d736f88 Mon Sep 17 00:00:00 2001
From: Mokhtar Naamani <mokhtar.naamani@gmail.com>
Date: Mon, 15 Jan 2024 12:23:07 +0400
Subject: [PATCH] runtime upgrade test: handle large state storage

---
 .../run-runtime-upgrade-tests.sh              |  9 +++--
 utils/api-scripts/src/fork-off.ts             | 39 ++++++++++---------
 2 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/tests/network-tests/run-runtime-upgrade-tests.sh b/tests/network-tests/run-runtime-upgrade-tests.sh
index 1611cc1a61..7a7500f293 100755
--- a/tests/network-tests/run-runtime-upgrade-tests.sh
+++ b/tests/network-tests/run-runtime-upgrade-tests.sh
@@ -6,9 +6,6 @@ cd $SCRIPT_PATH
 
 rm ./output.json || :
 
-# Log only to stderr
-# Only output from this script should be the container id of the node at the very end
-
 # Location that will be mounted to /spec in containers
 # This is where the initial balances files and generated chainspec files will be located.
 DATA_PATH=$PWD/data
@@ -118,6 +115,7 @@ function fork_off_init() {
 
     # download the raw storage state
     if ! [[ -f ${DATA_PATH}/storage.json ]]; then
+        echo >&2 "fetching state storage from $HTTP_RPC_ENDPOINT"
         curl $HTTP_RPC_ENDPOINT -H \
             "Content-type: application/json" -d \
             '{"jsonrpc":"2.0","id":1,"method":"state_getPairs","params":["0x"]}' \
@@ -125,7 +123,8 @@ function fork_off_init() {
         echo >&2 "storage trie downloaded at ${DATA_PATH}/storage.json"
     fi
 
-    yarn workspace api-scripts tsnode-strict src/fork-off.ts ${DATA_PATH} ${WS_RPC_ENDPOINT}
+    yarn workspace api-scripts tsnode-strict --max-old-space-size=6144 \
+        src/fork-off.ts ${DATA_PATH} ${WS_RPC_ENDPOINT}
 }
 
 #######################################
@@ -163,6 +162,8 @@ function main {
     # 2. clone live chainspec with fork it
     fork_off_init
     echo >&2 "storage downloaded & dumped into the raw chainspec"
+    return
+
     # 3. set path to new runtime.wasm
     set_new_runtime_wasm_path
     echo >&2 "new wasm path set"
diff --git a/utils/api-scripts/src/fork-off.ts b/utils/api-scripts/src/fork-off.ts
index 987ead54e4..4135a4d1df 100644
--- a/utils/api-scripts/src/fork-off.ts
+++ b/utils/api-scripts/src/fork-off.ts
@@ -1,7 +1,7 @@
 import { ApiPromise, WsProvider } from '@polkadot/api'
 import { xxhashAsHex } from '@polkadot/util-crypto'
-import fs from 'fs'
 import path from 'path'
+
 const bfj = require('bfj')
 
 /**
@@ -36,6 +36,12 @@ const skippedModulesPrefix = [
   'Instance3WorkingGroup', // empty content working group
   'Instance9WorkingGroup', // empty distribution working group
 ]
+const skippedKeys = [
+  // encoded storage key for `minAuctionDuration` key
+  '0xb5a494c92fa4747cc071573e93b32b87f9ad4eaa35a4c52d9289acbc42eba9d9',
+  // encoded storage key for `nftLimitsEnabled` key
+  '0xb5a494c92fa4747cc071573e93b32b87d2c14024f1b303fdc87019c4c1facfde'
+]
 
 async function main() {
   // paths & env variables
@@ -63,22 +69,20 @@ async function main() {
   })
 
   // blank starting chainspec guaranteed to exist
-  const storage: Storage = await bfj.read(storagePath)
+  console.error('Loading state storage')
+  const storage: Storage = await bfj.read(storagePath, { bufferLength: 4096 })
+  console.error('Loading raw chain spec')
   const chainSpec = await bfj.read(specPath)
 
   // Grab the items to be moved, then iterate through and insert into storage
-  storage.result
-    .filter((i) => prefixes.some((prefix) => i[0].startsWith(prefix)))
-    .forEach(([key, value]) => {
-      if (
-        // encoded storage key for `minAuctionDuration` key
-        key !== '0xb5a494c92fa4747cc071573e93b32b87f9ad4eaa35a4c52d9289acbc42eba9d9' &&
-        // encoded storage key for `nftLimitsEnabled` key
-        key !== '0xb5a494c92fa4747cc071573e93b32b87d2c14024f1b303fdc87019c4c1facfde'
-      ) {
-        chainSpec.genesis.raw.top[key] = value
-      }
-    })
+  // We are not using a storage.result.forEach to preserve memory
+  console.error('constructing final genesis chainspec')
+  for (let pair; pair = storage.result.pop();) {
+    const [key, value] = pair
+    if (skippedKeys.includes(key)) continue
+    if (!prefixes.some((prefix) => key.startsWith(prefix))) continue
+    chainSpec.genesis.raw.top[key] = value
+  }
 
   // Delete System.LastRuntimeUpgrade to ensure that the on_runtime_upgrade event is triggered
   delete chainSpec.genesis.raw.top['0x26aa394eea5630e07c48ae0c9558cef7f9cce9c888469bb1a0dceaa129672ef8']
@@ -86,12 +90,11 @@ async function main() {
   // To prevent the validator set from changing mid-test, set Staking.ForceEra to ForceNone ('0x02')
   chainSpec.genesis.raw.top['0x5f3e4907f716ac89b6347d15ececedcaf7dad0317324aecae8744b87fc95f2f3'] = '0x02'
 
-  fs.writeFileSync(forkedSpecPath, JSON.stringify(chainSpec, null, 4))
-
-  process.exit()
+  console.error('Writing chainspec to', forkedSpecPath)
+  await bfj.write(forkedSpecPath, chainSpec, { space: 4, bufferLength: 4096 })
 }
 
-main().catch(console.error)
+main().catch(console.error).finally(() => process.exit())
 
 interface Storage {
   'jsonrpc': string