save some more states after all (#1887)

Don't save states when replaying history, but do save states when
applying new blocks (!)

beacon_chain/block_pools/chain_dag.nim

@@ -36,7 +36,7 @@ proc putBlock*(
   dag.db.putBlock(signedBlock)
 
 proc updateStateData*(
-  dag: ChainDAGRef, state: var StateData, bs: BlockSlot,
+  dag: ChainDAGRef, state: var StateData, bs: BlockSlot, save: bool,
   cache: var StateCache) {.gcsafe.}
 
 template withState*(
@@ -48,7 +48,7 @@ template withState*(
   ##      while waiting for future to complete - catch this here somehow?
 
   var cache {.inject.} = blockSlot.blck.getStateCache(blockSlot.slot.epoch())
-  updateStateData(dag, stateData, blockSlot, cache)
+  updateStateData(dag, stateData, blockSlot, false, cache)
 
   template hashedState(): HashedBeaconState {.inject, used.} = stateData.data
   template state(): BeaconState {.inject, used.} = stateData.data.data
@@ -391,7 +391,7 @@ proc init*(T: type ChainDAGRef,
   doAssert res.updateFlags in [{}, {verifyFinalization}]
 
   var cache: StateCache
-  res.updateStateData(res.headState, headRef.atSlot(headRef.slot), cache)
+  res.updateStateData(res.headState, headRef.atSlot(headRef.slot), false, cache)
   # We presently save states on the epoch boundary - it means that the latest
   # state we loaded might be older than head block - nonetheless, it will be
   # from the same epoch as the head, thus the finalized and justified slots are
@@ -601,34 +601,29 @@ proc get*(dag: ChainDAGRef, root: Eth2Digest): Option[BlockData] =
     none(BlockData)
 
 proc advanceSlots(
-    dag: ChainDAGRef, state: var StateData, slot: Slot, cache: var StateCache) =
+    dag: ChainDAGRef, state: var StateData, slot: Slot, save: bool,
+    cache: var StateCache) =
   # Given a state, advance it zero or more slots by applying empty slot
   # processing - the state must be positions at a slot before or equal to the
   # target
   doAssert state.data.data.slot <= slot
-  if slot > state.data.data.slot:
-    doAssert process_slots(state.data, slot, cache, dag.updateFlags),
+  while state.data.data.slot < slot:
+    doAssert process_slots(
+        state.data, state.data.data.slot + 1, cache,
+        dag.updateFlags),
       "process_slots shouldn't fail when state slot is correct"
+    if save:
+      dag.putState(state)
 
 proc applyBlock(
     dag: ChainDAGRef,
     state: var StateData, blck: BlockData, flags: UpdateFlags,
-    cache: var StateCache, save: bool): bool =
+    cache: var StateCache): bool =
   # Apply a single block to the state - the state must be positioned at the
   # parent of the block with a slot lower than the one of the block being
   # applied
   doAssert state.blck == blck.refs.parent
 
-  # `state_transition` can handle empty slots, but we want to save the state
-  # before applying the block
-  dag.advanceSlots(state, blck.data.message.slot, cache)
-
-  if save:
-    # Save state before applying the block, in case the "raw" epoch state is
-    # needed for a different fork
-    # TODO if the block fails to apply, it can be removed from the database
-    dag.putState(state)
-
   var statePtr = unsafeAddr state # safe because `restore` is locally scoped
   func restore(v: var HashedBeaconState) =
     doAssert (addr(statePtr.data) == addr v)
@@ -643,7 +638,7 @@ proc applyBlock(
   ok
 
 proc updateStateData*(
-    dag: ChainDAGRef, state: var StateData, bs: BlockSlot,
+    dag: ChainDAGRef, state: var StateData, bs: BlockSlot, save: bool,
     cache: var StateCache) =
   ## Rewind or advance state such that it matches the given block and slot -
   ## this may include replaying from an earlier snapshot if blck is on a
@@ -658,7 +653,7 @@ proc updateStateData*(
   if state.blck == bs.blck and state.data.data.slot <= bs.slot:
     # The block is the same and we're at an early enough slot - advance the
     # state with empty slot processing until the slot is correct
-    dag.advanceSlots(state, bs.slot, cache)
+    dag.advanceSlots(state, bs.slot, save, cache)
 
     return
 
@@ -703,11 +698,11 @@ proc updateStateData*(
     # database, we can skip certain checks that have already been performed
     # before adding the block to the database.
     let ok =
-      dag.applyBlock(state, dag.get(ancestors[i]), {}, cache, false)
+      dag.applyBlock(state, dag.get(ancestors[i]), {}, cache)
     doAssert ok, "Blocks in database should never fail to apply.."
 
   # ...and make sure to process empty slots as requested
-  dag.advanceSlots(state, bs.slot, cache)
+  dag.advanceSlots(state, bs.slot, save, cache)
 
   beacon_state_rewinds.inc()
 
@@ -760,7 +755,7 @@ proc updateHead*(
   else:
     var cache = getStateCache(newHead, newHead.slot.epoch())
     updateStateData(
-      dag, dag.headState, newHead.atSlot(newHead.slot), cache)
+      dag, dag.headState, newHead.atSlot(newHead.slot), false, cache)
 
   dag.head = newHead
 

beacon_chain/block_pools/clearance.nim

@@ -198,7 +198,7 @@ proc addRawBlock*(
     #      but maybe we should use it as a hint that our clock is wrong?
     var cache = getStateCache(parent, blck.slot.epoch)
     updateStateData(
-      dag, dag.clearanceState, parent.atSlot(blck.slot), cache)
+      dag, dag.clearanceState, parent.atSlot(blck.slot), true, cache)
 
     let
       poolPtr = unsafeAddr dag # safe because restore is short-lived

research/block_sim.nim

@@ -176,7 +176,7 @@ cli do(slots = SLOTS_PER_EPOCH * 6,
     withTimer(timers[tReplay]):
       var cache = StateCache()
       chainDag.updateStateData(
-        replayState[], chainDag.head.atSlot(Slot(slots)), cache)
+        replayState[], chainDag.head.atSlot(Slot(slots)), false, cache)
 
   echo "Done!"
 

tests/test_block_pool.nim

@@ -292,39 +292,39 @@ suiteReport "Block pool processing" & preset():
 
     # move to specific block
     var cache = StateCache()
-    dag.updateStateData(tmpState[], bs1, cache)
+    dag.updateStateData(tmpState[], bs1, false, cache)
 
     check:
       tmpState.blck == b1Add[]
       tmpState.data.data.slot == bs1.slot
 
     # Skip slots
-    dag.updateStateData(tmpState[], bs1_3, cache) # skip slots
+    dag.updateStateData(tmpState[], bs1_3, false, cache) # skip slots
 
     check:
       tmpState.blck == b1Add[]
       tmpState.data.data.slot == bs1_3.slot
 
     # Move back slots, but not blocks
-    dag.updateStateData(tmpState[], bs1_3.parent(), cache)
+    dag.updateStateData(tmpState[], bs1_3.parent(), false, cache)
     check:
       tmpState.blck == b1Add[]
       tmpState.data.data.slot == bs1_3.parent().slot
 
     # Move to different block and slot
-    dag.updateStateData(tmpState[], bs2_3, cache)
+    dag.updateStateData(tmpState[], bs2_3, false, cache)
     check:
       tmpState.blck == b2Add[]
       tmpState.data.data.slot == bs2_3.slot
 
     # Move back slot and block
-    dag.updateStateData(tmpState[], bs1, cache)
+    dag.updateStateData(tmpState[], bs1, false, cache)
     check:
       tmpState.blck == b1Add[]
       tmpState.data.data.slot == bs1.slot
 
     # Move back to genesis
-    dag.updateStateData(tmpState[], bs1.parent(), cache)
+    dag.updateStateData(tmpState[], bs1.parent(), false, cache)
     check:
       tmpState.blck == b1Add[].parent
       tmpState.data.data.slot == bs1.parent.slot