Store states efficiently in the hot database (#746)

* Sparse hot DB and block root tree

* Fix store_tests

* Ensure loads of hot states on boundaries are fast

* Milder error for unaligned finalized blocks

beacon_node/store/src/block_root_tree.rs

@@ -0,0 +1,364 @@
+use itertools::Itertools;
+use parking_lot::RwLock;
+use ssz_derive::{Decode, Encode};
+use std::collections::{HashMap, HashSet};
+use std::iter::{self, FromIterator};
+use types::{Hash256, Slot};
+
+/// In-memory cache of all block roots post-finalization. Includes short-lived forks.
+///
+/// Used by fork choice to avoid reconstructing hot states just for their block roots.
+// NOTE: could possibly be streamlined by combining with the head tracker and/or fork choice
+#[derive(Debug)]
+pub struct BlockRootTree {
+    nodes: RwLock<HashMap<Hash256, Node>>,
+}
+
+impl Clone for BlockRootTree {
+    fn clone(&self) -> Self {
+        Self {
+            nodes: RwLock::new(self.nodes.read().clone()),
+        }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub enum BlockRootTreeError {
+    PrevUnknown(Hash256),
+}
+
+/// Data for a single `block_root` in the tree.
+#[derive(Debug, Clone, Encode, Decode)]
+struct Node {
+    /// Hash of the preceding block (should be the parent block).
+    ///
+    /// A `previous` of `Hash256::zero` indicates the root of the tree.
+    previous: Hash256,
+    /// Slot of this node's block.
+    slot: Slot,
+}
+
+impl BlockRootTree {
+    /// Create a new block root tree where `(root_hash, root_slot)` is considered finalized.
+    ///
+    /// All subsequent blocks added should descend from the root block.
+    pub fn new(root_hash: Hash256, root_slot: Slot) -> Self {
+        Self {
+            nodes: RwLock::new(HashMap::from_iter(iter::once((
+                root_hash,
+                Node {
+                    previous: Hash256::zero(),
+                    slot: root_slot,
+                },
+            )))),
+        }
+    }
+
+    /// Check if `block_root` exists in the tree.
+    pub fn is_known_block_root(&self, block_root: &Hash256) -> bool {
+        self.nodes.read().contains_key(block_root)
+    }
+
+    /// Add a new `block_root` to the tree.
+    ///
+    /// Will return an error if `prev_block_root` doesn't exist in the tree.
+    pub fn add_block_root(
+        &self,
+        block_root: Hash256,
+        prev_block_root: Hash256,
+        block_slot: Slot,
+    ) -> Result<(), BlockRootTreeError> {
+        let mut nodes = self.nodes.write();
+        if nodes.contains_key(&prev_block_root) {
+            nodes.insert(
+                block_root,
+                Node {
+                    previous: prev_block_root,
+                    slot: block_slot,
+                },
+            );
+            Ok(())
+        } else {
+            Err(BlockRootTreeError::PrevUnknown(prev_block_root))
+        }
+    }
+
+    /// Create a reverse iterator from `block_root` (inclusive).
+    ///
+    /// Will skip slots, see `every_slot_iter_from` for a non-skipping variant.
+    pub fn iter_from(&self, block_root: Hash256) -> BlockRootTreeIter {
+        BlockRootTreeIter {
+            tree: self,
+            current_block_root: block_root,
+        }
+    }
+
+    /// Create a reverse iterator that yields a block root for every slot.
+    ///
+    /// E.g. if slot 6 is skipped, this iterator will return the block root from slot 5 at slot 6.
+    pub fn every_slot_iter_from<'a>(
+        &'a self,
+        block_root: Hash256,
+    ) -> impl Iterator<Item = (Hash256, Slot)> + 'a {
+        let mut block_roots = self.iter_from(block_root).peekable();
+
+        // Include the value for the first `block_root` if any, then fill in the skipped slots
+        // between each pair of previous block roots by duplicating the older root.
+        block_roots
+            .peek()
+            .cloned()
+            .into_iter()
+            .chain(block_roots.tuple_windows().flat_map(
+                |((_, high_slot), (low_hash, low_slot))| {
+                    (low_slot.as_u64()..high_slot.as_u64())
+                        .rev()
+                        .map(move |slot| (low_hash, Slot::new(slot)))
+                },
+            ))
+    }
+
+    /// Prune the tree.
+    ///
+    /// Only keep block roots descended from `finalized_root`, which lie on a chain leading
+    /// to one of the heads contained in `heads`.
+    pub fn prune_to(&self, finalized_root: Hash256, heads: impl IntoIterator<Item = Hash256>) {
+        let mut keep = HashSet::new();
+        keep.insert(finalized_root);
+
+        for head_block_root in heads.into_iter() {
+            // Iterate backwards until we reach a portion of the chain that we've already decided
+            // to keep. This also discards the pre-finalization block roots.
+            let mut keep_head = false;
+
+            let head_blocks = self
+                .iter_from(head_block_root)
+                .map(|(block_root, _)| block_root)
+                .inspect(|block_root| {
+                    if block_root == &finalized_root {
+                        keep_head = true;
+                    }
+                })
+                .take_while(|block_root| !keep.contains(&block_root))
+                .collect::<HashSet<_>>();
+
+            // If the head descends from the finalized root, keep it. Else throw it out.
+            if keep_head {
+                keep.extend(head_blocks);
+            }
+        }
+
+        self.nodes
+            .write()
+            .retain(|block_root, _| keep.contains(block_root));
+    }
+
+    pub fn as_ssz_container(&self) -> SszBlockRootTree {
+        SszBlockRootTree {
+            nodes: Vec::from_iter(self.nodes.read().clone()),
+        }
+    }
+}
+
+/// Simple (skipping) iterator for `BlockRootTree`.
+#[derive(Debug)]
+pub struct BlockRootTreeIter<'a> {
+    tree: &'a BlockRootTree,
+    current_block_root: Hash256,
+}
+
+impl<'a> Iterator for BlockRootTreeIter<'a> {
+    type Item = (Hash256, Slot);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Genesis
+        if self.current_block_root.is_zero() {
+            None
+        } else {
+            let block_root = self.current_block_root;
+            self.tree.nodes.read().get(&block_root).map(|node| {
+                self.current_block_root = node.previous;
+                (block_root, node.slot)
+            })
+        }
+    }
+}
+
+/// Serializable version of `BlockRootTree` that can be persisted to disk.
+#[derive(Debug, Clone, Encode, Decode)]
+pub struct SszBlockRootTree {
+    nodes: Vec<(Hash256, Node)>,
+}
+
+impl Into<BlockRootTree> for SszBlockRootTree {
+    fn into(self) -> BlockRootTree {
+        BlockRootTree {
+            nodes: RwLock::new(HashMap::from_iter(self.nodes)),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn int_hash(x: u64) -> Hash256 {
+        Hash256::from_low_u64_be(x)
+    }
+
+    fn check_iter_from(
+        block_tree: &BlockRootTree,
+        start_block_root: Hash256,
+        expected: &[(Hash256, Slot)],
+    ) {
+        assert_eq!(
+            &block_tree.iter_from(start_block_root).collect::<Vec<_>>()[..],
+            expected
+        );
+    }
+
+    fn check_every_slot_iter_from(
+        block_tree: &BlockRootTree,
+        start_block_root: Hash256,
+        expected: &[(Hash256, Slot)],
+    ) {
+        assert_eq!(
+            &block_tree
+                .every_slot_iter_from(start_block_root)
+                .collect::<Vec<_>>()[..],
+            expected
+        );
+    }
+
+    #[test]
+    fn single_chain() {
+        let block_tree = BlockRootTree::new(int_hash(1), Slot::new(1));
+        for i in 2..100 {
+            block_tree
+                .add_block_root(int_hash(i), int_hash(i - 1), Slot::new(i))
+                .expect("add_block_root ok");
+
+            let expected = (1..i + 1)
+                .rev()
+                .map(|j| (int_hash(j), Slot::new(j)))
+                .collect::<Vec<_>>();
+
+            check_iter_from(&block_tree, int_hash(i), &expected);
+            check_every_slot_iter_from(&block_tree, int_hash(i), &expected);
+
+            // Still OK after pruning.
+            block_tree.prune_to(int_hash(1), vec![int_hash(i)]);
+
+            check_iter_from(&block_tree, int_hash(i), &expected);
+            check_every_slot_iter_from(&block_tree, int_hash(i), &expected);
+        }
+    }
+
+    #[test]
+    fn skips_of_2() {
+        let block_tree = BlockRootTree::new(int_hash(1), Slot::new(1));
+        let step_length = 2u64;
+        for i in (1 + step_length..100).step_by(step_length as usize) {
+            block_tree
+                .add_block_root(int_hash(i), int_hash(i - step_length), Slot::new(i))
+                .expect("add_block_root ok");
+
+            let sparse_expected = (1..i + 1)
+                .rev()
+                .step_by(step_length as usize)
+                .map(|j| (int_hash(j), Slot::new(j)))
+                .collect_vec();
+            let every_slot_expected = (1..i + 1)
+                .rev()
+                .map(|j| {
+                    let nearest = 1 + (j - 1) / step_length * step_length;
+                    (int_hash(nearest), Slot::new(j))
+                })
+                .collect_vec();
+
+            check_iter_from(&block_tree, int_hash(i), &sparse_expected);
+            check_every_slot_iter_from(&block_tree, int_hash(i), &every_slot_expected);
+
+            // Still OK after pruning.
+            block_tree.prune_to(int_hash(1), vec![int_hash(i)]);
+
+            check_iter_from(&block_tree, int_hash(i), &sparse_expected);
+            check_every_slot_iter_from(&block_tree, int_hash(i), &every_slot_expected);
+        }
+    }
+
+    #[test]
+    fn prune_small_fork() {
+        let tree = BlockRootTree::new(int_hash(1), Slot::new(1));
+        // Space between fork hash values
+        let offset = 1000;
+        let num_blocks = 50;
+
+        let fork1_start = 2;
+        let fork2_start = 2 + offset;
+
+        tree.add_block_root(int_hash(fork1_start), int_hash(1), Slot::new(2))
+            .expect("add first block of left fork");
+        tree.add_block_root(int_hash(fork2_start), int_hash(1), Slot::new(2))
+            .expect("add first block of right fork");
+
+        for i in 3..num_blocks {
+            tree.add_block_root(int_hash(i), int_hash(i - 1), Slot::new(i))
+                .expect("add block to left fork");
+            tree.add_block_root(int_hash(i + offset), int_hash(i + offset - 1), Slot::new(i))
+                .expect("add block to right fork");
+        }
+
+        let root = (int_hash(1), Slot::new(1));
+
+        let (all_fork1_blocks, all_fork2_blocks): (Vec<_>, Vec<_>) = (2..num_blocks)
+            .rev()
+            .map(|i| {
+                (
+                    (int_hash(i), Slot::new(i)),
+                    (int_hash(i + offset), Slot::new(i)),
+                )
+            })
+            .chain(iter::once((root, root)))
+            .unzip();
+
+        let fork1_head = int_hash(num_blocks - 1);
+        let fork2_head = int_hash(num_blocks + offset - 1);
+
+        // Check that pruning with both heads preserves both chains.
+        let both_tree = tree.clone();
+        both_tree.prune_to(root.0, vec![fork1_head, fork2_head]);
+        check_iter_from(&both_tree, fork1_head, &all_fork1_blocks);
+        check_iter_from(&both_tree, fork2_head, &all_fork2_blocks);
+
+        // Check that pruning to either of the single chains leaves just that chain in the tree.
+        let fork1_tree = tree.clone();
+        fork1_tree.prune_to(root.0, vec![fork1_head]);
+        check_iter_from(&fork1_tree, fork1_head, &all_fork1_blocks);
+        check_iter_from(&fork1_tree, fork2_head, &[]);
+
+        let fork2_tree = tree.clone();
+        fork2_tree.prune_to(root.0, vec![fork2_head]);
+        check_iter_from(&fork2_tree, fork1_head, &[]);
+        check_iter_from(&fork2_tree, fork2_head, &all_fork2_blocks);
+
+        // Check that advancing the finalized root onto one side completely removes the other
+        // side.
+        let fin_tree = tree.clone();
+        let prune_point = num_blocks / 2;
+        let remaining_fork1_blocks = all_fork1_blocks
+            .clone()
+            .into_iter()
+            .take_while(|(_, slot)| *slot >= prune_point)
+            .collect_vec();
+        fin_tree.prune_to(int_hash(prune_point), vec![fork1_head, fork2_head]);
+        check_iter_from(&fin_tree, fork1_head, &remaining_fork1_blocks);
+        check_iter_from(&fin_tree, fork2_head, &[]);
+    }
+
+    #[test]
+    fn iter_zero() {
+        let block_tree = BlockRootTree::new(int_hash(0), Slot::new(0));
+        assert_eq!(block_tree.iter_from(int_hash(0)).count(), 0);
+        assert_eq!(block_tree.every_slot_iter_from(int_hash(0)).count(), 0);
+    }
+}

beacon_node/store/src/errors.rs

@@ -1,5 +1,5 @@
 use crate::chunked_vector::ChunkError;
-use crate::hot_cold_store::HotColdDbError;
+use crate::hot_cold_store::HotColdDBError;
 use ssz::DecodeError;
 use types::BeaconStateError;
 
@@ -9,7 +9,7 @@ pub enum Error {
     VectorChunkError(ChunkError),
     BeaconStateError(BeaconStateError),
     PartialBeaconStateError,
-    HotColdDbError(HotColdDbError),
+    HotColdDBError(HotColdDBError),
     DBError { message: String },
 }
 
@@ -25,9 +25,9 @@ impl From<ChunkError> for Error {
     }
 }
 
-impl From<HotColdDbError> for Error {
-    fn from(e: HotColdDbError) -> Error {
-        Error::HotColdDbError(e)
+impl From<HotColdDBError> for Error {
+    fn from(e: HotColdDBError) -> Error {
+        Error::HotColdDBError(e)
     }
 }
 

beacon_node/store/src/hot_cold_store.rs

@@ -50,7 +50,10 @@ pub struct HotColdDB<E: EthSpec> {
 }
 
 #[derive(Debug, PartialEq)]
-pub enum HotColdDbError {
+pub enum HotColdDBError {
+    /// Recoverable error indicating that the database freeze point couldn't be updated
+    /// due to the finalized block not lying on an epoch boundary (should be infrequent).
+    FreezeSlotUnaligned(Slot),
     FreezeSlotError {
         current_split_slot: Slot,
         proposed_split_slot: Slot,
@@ -58,13 +61,12 @@ pub enum HotColdDbError {
     MissingStateToFreeze(Hash256),
     MissingRestorePointHash(u64),
     MissingRestorePoint(Hash256),
-    MissingStateSlot(Hash256),
+    MissingColdStateSummary(Hash256),
+    MissingHotStateSummary(Hash256),
+    MissingEpochBoundaryState(Hash256),
     MissingSplitState(Hash256, Slot),
+    HotStateSummaryError(BeaconStateError),
     RestorePointDecodeError(ssz::DecodeError),
-    RestorePointReplayFailure {
-        expected_state_root: Hash256,
-        observed_state_root: Hash256,
-    },
     BlockReplayBeaconError(BeaconStateError),
     BlockReplaySlotError(SlotProcessingError),
     BlockReplayBlockError(BlockProcessingError),
@@ -98,9 +100,9 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
     /// Store a state in the store.
     fn put_state(&self, state_root: &Hash256, state: &BeaconState<E>) -> Result<(), Error> {
         if state.slot < self.get_split_slot() {
-            self.store_archive_state(state_root, state)
+            self.store_cold_state(state_root, state)
         } else {
-            self.hot_db.put_state(state_root, state)
+            self.store_hot_state(state_root, state)
         }
     }
 
@@ -112,20 +114,14 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
     ) -> Result<Option<BeaconState<E>>, Error> {
         if let Some(slot) = slot {
             if slot < self.get_split_slot() {
-                self.load_archive_state(state_root, slot).map(Some)
+                self.load_cold_state_by_slot(slot).map(Some)
             } else {
-                self.hot_db.get_state(state_root, None)
+                self.load_hot_state(state_root)
             }
         } else {
-            match self.hot_db.get_state(state_root, None)? {
+            match self.load_hot_state(state_root)? {
                 Some(state) => Ok(Some(state)),
-                None => {
-                    // Look-up the state in the freezer DB. We don't know the slot, so we must
-                    // look it up separately and then use it to reconstruct the state from a
-                    // restore point.
-                    let slot = self.load_state_slot(state_root)?;
-                    self.load_archive_state(state_root, slot).map(Some)
-                }
+                None => self.load_cold_state(state_root),
             }
         }
     }
@@ -142,17 +138,24 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
             "slot" => frozen_head.slot
         );
 
-        // 1. Copy all of the states between the head and the split slot, from the hot DB
-        // to the cold DB.
+        // 0. Check that the migration is sensible.
+        // The new frozen head must increase the current split slot, and lie on an epoch
+        // boundary (in order for the hot state summary scheme to work).
         let current_split_slot = store.get_split_slot();
 
         if frozen_head.slot < current_split_slot {
-            Err(HotColdDbError::FreezeSlotError {
+            Err(HotColdDBError::FreezeSlotError {
                 current_split_slot,
                 proposed_split_slot: frozen_head.slot,
             })?;
         }
 
+        if frozen_head.slot % E::slots_per_epoch() != 0 {
+            Err(HotColdDBError::FreezeSlotUnaligned(frozen_head.slot))?;
+        }
+
+        // 1. Copy all of the states between the head and the split slot, from the hot DB
+        // to the cold DB.
         let state_root_iter = StateRootsIterator::new(store.clone(), frozen_head);
 
         let mut to_delete = vec![];
@@ -163,16 +166,20 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
                 let state: BeaconState<E> = store
                     .hot_db
                     .get_state(&state_root, None)?
-                    .ok_or_else(|| HotColdDbError::MissingStateToFreeze(state_root))?;
+                    .ok_or_else(|| HotColdDBError::MissingStateToFreeze(state_root))?;
 
-                store.store_archive_state(&state_root, &state)?;
+                store.store_cold_state(&state_root, &state)?;
             }
 
             // Store a pointer from this state root to its slot, so we can later reconstruct states
             // from their state root alone.
-            store.store_state_slot(&state_root, slot)?;
+            store.store_cold_state_slot(&state_root, slot)?;
 
-            to_delete.push(state_root);
+            // Delete the old summary, and the full state if we lie on an epoch boundary.
+            to_delete.push((DBColumn::BeaconStateSummary, state_root));
+            if slot % E::slots_per_epoch() == 0 {
+                to_delete.push((DBColumn::BeaconState, state_root));
+            }
         }
 
         // 2. Update the split slot
@@ -183,10 +190,10 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
         store.store_split()?;
 
         // 3. Delete from the hot DB
-        for state_root in to_delete {
+        for (column, state_root) in to_delete {
             store
                 .hot_db
-                .key_delete(DBColumn::BeaconState.into(), state_root.as_bytes())?;
+                .key_delete(column.into(), state_root.as_bytes())?;
         }
 
         debug!(
@@ -207,6 +214,38 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
     ) -> Self::ForwardsBlockRootsIterator {
         HybridForwardsBlockRootsIterator::new(store, start_slot, end_state, end_block_root, spec)
     }
+
+    /// Load an epoch boundary state by using the hot state summary look-up.
+    ///
+    /// Will fall back to the cold DB if a hot state summary is not found.
+    fn load_epoch_boundary_state(
+        &self,
+        state_root: &Hash256,
+    ) -> Result<Option<BeaconState<E>>, Error> {
+        if let Some(HotStateSummary {
+            epoch_boundary_state_root,
+            ..
+        }) = self.load_hot_state_summary(state_root)?
+        {
+            let state = self
+                .hot_db
+                .get_state(&epoch_boundary_state_root, None)?
+                .ok_or_else(|| {
+                    HotColdDBError::MissingEpochBoundaryState(epoch_boundary_state_root)
+                })?;
+            Ok(Some(state))
+        } else {
+            // Try the cold DB
+            match self.load_cold_state_slot(state_root)? {
+                Some(state_slot) => {
+                    let epoch_boundary_slot =
+                        state_slot / E::slots_per_epoch() * E::slots_per_epoch();
+                    self.load_cold_state_by_slot(epoch_boundary_slot).map(Some)
+                }
+                None => Ok(None),
+            }
+        }
+    }
 }
 
 impl<E: EthSpec> HotColdDB<E> {
@@ -240,10 +279,69 @@ impl<E: EthSpec> HotColdDB<E> {
         Ok(db)
     }
 
+    /// Store a post-finalization state efficiently in the hot database.
+    ///
+    /// On an epoch boundary, store a full state. On an intermediate slot, store
+    /// just a backpointer to the nearest epoch boundary.
+    pub fn store_hot_state(
+        &self,
+        state_root: &Hash256,
+        state: &BeaconState<E>,
+    ) -> Result<(), Error> {
+        // On the epoch boundary, store the full state.
+        if state.slot % E::slots_per_epoch() == 0 {
+            trace!(
+                self.log,
+                "Storing full state on epoch boundary";
+                "slot" => state.slot.as_u64(),
+                "state_root" => format!("{:?}", state_root)
+            );
+            self.hot_db.put_state(state_root, state)?;
+        }
+
+        // Store a summary of the state.
+        // We store one even for the epoch boundary states, as we may need their slots
+        // when doing a look up by state root.
+        self.store_hot_state_summary(state_root, state)?;
+
+        Ok(())
+    }
+
+    /// Load a post-finalization state from the hot database.
+    ///
+    /// Will replay blocks from the nearest epoch boundary.
+    pub fn load_hot_state(&self, state_root: &Hash256) -> Result<Option<BeaconState<E>>, Error> {
+        if let Some(HotStateSummary {
+            slot,
+            latest_block_root,
+            epoch_boundary_state_root,
+        }) = self.load_hot_state_summary(state_root)?
+        {
+            let state: BeaconState<E> = self
+                .hot_db
+                .get_state(&epoch_boundary_state_root, None)?
+                .ok_or_else(|| {
+                    HotColdDBError::MissingEpochBoundaryState(epoch_boundary_state_root)
+                })?;
+
+            // Optimization to avoid even *thinking* about replaying blocks if we're already
+            // on an epoch boundary.
+            if slot % E::slots_per_epoch() == 0 {
+                Ok(Some(state))
+            } else {
+                let blocks = self.load_blocks_to_replay(state.slot, slot, latest_block_root)?;
+                self.replay_blocks(state, blocks, slot).map(Some)
+            }
+        } else {
+            Ok(None)
+        }
+    }
+
     /// Store a pre-finalization state in the freezer database.
     ///
-    /// Will return an error if the state does not lie on a restore point boundary.
-    pub fn store_archive_state(
+    /// Will log a warning and not store anything if the state does not lie on a restore point
+    /// boundary.
+    pub fn store_cold_state(
         &self,
         state_root: &Hash256,
         state: &BeaconState<E>,
@@ -283,25 +381,32 @@ impl<E: EthSpec> HotColdDB<E> {
         Ok(())
     }
 
+    /// Try to load a pre-finalization state from the freezer database.
+    ///
+    /// Return `None` if no state with `state_root` lies in the freezer.
+    pub fn load_cold_state(&self, state_root: &Hash256) -> Result<Option<BeaconState<E>>, Error> {
+        match self.load_cold_state_slot(state_root)? {
+            Some(slot) => self.load_cold_state_by_slot(slot).map(Some),
+            None => Ok(None),
+        }
+    }
+
     /// Load a pre-finalization state from the freezer database.
     ///
     /// Will reconstruct the state if it lies between restore points.
-    pub fn load_archive_state(
-        &self,
-        state_root: &Hash256,
-        slot: Slot,
-    ) -> Result<BeaconState<E>, Error> {
+    pub fn load_cold_state_by_slot(&self, slot: Slot) -> Result<BeaconState<E>, Error> {
         if slot % self.slots_per_restore_point == 0 {
-            self.load_restore_point(state_root)
+            let restore_point_idx = slot.as_u64() / self.slots_per_restore_point;
+            self.load_restore_point_by_index(restore_point_idx)
         } else {
-            self.load_intermediate_state(state_root, slot)
+            self.load_cold_intermediate_state(slot)
         }
     }
 
     /// Load a restore point state by its `state_root`.
     fn load_restore_point(&self, state_root: &Hash256) -> Result<BeaconState<E>, Error> {
         let mut partial_state = PartialBeaconState::db_get(&self.cold_db, state_root)?
-            .ok_or_else(|| HotColdDbError::MissingRestorePoint(*state_root))?;
+            .ok_or_else(|| HotColdDBError::MissingRestorePoint(*state_root))?;
 
         // Fill in the fields of the partial state.
         partial_state.load_block_roots(&self.cold_db, &self.spec)?;
@@ -321,12 +426,8 @@ impl<E: EthSpec> HotColdDB<E> {
         self.load_restore_point(&state_root)
     }
 
-    /// Load a state that lies between restore points.
-    fn load_intermediate_state(
-        &self,
-        state_root: &Hash256,
-        slot: Slot,
-    ) -> Result<BeaconState<E>, Error> {
+    /// Load a frozen state that lies between restore points.
+    fn load_cold_intermediate_state(&self, slot: Slot) -> Result<BeaconState<E>, Error> {
         // 1. Load the restore points either side of the intermediate state.
         let low_restore_point_idx = slot.as_u64() / self.slots_per_restore_point;
         let high_restore_point_idx = low_restore_point_idx + 1;
@@ -341,7 +442,7 @@ impl<E: EthSpec> HotColdDB<E> {
             >= split.slot.as_u64()
         {
             self.get_state(&split.state_root, Some(split.slot))?
-                .ok_or_else(|| HotColdDbError::MissingSplitState(split.state_root, split.slot))?
+                .ok_or_else(|| HotColdDBError::MissingSplitState(split.state_root, split.slot))?
         } else {
             self.load_restore_point_by_index(high_restore_point_idx)?
         };
@@ -354,22 +455,7 @@ impl<E: EthSpec> HotColdDB<E> {
         )?;
 
         // 3. Replay the blocks on top of the low restore point.
-        let mut state = self.replay_blocks(low_restore_point, blocks, slot)?;
-
-        // 4. Check that the state root is correct (should be quick with the cache already built).
-        // TODO: we could optimise out *all* the tree hashing when replaying blocks,
-        // in which case we could also drop this check.
-        let observed_state_root = state.update_tree_hash_cache()?;
-
-        if observed_state_root == *state_root {
-            Ok(state)
-        } else {
-            Err(HotColdDbError::RestorePointReplayFailure {
-                expected_state_root: *state_root,
-                observed_state_root,
-            }
-            .into())
-        }
+        self.replay_blocks(low_restore_point, blocks, slot)
     }
 
     /// Get a suitable block root for backtracking from `high_restore_point` to the state at `slot`.
@@ -379,12 +465,12 @@ impl<E: EthSpec> HotColdDB<E> {
         &self,
         high_restore_point: &BeaconState<E>,
         slot: Slot,
-    ) -> Result<Hash256, HotColdDbError> {
+    ) -> Result<Hash256, HotColdDBError> {
         high_restore_point
             .get_block_root(slot)
             .or_else(|_| high_restore_point.get_oldest_block_root())
             .map(|x| *x)
-            .map_err(HotColdDbError::RestorePointBlockHashError)
+            .map_err(HotColdDBError::RestorePointBlockHashError)
     }
 
     /// Load the blocks between `start_slot` and `end_slot` by backtracking from `end_block_hash`.
@@ -398,6 +484,7 @@ impl<E: EthSpec> HotColdDB<E> {
         end_block_hash: Hash256,
     ) -> Result<Vec<BeaconBlock<E>>, Error> {
         let mut blocks = ParentRootBlockIterator::new(self, end_block_hash)
+            .map(|(_, block)| block)
             // Include the block at the end slot (if any), it needs to be
             // replayed in order to construct the canonical state at `end_slot`.
             .filter(|block| block.slot <= end_slot)
@@ -420,12 +507,26 @@ impl<E: EthSpec> HotColdDB<E> {
     ) -> Result<BeaconState<E>, Error> {
         state
             .build_all_caches(&self.spec)
-            .map_err(HotColdDbError::BlockReplayBeaconError)?;
+            .map_err(HotColdDBError::BlockReplayBeaconError)?;
 
-        for block in blocks {
+        let state_root_from_prev_block = |i: usize, state: &BeaconState<E>| {
+            if i > 0 {
+                let prev_block = &blocks[i - 1];
+                if prev_block.slot == state.slot {
+                    Some(prev_block.state_root)
+                } else {
+                    None
+                }
+            } else {
+                None
+            }
+        };
+
+        for (i, block) in blocks.iter().enumerate() {
             while state.slot < block.slot {
-                per_slot_processing(&mut state, None, &self.spec)
-                    .map_err(HotColdDbError::BlockReplaySlotError)?;
+                let state_root = state_root_from_prev_block(i, &state);
+                per_slot_processing(&mut state, state_root, &self.spec)
+                    .map_err(HotColdDBError::BlockReplaySlotError)?;
             }
             per_block_processing(
                 &mut state,
@@ -434,12 +535,13 @@ impl<E: EthSpec> HotColdDB<E> {
                 BlockSignatureStrategy::NoVerification,
                 &self.spec,
             )
-            .map_err(HotColdDbError::BlockReplayBlockError)?;
+            .map_err(HotColdDBError::BlockReplayBlockError)?;
         }
 
         while state.slot < target_slot {
-            per_slot_processing(&mut state, None, &self.spec)
-                .map_err(HotColdDbError::BlockReplaySlotError)?;
+            let state_root = state_root_from_prev_block(blocks.len(), &state);
+            per_slot_processing(&mut state, state_root, &self.spec)
+                .map_err(HotColdDBError::BlockReplaySlotError)?;
         }
 
         Ok(state)
@@ -474,7 +576,7 @@ impl<E: EthSpec> HotColdDB<E> {
         let key = Self::restore_point_key(restore_point_index);
         RestorePointHash::db_get(&self.cold_db, &key)?
             .map(|r| r.state_root)
-            .ok_or(HotColdDbError::MissingRestorePointHash(restore_point_index).into())
+            .ok_or(HotColdDBError::MissingRestorePointHash(restore_point_index).into())
     }
 
     /// Store the state root of a restore point.
@@ -495,30 +597,63 @@ impl<E: EthSpec> HotColdDB<E> {
     }
 
     /// Load a frozen state's slot, given its root.
-    fn load_state_slot(&self, state_root: &Hash256) -> Result<Slot, Error> {
-        StateSlot::db_get(&self.cold_db, state_root)?
-            .map(|s| s.slot)
-            .ok_or_else(|| HotColdDbError::MissingStateSlot(*state_root).into())
+    fn load_cold_state_slot(&self, state_root: &Hash256) -> Result<Option<Slot>, Error> {
+        Ok(ColdStateSummary::db_get(&self.cold_db, state_root)?.map(|s| s.slot))
     }
 
     /// Store the slot of a frozen state.
-    fn store_state_slot(&self, state_root: &Hash256, slot: Slot) -> Result<(), Error> {
-        StateSlot { slot }
+    fn store_cold_state_slot(&self, state_root: &Hash256, slot: Slot) -> Result<(), Error> {
+        ColdStateSummary { slot }
             .db_put(&self.cold_db, state_root)
             .map_err(Into::into)
     }
 
+    /// Load a hot state's summary, given its root.
+    pub fn load_hot_state_summary(
+        &self,
+        state_root: &Hash256,
+    ) -> Result<Option<HotStateSummary>, Error> {
+        HotStateSummary::db_get(&self.hot_db, state_root)
+    }
+
+    /// Store a summary of a hot database state.
+    fn store_hot_state_summary(
+        &self,
+        state_root: &Hash256,
+        state: &BeaconState<E>,
+    ) -> Result<(), Error> {
+        // Fill in the state root on the latest block header if necessary (this happens on all
+        // slots where there isn't a skip).
+        let latest_block_root = state.get_latest_block_root(*state_root);
+        let epoch_boundary_slot = state.slot / E::slots_per_epoch() * E::slots_per_epoch();
+        let epoch_boundary_state_root = if epoch_boundary_slot == state.slot {
+            *state_root
+        } else {
+            *state
+                .get_state_root(epoch_boundary_slot)
+                .map_err(HotColdDBError::HotStateSummaryError)?
+        };
+
+        HotStateSummary {
+            slot: state.slot,
+            latest_block_root,
+            epoch_boundary_state_root,
+        }
+        .db_put(&self.hot_db, state_root)
+        .map_err(Into::into)
+    }
+
     /// Check that the restore point frequency is a divisor of the slots per historical root.
     ///
     /// This ensures that we have at least one restore point within range of our state
     /// root history when iterating backwards (and allows for more frequent restore points if
     /// desired).
-    fn verify_slots_per_restore_point(slots_per_restore_point: u64) -> Result<(), HotColdDbError> {
+    fn verify_slots_per_restore_point(slots_per_restore_point: u64) -> Result<(), HotColdDBError> {
         let slots_per_historical_root = E::SlotsPerHistoricalRoot::to_u64();
         if slots_per_restore_point > 0 && slots_per_historical_root % slots_per_restore_point == 0 {
             Ok(())
         } else {
-            Err(HotColdDbError::InvalidSlotsPerRestorePoint {
+            Err(HotColdDBError::InvalidSlotsPerRestorePoint {
                 slots_per_restore_point,
                 slots_per_historical_root,
             })
@@ -527,7 +662,7 @@ impl<E: EthSpec> HotColdDB<E> {
 }
 
 /// Struct for storing the split slot and state root in the database.
-#[derive(Clone, Copy, Default, Encode, Decode)]
+#[derive(Debug, Clone, Copy, Default, Encode, Decode)]
 struct Split {
     slot: Slot,
     state_root: Hash256,
@@ -547,15 +682,39 @@ impl SimpleStoreItem for Split {
     }
 }
 
-/// Struct for storing the slot of a state root in the database.
-#[derive(Clone, Copy, Default, Encode, Decode)]
-struct StateSlot {
+/// Struct for summarising a state in the hot database.
+///
+/// Allows full reconstruction by replaying blocks.
+#[derive(Debug, Clone, Copy, Default, Encode, Decode)]
+pub struct HotStateSummary {
+    slot: Slot,
+    latest_block_root: Hash256,
+    epoch_boundary_state_root: Hash256,
+}
+
+impl SimpleStoreItem for HotStateSummary {
+    fn db_column() -> DBColumn {
+        DBColumn::BeaconStateSummary
+    }
+
+    fn as_store_bytes(&self) -> Vec<u8> {
+        self.as_ssz_bytes()
+    }
+
+    fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        Ok(Self::from_ssz_bytes(bytes)?)
+    }
+}
+
+/// Struct for summarising a state in the freezer database.
+#[derive(Debug, Clone, Copy, Default, Encode, Decode)]
+struct ColdStateSummary {
     slot: Slot,
 }
 
-impl SimpleStoreItem for StateSlot {
+impl SimpleStoreItem for ColdStateSummary {
     fn db_column() -> DBColumn {
-        DBColumn::BeaconStateSlot
+        DBColumn::BeaconStateSummary
     }
 
     fn as_store_bytes(&self) -> Vec<u8> {
@@ -568,7 +727,7 @@ impl SimpleStoreItem for StateSlot {
 }
 
 /// Struct for storing the state root of a restore point in the database.
-#[derive(Clone, Copy, Default, Encode, Decode)]
+#[derive(Debug, Clone, Copy, Default, Encode, Decode)]
 struct RestorePointHash {
     state_root: Hash256,
 }

beacon_node/store/src/iter.rs

@@ -120,7 +120,7 @@ impl<'a, E: EthSpec, S: Store<E>> ParentRootBlockIterator<'a, E, S> {
 }
 
 impl<'a, E: EthSpec, S: Store<E>> Iterator for ParentRootBlockIterator<'a, E, S> {
-    type Item = BeaconBlock<E>;
+    type Item = (Hash256, BeaconBlock<E>);
 
     fn next(&mut self) -> Option<Self::Item> {
         // Stop once we reach the zero parent, otherwise we'll keep returning the genesis
@@ -128,9 +128,10 @@ impl<'a, E: EthSpec, S: Store<E>> Iterator for ParentRootBlockIterator<'a, E, S>
         if self.next_block_root.is_zero() {
             None
         } else {
-            let block: BeaconBlock<E> = self.store.get(&self.next_block_root).ok()??;
+            let block_root = self.next_block_root;
+            let block: BeaconBlock<E> = self.store.get(&block_root).ok()??;
             self.next_block_root = block.parent_root;
-            Some(block)
+            Some((block_root, block))
         }
     }
 }

beacon_node/store/src/lib.rs

@@ -11,6 +11,7 @@
 extern crate lazy_static;
 
 mod block_at_slot;
+pub mod block_root_tree;
 pub mod chunked_iter;
 pub mod chunked_vector;
 pub mod config;
@@ -28,6 +29,7 @@ pub mod migrate;
 
 use std::sync::Arc;
 
+pub use self::block_root_tree::{BlockRootTree, SszBlockRootTree};
 pub use self::config::StoreConfig;
 pub use self::hot_cold_store::HotColdDB as DiskStore;
 pub use self::leveldb_store::LevelDB as SimpleDiskStore;
@@ -128,6 +130,29 @@ pub trait Store<E: EthSpec>: Sync + Send + Sized + 'static {
         end_block_root: Hash256,
         spec: &ChainSpec,
     ) -> Self::ForwardsBlockRootsIterator;
+
+    /// Load the most recent ancestor state of `state_root` which lies on an epoch boundary.
+    ///
+    /// If `state_root` corresponds to an epoch boundary state, then that state itself should be
+    /// returned.
+    fn load_epoch_boundary_state(
+        &self,
+        state_root: &Hash256,
+    ) -> Result<Option<BeaconState<E>>, Error> {
+        // The default implementation is not very efficient, but isn't used in prod.
+        // See `HotColdDB` for the optimized implementation.
+        if let Some(state) = self.get_state(state_root, None)? {
+            let epoch_boundary_slot = state.slot / E::slots_per_epoch() * E::slots_per_epoch();
+            if state.slot == epoch_boundary_slot {
+                Ok(Some(state))
+            } else {
+                let epoch_boundary_state_root = state.get_state_root(epoch_boundary_slot)?;
+                self.get_state(epoch_boundary_state_root, Some(epoch_boundary_slot))
+            }
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 /// A unique column identifier.
@@ -140,8 +165,8 @@ pub enum DBColumn {
     BeaconChain,
     /// For the table mapping restore point numbers to state roots.
     BeaconRestorePoint,
-    /// For the mapping from state roots to their slots.
-    BeaconStateSlot,
+    /// For the mapping from state roots to their slots or summaries.
+    BeaconStateSummary,
     BeaconBlockRoots,
     BeaconStateRoots,
     BeaconHistoricalRoots,
@@ -157,7 +182,7 @@ impl Into<&'static str> for DBColumn {
             DBColumn::BeaconState => "ste",
             DBColumn::BeaconChain => "bch",
             DBColumn::BeaconRestorePoint => "brp",
-            DBColumn::BeaconStateSlot => "bss",
+            DBColumn::BeaconStateSummary => "bss",
             DBColumn::BeaconBlockRoots => "bbr",
             DBColumn::BeaconStateRoots => "bsr",
             DBColumn::BeaconHistoricalRoots => "bhr",

beacon_node/store/src/migrate.rs

@@ -1,6 +1,8 @@
-use crate::{DiskStore, MemoryStore, SimpleDiskStore, Store};
+use crate::{
+    hot_cold_store::HotColdDBError, DiskStore, Error, MemoryStore, SimpleDiskStore, Store,
+};
 use parking_lot::Mutex;
-use slog::warn;
+use slog::{info, warn};
 use std::mem;
 use std::sync::mpsc;
 use std::sync::Arc;
@@ -127,12 +129,22 @@ impl<E: EthSpec> BackgroundMigrator<E> {
         let (tx, rx) = mpsc::channel();
         let thread = thread::spawn(move || {
             while let Ok((state_root, state)) = rx.recv() {
-                if let Err(e) = DiskStore::freeze_to_state(db.clone(), state_root, &state) {
-                    warn!(
-                        db.log,
-                        "Database migration failed";
-                        "error" => format!("{:?}", e)
-                    );
+                match DiskStore::freeze_to_state(db.clone(), state_root, &state) {
+                    Ok(()) => {}
+                    Err(Error::HotColdDBError(HotColdDBError::FreezeSlotUnaligned(slot))) => {
+                        info!(
+                            db.log,
+                            "Database migration postponed, unaligned finalized block";
+                            "slot" => slot.as_u64()
+                        );
+                    }
+                    Err(e) => {
+                        warn!(
+                            db.log,
+                            "Database migration failed";
+                            "error" => format!("{:?}", e)
+                        );
+                    }
                 }
             }
         });