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/beacon_chain/src/beacon_chain.rs

@@ -31,7 +31,7 @@ use std::time::{Duration, Instant};
 use store::iter::{
     BlockRootsIterator, ReverseBlockRootIterator, ReverseStateRootIterator, StateRootsIterator,
 };
-use store::{Error as DBError, Migrate, Store};
+use store::{BlockRootTree, Error as DBError, Migrate, Store};
 use tree_hash::TreeHash;
 use types::*;
 
@@ -149,6 +149,8 @@ pub struct BeaconChain<T: BeaconChainTypes> {
     pub(crate) head_tracker: HeadTracker,
     /// Provides a small cache of `BeaconState` and `BeaconBlock`.
     pub(crate) checkpoint_cache: CheckPointCache<T::EthSpec>,
+    /// Cache of block roots for all known forks post-finalization.
+    pub block_root_tree: Arc<BlockRootTree>,
     /// Logging to CLI, etc.
     pub(crate) log: Logger,
 }
@@ -170,8 +172,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 .ok_or_else(|| Error::MissingBeaconBlock(beacon_block_root))?;
             let beacon_state_root = beacon_block.state_root;
             let beacon_state = self
-                .store
-                .get_state(&beacon_state_root, Some(beacon_block.slot))?
+                .get_state_caching(&beacon_state_root, Some(beacon_block.slot))?
                 .ok_or_else(|| Error::MissingBeaconState(beacon_state_root))?;
 
             CheckPoint {
@@ -189,6 +190,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             genesis_block_root: self.genesis_block_root,
             ssz_head_tracker: self.head_tracker.to_ssz_container(),
             fork_choice: self.fork_choice.as_ssz_container(),
+            block_root_tree: self.block_root_tree.as_ssz_container(),
         };
 
         let key = Hash256::from_slice(&BEACON_CHAIN_DB_KEY.as_bytes());
@@ -411,7 +413,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// ## Errors
     ///
     /// May return a database error.
