txFrame HashKeys (kMap) memory optimization (#3606)

* Move kMap HashKeys in ForkedChain.

* Only move hashkeys when batchsize is greater than 1.

* Add a test.

* Add more tests.

Author: bhartnett

execution_chain/core/chain/forked_chain.nim

@@ -437,8 +437,13 @@ proc validateBlock(c: ForkedChainRef,
       c.latestFinalizedBlockNumber)
 
   let
+    # As a memory optimization we move the HashKeys (kMap) stored in the
+    # parent txFrame to the new txFrame unless the block number is one
+    # greater than a block which is expected to be persisted based on the
+    # persistBatchSize
+    moveParentHashKeys = c.persistBatchSize > 1 and (blk.header.number mod c.persistBatchSize) != 1
     parentFrame = parent.txFrame
-    txFrame = parentFrame.txFrameBegin
+    txFrame = parentFrame.txFrameBegin(moveParentHashKeys)
 
   # TODO shortLog-equivalent for eth types
   debug "Validating block",

execution_chain/db/aristo/aristo_tx_frame.nim

@@ -91,9 +91,15 @@ proc buildSnapshot(txFrame: AristoTxRef, minLevel: int) =
 
   txFrame.snapshot.level = Opt.some(minLevel)
 
-proc txFrameBegin*(db: AristoDbRef, parent: AristoTxRef): AristoTxRef =
+proc txFrameBegin*(db: AristoDbRef, parent: AristoTxRef, moveParentHashKeys = false): AristoTxRef =
   let parent = if parent == nil: db.txRef else: parent
-  AristoTxRef(db: db, parent: parent, vTop: parent.vTop, level: parent.level + 1)
+
+  AristoTxRef(
+    db: db,
+    parent: parent,
+    kMap: if moveParentHashKeys: move(parent.kMap) else: default(parent.kMap.type),
+    vTop: parent.vTop,
+    level: parent.level + 1)
 
 proc dispose*(tx: AristoTxRef) =
   tx[].reset()

execution_chain/db/core_db/base.nim

@@ -455,16 +455,19 @@ proc txFrameBegin*(db: CoreDbRef): CoreDbTxRef =
   ##
   let
     kTx = db.kvt.txFrameBegin(nil)
-    aTx = db.mpt.txFrameBegin(nil)
+    aTx = db.mpt.txFrameBegin(nil, false)
 
   CoreDbTxRef(kTx: kTx, aTx: aTx)
 
-proc txFrameBegin*(parent: CoreDbTxRef): CoreDbTxRef =
+proc txFrameBegin*(
+    parent: CoreDbTxRef,
+    moveParentHashKeys = false
+  ): CoreDbTxRef =
   ## Constructor
   ##
   let
     kTx = parent.kTx.db.txFrameBegin(parent.kTx)
-    aTx = parent.aTx.db.txFrameBegin(parent.aTx)
+    aTx = parent.aTx.db.txFrameBegin(parent.aTx, moveParentHashKeys)
 
   CoreDbTxRef(kTx: kTx, aTx: aTx)
 

tests/test_aristo/test_tx_frame.nim

@@ -108,3 +108,96 @@ suite "Aristo TxFrame":
       check:
         tx.fetchAccountRecord(acc1[0]).isOk()
         tx.fetchAccountRecord(acc2[0]).isErr() # Doesn't exist in tx2
+
+  test "Frames using moveParentHashKeys parameter":
+    let
+      tx0 = db.txFrameBegin(db.baseTxFrame())
+      tx1 = db.txFrameBegin(tx0)
+
+    check:
+      tx0.mergeAccountRecord(acc1[0], acc1[1]).isOk()
+      tx1.mergeAccountRecord(acc2[0], acc2[1]).isOk()
+      tx0.fetchStateRoot() != tx1.fetchStateRoot()
+      tx0.kMap.len() == 0
+      tx1.kMap.len() == 1
+
+    # Check that the kMap hashkeys are moved correctly
+    # and that the stateroot is correct before and after each move.
+    let tx2 = db.txFrameBegin(tx1, moveParentHashKeys = true)
+    check:
+      tx1.kMap.len() == 0
+      tx2.kMap.len() == 1
+      tx1.fetchStateRoot() == tx2.fetchStateRoot()
+      # keys are recomputed when fetching state root even after moving
+      tx1.kMap.len() == 1
+      tx2.kMap.len() == 1
+
+    let tx3 = db.txFrameBegin(tx2, moveParentHashKeys = false)
+    check:
+      tx2.kMap.len() == 1
+      tx3.kMap.len() == 0
+      tx2.fetchStateRoot() == tx3.fetchStateRoot()
+
+  test "Frames using moveParentHashKeys parameter - moved from persist":
+    let
+      tx0 = db.txFrameBegin(db.baseTxFrame())
+      tx1 = db.txFrameBegin(tx0)
+
+    check:
+      tx0.mergeAccountRecord(acc1[0], acc1[1]).isOk()
+      tx1.mergeAccountRecord(acc2[0], acc2[1]).isOk()
+      tx0.fetchStateRoot() != tx1.fetchStateRoot()
+      tx0.kMap.len() == 0
+      tx1.kMap.len() == 1
+
+    let
+      tx2 = db.txFrameBegin(tx1, moveParentHashKeys = true)
+      stateRoot = tx1.fetchStateRoot().get()
+
+    # Check that we can still persist the moved from txFrame
+    tx1.checkpoint(1, skipSnapshot = true)
+    let batch = db.putBegFn().expect("working batch")
+    db.persist(batch, tx1, Opt.some(stateRoot))
+    check:
+      db.putEndFn(batch).isOk()
+
+    db.finish()
+
+    # Load the data
+    db.initInstance().expect("working backend")
+    let tx = db.baseTxFrame()
+    check:
+      tx.fetchAccountRecord(acc1[0]).isOk()
+      tx.fetchAccountRecord(acc2[0]).isOk()
+
+  test "Frames using moveParentHashKeys parameter - moved to persist":
+    let
+      tx0 = db.txFrameBegin(db.baseTxFrame())
+      tx1 = db.txFrameBegin(tx0)
+
+    check:
+      tx0.mergeAccountRecord(acc1[0], acc1[1]).isOk()
+      tx1.mergeAccountRecord(acc2[0], acc2[1]).isOk()
+      tx0.fetchStateRoot() != tx1.fetchStateRoot()
+      tx0.kMap.len() == 0
+      tx1.kMap.len() == 1
+
+    let
+      tx2 = db.txFrameBegin(tx1, moveParentHashKeys = true)
+      stateRoot = tx1.fetchStateRoot().get()
+
+    # Check that we can still persist the moved to txFrame
+    tx2.checkpoint(2, skipSnapshot = true)
+    let batch = db.putBegFn().expect("working batch")
+    db.persist(batch, tx2, Opt.some(stateRoot))
+    check:
+      db.putEndFn(batch).isOk()
+
+    db.finish()
+
+    # Load the data
+    db.initInstance().expect("working backend")
+    let tx = db.baseTxFrame()
+    check:
+      tx.fetchAccountRecord(acc1[0]).isOk()
+      tx.fetchAccountRecord(acc2[0]).isOk()