use state cache to avoid unnecessary block replay

Author: int88

beacon_node/store/src/hot_cold_store.rs

@@ -40,7 +40,10 @@ use std::sync::Arc;
 use std::time::Duration;
 use types::*;
 
-/// On-disk database that stores finalized states efficiently.
+/// The number of replayed states to be cached.
+const STATE_CACHE_SIZE: usize = 32;
+
+/// On-disk database that stores fnalized states efficiently.
 ///
 /// Stores vector fields like the `block_roots` and `state_roots` separately, and only stores
 /// intermittent "restore point" states pre-finalization.
@@ -62,6 +65,8 @@ pub struct HotColdDB<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> {
     pub hot_db: Hot,
     /// LRU cache of deserialized blocks. Updated whenever a block is loaded.
     block_cache: Mutex<LruCache<Hash256, SignedBeaconBlock<E>>>,
+    /// LRU cache of replayed states.
+    state_cache: Mutex<LruCache<Slot, BeaconState<E>>>,
     /// Chain spec.
     pub(crate) spec: ChainSpec,
     /// Logger.
@@ -129,6 +134,7 @@ impl<E: EthSpec> HotColdDB<E, MemoryStore<E>, MemoryStore<E>> {
             cold_db: MemoryStore::open(),
             hot_db: MemoryStore::open(),
             block_cache: Mutex::new(LruCache::new(config.block_cache_size)),
+            state_cache: Mutex::new(LruCache::new(STATE_CACHE_SIZE)),
             config,
             spec,
             log,
@@ -162,6 +168,7 @@ impl<E: EthSpec> HotColdDB<E, LevelDB<E>, LevelDB<E>> {
             cold_db: LevelDB::open(cold_path)?,
             hot_db: LevelDB::open(hot_path)?,
             block_cache: Mutex::new(LruCache::new(config.block_cache_size)),
+            state_cache: Mutex::new(LruCache::new(STATE_CACHE_SIZE)),
             config,
             spec,
             log,
@@ -579,6 +586,9 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
     /// (which are frozen, and won't be deleted), or valid descendents of the finalized checkpoint
     /// (which will be deleted by this function but shouldn't be).
     pub fn delete_state(&self, state_root: &Hash256, slot: Slot) -> Result<(), Error> {
+        // Delete the state from the cache.
+        self.state_cache.lock().pop(&slot);
+
         // Delete the state summary.
         self.hot_db
             .key_delete(DBColumn::BeaconStateSummary.into(), state_root.as_bytes())?;
@@ -977,40 +987,60 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
 
     /// Load a frozen state that lies between restore points.
     fn load_cold_intermediate_state(&self, slot: Slot) -> Result<BeaconState<E>, Error> {
+        if let Some(state) = self.state_cache.lock().get(&slot) {
+            return Ok(state.clone());
+        }
+
         // 1. Load the restore points either side of the intermediate state.
         let low_restore_point_idx = slot.as_u64() / self.config.slots_per_restore_point;
         let high_restore_point_idx = low_restore_point_idx + 1;
 
+        // Use low restore point as the base state.
+        let mut low_slot: Slot = Slot::new(low_restore_point_idx * self.config.slots_per_restore_point);
+        let mut low_state: BeaconState<E> = self.load_restore_point_by_index(low_restore_point_idx)?;
+
+        // Try to get a more recent state from the cache to avoid massive blocks replay.
+        for (s, state) in self.state_cache.lock().iter() {
+            if s.as_u64() / self.config.slots_per_restore_point == low_restore_point_idx && *s < slot && low_slot < *s {
+                low_slot = *s;
+                low_state = state.clone();
+            }
+        }
+
         // Acquire the read lock, so that the split can't change while this is happening.
         let split = self.split.read_recursive();
 
-        let low_restore_point = self.load_restore_point_by_index(low_restore_point_idx)?;
         let high_restore_point = self.get_restore_point(high_restore_point_idx, &split)?;
 
-        // 2. Load the blocks from the high restore point back to the low restore point.
+        // 2. Load the blocks from the high restore point back to the low point.
         let blocks = self.load_blocks_to_replay(
-            low_restore_point.slot(),
+            low_slot,
             slot,
             self.get_high_restore_point_block_root(&high_restore_point, slot)?,
         )?;
 
-        // 3. Replay the blocks on top of the low restore point.
+        // 3. Replay the blocks on top of the low point.
         // Use a forwards state root iterator to avoid doing any tree hashing.
         // The state root of the high restore point should never be used, so is safely set to 0.
         let state_root_iter = self.forwards_state_roots_iterator_until(
-            low_restore_point.slot(),
+            low_slot,
             slot,
             || (high_restore_point, Hash256::zero()),
             &self.spec,
         )?;
 
-        self.replay_blocks(
-            low_restore_point,
+        let state = self.replay_blocks(
+            low_state.clone(),
             blocks,
             slot,
             Some(state_root_iter),
             StateRootStrategy::Accurate,
-        )
+        )?;
+
+        // If state is not error, put it in the cache.
+        self.state_cache.lock().put(slot, state.clone());
+
+        Ok(state)
     }
 
     /// Get the restore point with the given index, or if it is out of bounds, the split state.