-    fn get_block_caching(
+    pub fn get_block_caching(
         &self,
         block_root: &Hash256,
     ) -> Result<Option<BeaconBlock<T::EthSpec>>, Error> {
@@ -427,7 +429,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// ## Errors
     ///
     /// May return a database error.
-    fn get_state_caching(
+    pub fn get_state_caching(
         &self,
         state_root: &Hash256,
         slot: Option<Slot>,
@@ -448,7 +450,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// ## Errors
     ///
     /// May return a database error.
-    fn get_state_caching_only_with_committee_caches(
+    pub fn get_state_caching_only_with_committee_caches(
         &self,
         state_root: &Hash256,
         slot: Option<Slot>,
@@ -888,37 +890,40 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         let result = if let Some(attestation_head_block) =
             self.get_block_caching(&attestation.data.beacon_block_root)?
         {
-            // Use the `data.beacon_block_root` to load the state from the latest non-skipped
-            // slot preceding the attestation's creation.
-            //
-            // This state is guaranteed to be in the same chain as the attestation, but it's
-            // not guaranteed to be from the same slot or epoch as the attestation.
-            let mut state: BeaconState<T::EthSpec> = self
-                .get_state_caching_only_with_committee_caches(
-                    &attestation_head_block.state_root,
-                    Some(attestation_head_block.slot),
-                )?
-                .ok_or_else(|| Error::MissingBeaconState(attestation_head_block.state_root))?;
-
-            // Ensure the state loaded from the database matches the state of the attestation
-            // head block.
-            //
-            // The state needs to be advanced from the current slot through to the epoch in
-            // which the attestation was created in. It would be an error to try and use
-            // `state.get_attestation_data_slot(..)` because the state matching the
-            // `data.beacon_block_root` isn't necessarily in a nearby epoch to the attestation
-            // (e.g., if there were lots of skip slots since the head of the chain and the
-            // epoch creation epoch).
-            for _ in state.slot.as_u64()
-                ..attestation
-                    .data
-                    .target
-                    .epoch
-                    .start_slot(T::EthSpec::slots_per_epoch())
-                    .as_u64()
+            // If the attestation points to a block in the same epoch in which it was made,
+            // then it is sufficient to load the state from that epoch's boundary, because
+            // the epoch-variable fields like the justified checkpoints cannot have changed
+            // between the epoch boundary and when the attestation was made. If conversely,
+            // the attestation points to a block in a prior epoch, then it is necessary to
+            // load the full state corresponding to its block, and transition it to the
+            // attestation's epoch.
+            let attestation_epoch = attestation.data.target.epoch;
+            let slots_per_epoch = T::EthSpec::slots_per_epoch();
+            let mut state = if attestation_epoch
+                == attestation_head_block.slot.epoch(slots_per_epoch)
             {
-                per_slot_processing(&mut state, None, &self.spec)?;
-            }
+                self.store
+                    .load_epoch_boundary_state(&attestation_head_block.state_root)?
+                    .ok_or_else(|| Error::MissingBeaconState(attestation_head_block.state_root))?
+            } else {
+                let mut state = self
+                    .store
+                    .get_state(
+                        &attestation_head_block.state_root,
+                        Some(attestation_head_block.slot),
+                    )?
+                    .ok_or_else(|| Error::MissingBeaconState(attestation_head_block.state_root))?;
+
+                // Fastforward the state to the epoch in which the attestation was made.
+                // NOTE: this looks like a potential DoS vector, we should probably limit
+                // the amount we're willing to fastforward without a valid signature.
+                for _ in state.slot.as_u64()..attestation_epoch.start_slot(slots_per_epoch).as_u64()
+                {
+                    per_slot_processing(&mut state, None, &self.spec)?;
+                }
+
+                state
+            };
 
             state.build_committee_cache(RelativeEpoch::Current, &self.spec)?;
 
@@ -1242,14 +1247,16 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             });
         }
 
+        // Check if the block is already known. We know it is post-finalization, so it is
+        // sufficient to check the block root tree.
+        if self.block_root_tree.is_known_block_root(&block_root) {
+            return Ok(BlockProcessingOutcome::BlockIsAlreadyKnown);
+        }
+
         // Records the time taken to load the block and state from the database during block
         // processing.
         let db_read_timer = metrics::start_timer(&metrics::BLOCK_PROCESSING_DB_READ);
 
-        if self.store.exists::<BeaconBlock<T::EthSpec>>(&block_root)? {
-            return Ok(BlockProcessingOutcome::BlockIsAlreadyKnown);
-        }
-
         // Load the blocks parent block from the database, returning invalid if that block is not
         // found.
         let parent_block: BeaconBlock<T::EthSpec> =
@@ -1381,6 +1388,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         metrics::stop_timer(db_write_timer);
 
+        self.block_root_tree
+            .add_block_root(block_root, block.parent_root, block.slot)?;
+
         self.head_tracker.register_block(block_root, &block);
 
         let fork_choice_register_timer =
@@ -1692,8 +1702,10 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 .process_finalization(&finalized_block, finalized_block_root)?;
 
             let finalized_state = self
-                .store
-                .get_state(&finalized_block.state_root, Some(finalized_block.slot))?
+                .get_state_caching_only_with_committee_caches(
+                    &finalized_block.state_root,
+                    Some(finalized_block.slot),
+                )?
                 .ok_or_else(|| Error::MissingBeaconState(finalized_block.state_root))?;
 
             self.op_pool.prune_all(&finalized_state, &self.spec);
@@ -1706,6 +1718,12 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 max_finality_distance,
             );
 
+            // Prune in-memory block root tree.
+            self.block_root_tree.prune_to(
+                finalized_block_root,
+                self.heads().into_iter().map(|(block_root, _)| block_root),
+            );
+
             let _ = self.event_handler.register(EventKind::BeaconFinalization {
                 epoch: new_finalized_epoch,
                 root: finalized_block_root,

beacon_node/beacon_chain/src/builder.rs

@@ -16,7 +16,7 @@ use slot_clock::{SlotClock, TestingSlotClock};
 use std::marker::PhantomData;
 use std::sync::Arc;
 use std::time::Duration;
-use store::Store;
+use store::{BlockRootTree, Store};
 use types::{BeaconBlock, BeaconState, ChainSpec, EthSpec, Hash256, Slot};
 
 /// An empty struct used to "witness" all the `BeaconChainTypes` traits. It has no user-facing
@@ -92,6 +92,7 @@ pub struct BeaconChainBuilder<T: BeaconChainTypes> {
     slot_clock: Option<T::SlotClock>,
     persisted_beacon_chain: Option<PersistedBeaconChain<T>>,
     head_tracker: Option<HeadTracker>,
+    block_root_tree: Option<Arc<BlockRootTree>>,
     spec: ChainSpec,
     log: Option<Logger>,
 }
@@ -134,6 +135,7 @@ where
             slot_clock: None,
             persisted_beacon_chain: None,
             head_tracker: None,
+            block_root_tree: None,
             spec: TEthSpec::default_spec(),
             log: None,
         }
@@ -224,6 +226,7 @@ where
             HeadTracker::from_ssz_container(&p.ssz_head_tracker)
                 .map_err(|e| format!("Failed to decode head tracker for database: {:?}", e))?,
         );
+        self.block_root_tree = Some(Arc::new(p.block_root_tree.clone().into()));
         self.persisted_beacon_chain = Some(p);
 
         Ok(self)
@@ -266,6 +269,11 @@ where
             )
         })?;
 
+        self.block_root_tree = Some(Arc::new(BlockRootTree::new(
+            beacon_block_root,
+            beacon_block.slot,
+        )));
+
         self.finalized_checkpoint = Some(CheckPoint {
             beacon_block_root,
             beacon_block,
@@ -375,6 +383,9 @@ where
                 .event_handler
                 .ok_or_else(|| "Cannot build without an event handler".to_string())?,
             head_tracker: self.head_tracker.unwrap_or_default(),
+            block_root_tree: self
+                .block_root_tree
+                .ok_or_else(|| "Cannot build without a block root tree".to_string())?,
             checkpoint_cache: CheckPointCache::default(),
             log: log.clone(),
         };
@@ -425,10 +436,16 @@ where
             .clone()
             .ok_or_else(|| "reduced_tree_fork_choice requires a store")?;
 
+        let block_root_tree = self
+            .block_root_tree
+            .clone()
+            .ok_or_else(|| "reduced_tree_fork_choice requires a block root tree")?;
+
         let fork_choice = if let Some(persisted_beacon_chain) = &self.persisted_beacon_chain {
             ForkChoice::from_ssz_container(
                 persisted_beacon_chain.fork_choice.clone(),
                 store.clone(),
+                block_root_tree,
             )
             .map_err(|e| format!("Unable to decode fork choice from db: {:?}", e))?
         } else {
@@ -442,11 +459,12 @@ where
 
             let backend = ThreadSafeReducedTree::new(
                 store.clone(),
+                block_root_tree,
                 &finalized_checkpoint.beacon_block,
                 finalized_checkpoint.beacon_block_root,
             );
 
-            ForkChoice::new(store, backend, genesis_block_root, self.spec.genesis_slot)
+            ForkChoice::new(backend, genesis_block_root, self.spec.genesis_slot)
         };
 
         self.fork_choice = Some(fork_choice);

beacon_node/beacon_chain/src/errors.rs

@@ -5,6 +5,7 @@ use state_processing::per_block_processing::errors::AttestationValidationError;
 use state_processing::BlockProcessingError;
 use state_processing::SlotProcessingError;
 use std::time::Duration;
+use store::block_root_tree::BlockRootTreeError;
 use types::*;
 
 macro_rules! easy_from_to {
@@ -49,11 +50,13 @@ pub enum BeaconChainError {
     InvariantViolated(String),
     SszTypesError(SszTypesError),
     CanonicalHeadLockTimeout,
+    BlockRootTreeError(BlockRootTreeError),
 }
 
 easy_from_to!(SlotProcessingError, BeaconChainError);
 easy_from_to!(AttestationValidationError, BeaconChainError);
 easy_from_to!(SszTypesError, BeaconChainError);
+easy_from_to!(BlockRootTreeError, BeaconChainError);
 
 #[derive(Debug, PartialEq)]
 pub enum BlockProductionError {

beacon_node/beacon_chain/src/fork_choice.rs

@@ -4,7 +4,7 @@ use parking_lot::RwLock;
 use ssz_derive::{Decode, Encode};
 use state_processing::{common::get_attesting_indices, per_slot_processing};
 use std::sync::Arc;
-use store::{Error as StoreError, Store};
+use store::{BlockRootTree, Error as StoreError, Store};
 use types::{
     Attestation, BeaconBlock, BeaconState, BeaconStateError, Checkpoint, EthSpec, Hash256, Slot,
 };
@@ -22,7 +22,6 @@ pub enum Error {
 }
 
 pub struct ForkChoice<T: BeaconChainTypes> {
-    store: Arc<T::Store>,
     backend: T::LmdGhost,
     /// Used for resolving the `0x00..00` alias back to genesis.
     ///
@@ -36,7 +35,6 @@ pub struct ForkChoice<T: BeaconChainTypes> {
 }
 
 impl<T: BeaconChainTypes> PartialEq for ForkChoice<T> {
-    /// This implementation ignores the `store`.
     fn eq(&self, other: &Self) -> bool {
         self.backend == other.backend
             && self.genesis_block_root == other.genesis_block_root
@@ -50,18 +48,12 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
     ///
     /// "Genesis" does not necessarily need to be the absolute genesis, it can be some finalized
     /// block.
-    pub fn new(
-        store: Arc<T::Store>,
-        backend: T::LmdGhost,
-        genesis_block_root: Hash256,
-        genesis_slot: Slot,
-    ) -> Self {
+    pub fn new(backend: T::LmdGhost, genesis_block_root: Hash256, genesis_slot: Slot) -> Self {
         let justified_checkpoint = Checkpoint {
             epoch: genesis_slot.epoch(T::EthSpec::slots_per_epoch()),
             root: genesis_block_root,
         };
         Self {
-            store: store.clone(),
             backend,
             genesis_block_root,
             justified_checkpoint: RwLock::new(justified_checkpoint.clone()),
@@ -149,8 +141,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
             };
 
             let mut state: BeaconState<T::EthSpec> = chain
-                .store
-                .get_state(&block.state_root, Some(block.slot))?
+                .get_state_caching_only_with_committee_caches(&block.state_root, Some(block.slot))?
                 .ok_or_else(|| Error::MissingState(block.state_root))?;
 
             // Fast-forward the state to the start slot of the epoch where it was justified.
@@ -201,10 +192,7 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
         for attestation in &block.body.attestations {
             // If the `data.beacon_block_root` block is not known to us, simply ignore the latest
             // vote.
-            if let Some(block) = self
-                .store
-                .get::<BeaconBlock<T::EthSpec>>(&attestation.data.beacon_block_root)?
-            {
+            if let Some(block) = chain.get_block_caching(&attestation.data.beacon_block_root)? {
                 self.process_attestation(state, attestation, &block)?;
             }
         }
@@ -316,11 +304,14 @@ impl<T: BeaconChainTypes> ForkChoice<T> {
     /// Instantiates `Self` from a prior `SszForkChoice`.
     ///
     /// The created `Self` will have the same state as the `Self` that created the `SszForkChoice`.
-    pub fn from_ssz_container(ssz_container: SszForkChoice, store: Arc<T::Store>) -> Result<Self> {
-        let backend = LmdGhost::from_bytes(&ssz_container.backend_bytes, store.clone())?;
+    pub fn from_ssz_container(
+        ssz_container: SszForkChoice,
+        store: Arc<T::Store>,
+        block_root_tree: Arc<BlockRootTree>,
+    ) -> Result<Self> {
+        let backend = LmdGhost::from_bytes(&ssz_container.backend_bytes, store, block_root_tree)?;
 
         Ok(Self {
-            store,
             backend,
             genesis_block_root: ssz_container.genesis_block_root,
             justified_checkpoint: RwLock::new(ssz_container.justified_checkpoint),

beacon_node/beacon_chain/src/persisted_beacon_chain.rs

@@ -4,7 +4,7 @@ use crate::{BeaconChainTypes, CheckPoint};
 use operation_pool::PersistedOperationPool;
 use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
-use store::{DBColumn, Error as StoreError, SimpleStoreItem};
+use store::{DBColumn, Error as StoreError, SimpleStoreItem, SszBlockRootTree};
 use types::Hash256;
 
 /// 32-byte key for accessing the `PersistedBeaconChain`.
@@ -18,6 +18,7 @@ pub struct PersistedBeaconChain<T: BeaconChainTypes> {
     pub genesis_block_root: Hash256,
     pub ssz_head_tracker: SszHeadTracker,
     pub fork_choice: SszForkChoice,
+    pub block_root_tree: SszBlockRootTree,
 }
 
 impl<T: BeaconChainTypes> SimpleStoreItem for PersistedBeaconChain<T> {