Add LRU cache to database (#837)

* Add LRU caches to store

* Improvements to LRU caches

* Take state by value in `Store::put_state`

* Store blocks by value, configurable cache sizes

* Use a StateBatch to efficiently store skip states

* Fix store tests

* Add CloneConfig test, remove unused metrics

* Use Mutexes instead of RwLocks for LRU caches

beacon_node/store/src/config.rs

@@ -1,36 +1,28 @@
 use serde_derive::{Deserialize, Serialize};
-use std::path::PathBuf;
 use types::{EthSpec, MinimalEthSpec};
 
-/// Default directory name for the freezer database under the top-level data dir.
-const DEFAULT_FREEZER_DB_DIR: &str = "freezer_db";
-
-/// Default value for the freezer DB's restore point frequency.
 pub const DEFAULT_SLOTS_PER_RESTORE_POINT: u64 = 2048;
+pub const DEFAULT_BLOCK_CACHE_SIZE: usize = 5;
+pub const DEFAULT_STATE_CACHE_SIZE: usize = 5;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+/// Database configuration parameters.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct StoreConfig {
-    /// Name of the directory inside the data directory where the main "hot" DB is located.
-    pub db_name: String,
-    /// Path where the freezer database will be located.
-    pub freezer_db_path: Option<PathBuf>,
     /// Number of slots to wait between storing restore points in the freezer database.
     pub slots_per_restore_point: u64,
+    /// Maximum number of blocks to store in the in-memory block cache.
+    pub block_cache_size: usize,
+    /// Maximum number of states to store in the in-memory state cache.
+    pub state_cache_size: usize,
 }
 
 impl Default for StoreConfig {
     fn default() -> Self {
         Self {
-            db_name: "chain_db".to_string(),
-            freezer_db_path: None,
             // Safe default for tests, shouldn't ever be read by a CLI node.
             slots_per_restore_point: MinimalEthSpec::slots_per_historical_root() as u64,
+            block_cache_size: DEFAULT_BLOCK_CACHE_SIZE,
+            state_cache_size: DEFAULT_STATE_CACHE_SIZE,
         }
     }
 }
-
-impl StoreConfig {
-    pub fn default_freezer_db_dir(&self) -> &'static str {
-        DEFAULT_FREEZER_DB_DIR
-    }
-}

beacon_node/store/src/hot_cold_store.rs

@@ -1,12 +1,16 @@
 use crate::chunked_vector::{
     store_updated_vector, BlockRoots, HistoricalRoots, RandaoMixes, StateRoots,
 };
+use crate::config::StoreConfig;
 use crate::forwards_iter::HybridForwardsBlockRootsIterator;
+use crate::impls::beacon_state::store_full_state;
 use crate::iter::{ParentRootBlockIterator, StateRootsIterator};
+use crate::metrics;
 use crate::{
     leveldb_store::LevelDB, DBColumn, Error, PartialBeaconState, SimpleStoreItem, Store, StoreItem,
 };
-use parking_lot::RwLock;
+use lru::LruCache;
+use parking_lot::{Mutex, RwLock};
 use slog::{debug, trace, warn, Logger};
 use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
@@ -18,6 +22,7 @@ use std::convert::TryInto;
 use std::marker::PhantomData;
 use std::path::Path;
 use std::sync::Arc;
+use types::beacon_state::CloneConfig;
 use types::*;
 
 /// 32-byte key for accessing the `split` of the freezer DB.
@@ -33,14 +38,17 @@ pub struct HotColdDB<E: EthSpec> {
     /// States with slots less than `split.slot` are in the cold DB, while states with slots
     /// greater than or equal are in the hot DB.
     split: RwLock<Split>,
-    /// Number of slots per restore point state in the freezer database.
-    slots_per_restore_point: u64,
+    config: StoreConfig,
     /// Cold database containing compact historical data.
     pub(crate) cold_db: LevelDB<E>,
     /// Hot database containing duplicated but quick-to-access recent data.
     ///
     /// The hot database also contains all blocks.
     pub(crate) hot_db: LevelDB<E>,
+    /// LRU cache of deserialized blocks. Updated whenever a block is loaded.
+    block_cache: Mutex<LruCache<Hash256, BeaconBlock<E>>>,
+    /// LRU cache of deserialized states. Updated whenever a state is loaded.
+    state_cache: Mutex<LruCache<Hash256, BeaconState<E>>>,
     /// Chain spec.
     spec: ChainSpec,
     /// Logger.
@@ -98,10 +106,42 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
         self.hot_db.key_delete(column, key)
     }
 
+    /// Store a block and update the LRU cache.
+    fn put_block(&self, block_root: &Hash256, block: BeaconBlock<E>) -> Result<(), Error> {
+        // Store on disk.
+        self.put(block_root, &block)?;
+
+        // Update cache.
+        self.block_cache.lock().put(*block_root, block);
+
+        Ok(())
+    }
+
+    /// Fetch a block from the store.
+    fn get_block(&self, block_root: &Hash256) -> Result<Option<BeaconBlock<E>>, Error> {
+        metrics::inc_counter(&metrics::BEACON_BLOCK_GET_COUNT);
+
+        // Check the cache.
+        if let Some(block) = self.block_cache.lock().get(block_root) {
+            metrics::inc_counter(&metrics::BEACON_BLOCK_CACHE_HIT_COUNT);
+            return Ok(Some(block.clone()));
+        }
+
+        // Fetch from database.
+        match self.get::<BeaconBlock<E>>(block_root)? {
+            Some(block) => {
+                // Add to cache.
+                self.block_cache.lock().put(*block_root, block.clone());
+                Ok(Some(block))
+            }
+            None => Ok(None),
+        }
+    }
+
     /// Store a state in the store.
-    fn put_state(&self, state_root: &Hash256, state: &BeaconState<E>) -> Result<(), Error> {
+    fn put_state(&self, state_root: &Hash256, state: BeaconState<E>) -> Result<(), Error> {
         if state.slot < self.get_split_slot() {
-            self.store_cold_state(state_root, state)
+            self.store_cold_state(state_root, &state)
         } else {
             self.store_hot_state(state_root, state)
         }
@@ -113,14 +153,28 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
         state_root: &Hash256,
         slot: Option<Slot>,
     ) -> Result<Option<BeaconState<E>>, Error> {
+        self.get_state_with(state_root, slot, CloneConfig::all())
+    }
+
+    /// Get a state from the store.
+    ///
+    /// Fetch a state from the store, controlling which cache fields are cloned.
+    fn get_state_with(
+        &self,
+        state_root: &Hash256,
+        slot: Option<Slot>,
+        clone_config: CloneConfig,
+    ) -> Result<Option<BeaconState<E>>, Error> {
+        metrics::inc_counter(&metrics::BEACON_STATE_GET_COUNT);
+
         if let Some(slot) = slot {
             if slot < self.get_split_slot() {
                 self.load_cold_state_by_slot(slot).map(Some)
             } else {
-                self.load_hot_state(state_root)
+                self.load_hot_state(state_root, clone_config)
             }
         } else {
-            match self.load_hot_state(state_root)? {
+            match self.load_hot_state(state_root, clone_config)? {
                 Some(state) => Ok(Some(state)),
                 None => self.load_cold_state(state_root),
             }
@@ -164,7 +218,7 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
         for (state_root, slot) in
             state_root_iter.take_while(|&(_, slot)| slot >= current_split_slot)
         {
-            if slot % store.slots_per_restore_point == 0 {
+            if slot % store.config.slots_per_restore_point == 0 {
                 let state: BeaconState<E> = store
                     .hot_db
                     .get_state(&state_root, None)?
@@ -229,9 +283,12 @@ impl<E: EthSpec> Store<E> for HotColdDB<E> {
             ..
         }) = self.load_hot_state_summary(state_root)?
         {
+            // NOTE: minor inefficiency here because we load an unnecessary hot state summary
             let state = self
-                .hot_db
-                .get_state(&epoch_boundary_state_root, None)?
+                .load_hot_state(
+                    &epoch_boundary_state_root,
+                    CloneConfig::committee_caches_only(),
+                )?
                 .ok_or_else(|| {
                     HotColdDBError::MissingEpochBoundaryState(epoch_boundary_state_root)
                 })?;
@@ -257,17 +314,19 @@ impl<E: EthSpec> HotColdDB<E> {
     pub fn open(
         hot_path: &Path,
         cold_path: &Path,
-        slots_per_restore_point: u64,
+        config: StoreConfig,
         spec: ChainSpec,
         log: Logger,
     ) -> Result<Self, Error> {
-        Self::verify_slots_per_restore_point(slots_per_restore_point)?;
+        Self::verify_slots_per_restore_point(config.slots_per_restore_point)?;
 
         let db = HotColdDB {
             split: RwLock::new(Split::default()),
-            slots_per_restore_point,
             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(config.state_cache_size)),
+            config,
             spec,
             log,
             _phantom: PhantomData,
@@ -288,7 +347,7 @@ impl<E: EthSpec> HotColdDB<E> {
     pub fn store_hot_state(
         &self,
         state_root: &Hash256,
-        state: &BeaconState<E>,
+        state: BeaconState<E>,
     ) -> Result<(), Error> {
         // On the epoch boundary, store the full state.
         if state.slot % E::slots_per_epoch() == 0 {
@@ -298,13 +357,16 @@ impl<E: EthSpec> HotColdDB<E> {
                 "slot" => state.slot.as_u64(),
                 "state_root" => format!("{:?}", state_root)
             );
-            self.hot_db.put_state(state_root, state)?;
+            store_full_state(&self.hot_db, 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)?;
+        self.put_state_summary(state_root, HotStateSummary::new(state_root, &state)?)?;
+
+        // Store the state in the cache.
+        self.state_cache.lock().put(*state_root, state);
 
         Ok(())
     }
@@ -312,14 +374,31 @@ impl<E: EthSpec> HotColdDB<E> {
     /// 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> {
+    pub fn load_hot_state(
+        &self,
+        state_root: &Hash256,
+        clone_config: CloneConfig,
+    ) -> Result<Option<BeaconState<E>>, Error> {
+        metrics::inc_counter(&metrics::BEACON_STATE_HOT_GET_COUNT);
+
+        // Check the cache.
+        if let Some(state) = self.state_cache.lock().get(state_root) {
+            metrics::inc_counter(&metrics::BEACON_STATE_CACHE_HIT_COUNT);
+
+            let timer = metrics::start_timer(&metrics::BEACON_STATE_CACHE_CLONE_TIME);
+            let state = state.clone_with(clone_config);
+            metrics::stop_timer(timer);
+
+            return Ok(Some(state));
+        }
+
         if let Some(HotStateSummary {
             slot,
             latest_block_root,
             epoch_boundary_state_root,
         }) = self.load_hot_state_summary(state_root)?
         {
-            let state: BeaconState<E> = self
+            let boundary_state = self
                 .hot_db
                 .get_state(&epoch_boundary_state_root, None)?
                 .ok_or_else(|| {
@@ -328,12 +407,18 @@ impl<E: EthSpec> HotColdDB<E> {
 
             // 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))
+            let state = if slot % E::slots_per_epoch() == 0 {
+                boundary_state
             } else {
-                let blocks = self.load_blocks_to_replay(state.slot, slot, latest_block_root)?;
-                self.replay_blocks(state, blocks, slot).map(Some)
-            }
+                let blocks =
+                    self.load_blocks_to_replay(boundary_state.slot, slot, latest_block_root)?;
+                self.replay_blocks(boundary_state, blocks, slot)?
+            };
+
+            // Update the LRU cache.
+            self.state_cache.lock().put(*state_root, state.clone());
+
+            Ok(Some(state))
         } else {
             Ok(None)
         }
@@ -348,7 +433,7 @@ impl<E: EthSpec> HotColdDB<E> {
         state_root: &Hash256,
         state: &BeaconState<E>,
     ) -> Result<(), Error> {
-        if state.slot % self.slots_per_restore_point != 0 {
+        if state.slot % self.config.slots_per_restore_point != 0 {
             warn!(
                 self.log,
                 "Not storing non-restore point state in freezer";
@@ -377,7 +462,7 @@ impl<E: EthSpec> HotColdDB<E> {
         store_updated_vector(RandaoMixes, db, state, &self.spec)?;
 
         // 3. Store restore point.
-        let restore_point_index = state.slot.as_u64() / self.slots_per_restore_point;
+        let restore_point_index = state.slot.as_u64() / self.config.slots_per_restore_point;
         self.store_restore_point_hash(restore_point_index, *state_root)?;
 
         Ok(())
@@ -397,8 +482,8 @@ impl<E: EthSpec> HotColdDB<E> {
     ///
     /// Will reconstruct the state if it lies between restore points.
     pub fn load_cold_state_by_slot(&self, slot: Slot) -> Result<BeaconState<E>, Error> {
-        if slot % self.slots_per_restore_point == 0 {
-            let restore_point_idx = slot.as_u64() / self.slots_per_restore_point;
+        if slot % self.config.slots_per_restore_point == 0 {
+            let restore_point_idx = slot.as_u64() / self.config.slots_per_restore_point;
             self.load_restore_point_by_index(restore_point_idx)
         } else {
             self.load_cold_intermediate_state(slot)
@@ -431,7 +516,7 @@ impl<E: EthSpec> HotColdDB<E> {
     /// 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 low_restore_point_idx = slot.as_u64() / self.config.slots_per_restore_point;
         let high_restore_point_idx = low_restore_point_idx + 1;
 
         // Acquire the read lock, so that the split can't change while this is happening.
@@ -440,7 +525,7 @@ impl<E: EthSpec> HotColdDB<E> {
         let low_restore_point = self.load_restore_point_by_index(low_restore_point_idx)?;
         // If the slot of the high point lies outside the freezer, use the split state
         // as the upper restore point.
-        let high_restore_point = if high_restore_point_idx * self.slots_per_restore_point
+        let high_restore_point = if high_restore_point_idx * self.config.slots_per_restore_point
             >= split.slot.as_u64()
         {
             self.get_state(&split.state_root, Some(split.slot))?
@@ -553,7 +638,8 @@ impl<E: EthSpec> HotColdDB<E> {
 
     /// Fetch the slot of the most recently stored restore point.
     pub fn get_latest_restore_point_slot(&self) -> Slot {
-        (self.get_split_slot() - 1) / self.slots_per_restore_point * self.slots_per_restore_point
+        (self.get_split_slot() - 1) / self.config.slots_per_restore_point
+            * self.config.slots_per_restore_point
     }
 
     /// Load the split point from disk.
@@ -615,33 +701,6 @@ impl<E: EthSpec> HotColdDB<E> {
         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 valid.
     ///
     /// Specifically, check that it is:
@@ -718,6 +777,29 @@ impl SimpleStoreItem for HotStateSummary {
     }
 }
 
+impl HotStateSummary {
+    /// Construct a new summary of the given state.
+    pub fn new<E: EthSpec>(state_root: &Hash256, state: &BeaconState<E>) -> Result<Self, 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)?
+        };
+
+        Ok(HotStateSummary {
+            slot: state.slot,
+            latest_block_root,
+            epoch_boundary_state_root,
+        })
+    }
+}
+
 /// Struct for summarising a state in the freezer database.
 #[derive(Debug, Clone, Copy, Default, Encode, Decode)]
 struct ColdStateSummary {

beacon_node/store/src/impls/beacon_state.rs

@@ -2,7 +2,7 @@ use crate::*;
 use ssz::{Decode, DecodeError, Encode};
 use ssz_derive::{Decode, Encode};
 use std::convert::TryInto;
-use types::beacon_state::{CommitteeCache, CACHED_EPOCHS};
+use types::beacon_state::{CloneConfig, CommitteeCache, CACHED_EPOCHS};
 
 pub fn store_full_state<S: Store<E>, E: EthSpec>(
     store: &S,
@@ -58,7 +58,7 @@ impl<T: EthSpec> StorageContainer<T> {
     /// Create a new instance for storing a `BeaconState`.
     pub fn new(state: &BeaconState<T>) -> Self {
         Self {
-            state: state.clone_without_caches(),
+            state: state.clone_with(CloneConfig::none()),
             committee_caches: state.committee_caches.to_vec(),
         }
     }

beacon_node/store/src/iter.rs

@@ -343,7 +343,7 @@ mod test {
 
         let state_a_root = hashes.next().unwrap();
         state_b.state_roots[0] = state_a_root;
-        store.put_state(&state_a_root, &state_a).unwrap();
+        store.put_state(&state_a_root, state_a).unwrap();
 
         let iter = BlockRootsIterator::new(store, &state_b);
 
@@ -391,8 +391,8 @@ mod test {
         let state_a_root = Hash256::from_low_u64_be(slots_per_historical_root as u64);
         let state_b_root = Hash256::from_low_u64_be(slots_per_historical_root as u64 * 2);
 
-        store.put_state(&state_a_root, &state_a).unwrap();
-        store.put_state(&state_b_root, &state_b).unwrap();
+        store.put_state(&state_a_root, state_a).unwrap();
+        store.put_state(&state_b_root, state_b.clone()).unwrap();
 
         let iter = StateRootsIterator::new(store, &state_b);
 

beacon_node/store/src/leveldb_store.rs

@@ -123,8 +123,8 @@ impl<E: EthSpec> Store<E> for LevelDB<E> {
     }
 
     /// Store a state in the store.
-    fn put_state(&self, state_root: &Hash256, state: &BeaconState<E>) -> Result<(), Error> {
-        store_full_state(self, state_root, state)
+    fn put_state(&self, state_root: &Hash256, state: BeaconState<E>) -> Result<(), Error> {
+        store_full_state(self, state_root, &state)
     }
 
     /// Fetch a state from the store.

beacon_node/store/src/lib.rs

@@ -22,6 +22,7 @@ mod leveldb_store;
 mod memory_store;
 mod metrics;
 mod partial_beacon_state;
+mod state_batch;
 
 pub mod iter;
 pub mod migrate;
@@ -29,7 +30,7 @@ pub mod migrate;
 use std::sync::Arc;
 
 pub use self::config::StoreConfig;
-pub use self::hot_cold_store::HotColdDB as DiskStore;
+pub use self::hot_cold_store::{HotColdDB as DiskStore, HotStateSummary};
 pub use self::leveldb_store::LevelDB as SimpleDiskStore;
 pub use self::memory_store::MemoryStore;
 pub use self::migrate::Migrate;
@@ -37,6 +38,8 @@ pub use self::partial_beacon_state::PartialBeaconState;
 pub use errors::Error;
 pub use impls::beacon_state::StorageContainer as BeaconStateStorageContainer;
 pub use metrics::scrape_for_metrics;
+pub use state_batch::StateBatch;
+pub use types::beacon_state::CloneConfig;
 pub use types::*;
 
 /// An object capable of storing and retrieving objects implementing `StoreItem`.
@@ -79,8 +82,29 @@ pub trait Store<E: EthSpec>: Sync + Send + Sized + 'static {
         I::db_delete(self, key)
     }
 
+    /// Store a block in the store.
+    fn put_block(&self, block_root: &Hash256, block: BeaconBlock<E>) -> Result<(), Error> {
+        self.put(block_root, &block)
+    }
+
+    /// Fetch a block from the store.
+    fn get_block(&self, block_root: &Hash256) -> Result<Option<BeaconBlock<E>>, Error> {
+        self.get(block_root)
+    }
+
     /// Store a state in the store.
-    fn put_state(&self, state_root: &Hash256, state: &BeaconState<E>) -> Result<(), Error>;
+    fn put_state(&self, state_root: &Hash256, state: BeaconState<E>) -> Result<(), Error>;
+
+    /// Store a state summary in the store.
+    // NOTE: this is a hack for the HotColdDb, we could consider splitting this
+    // trait and removing the generic `S: Store` types everywhere?
+    fn put_state_summary(
+        &self,
+        state_root: &Hash256,
+        summary: HotStateSummary,
+    ) -> Result<(), Error> {
+        summary.db_put(self, state_root).map_err(Into::into)
+    }
 
     /// Fetch a state from the store.
     fn get_state(
@@ -89,6 +113,17 @@ pub trait Store<E: EthSpec>: Sync + Send + Sized + 'static {
         slot: Option<Slot>,
     ) -> Result<Option<BeaconState<E>>, Error>;
 
+    /// Fetch a state from the store, controlling which cache fields are cloned.
+    fn get_state_with(
+        &self,
+        state_root: &Hash256,
+        slot: Option<Slot>,
+        _clone_config: CloneConfig,
+    ) -> Result<Option<BeaconState<E>>, Error> {
+        // Default impl ignores config. Overriden in `HotColdDb`.
+        self.get_state(state_root, slot)
+    }
+
     /// Given the root of an existing block in the store (`start_block_root`), return a parent
     /// block with the specified `slot`.
     ///
@@ -315,13 +350,12 @@ mod tests {
 
         let hot_dir = tempdir().unwrap();
         let cold_dir = tempdir().unwrap();
-        let slots_per_restore_point = MinimalEthSpec::slots_per_historical_root() as u64;
         let spec = MinimalEthSpec::default_spec();
         let log = NullLoggerBuilder.build().unwrap();
         let store = DiskStore::open(
             &hot_dir.path(),
             &cold_dir.path(),
-            slots_per_restore_point,
+            StoreConfig::default(),
             spec,
             log,
         )

beacon_node/store/src/memory_store.rs

@@ -76,8 +76,8 @@ impl<E: EthSpec> Store<E> for MemoryStore<E> {
     }
 
     /// Store a state in the store.
-    fn put_state(&self, state_root: &Hash256, state: &BeaconState<E>) -> Result<(), Error> {
-        store_full_state(self, state_root, state)
+    fn put_state(&self, state_root: &Hash256, state: BeaconState<E>) -> Result<(), Error> {
+        store_full_state(self, state_root, &state)
     }
 
     /// Fetch a state from the store.

beacon_node/store/src/metrics.rs

@@ -46,6 +46,22 @@ lazy_static! {
     /*
      * Beacon State
      */
+    pub static ref BEACON_STATE_GET_COUNT: Result<IntCounter> = try_create_int_counter(
+        "store_beacon_state_get_total",
+        "Total number of beacon states requested from the store (cache or DB)"
+    );
+    pub static ref BEACON_STATE_HOT_GET_COUNT: Result<IntCounter> = try_create_int_counter(
+        "store_beacon_state_hot_get_total",
+        "Total number of hot beacon states requested from the store (cache or DB)"
+    );
+    pub static ref BEACON_STATE_CACHE_HIT_COUNT: Result<IntCounter> = try_create_int_counter(
+        "store_beacon_state_cache_hit_total",
+        "Number of hits to the store's state cache"
+    );
+    pub static ref BEACON_STATE_CACHE_CLONE_TIME: Result<Histogram> = try_create_histogram(
+        "store_beacon_state_cache_clone_time",
+        "Time to load a beacon block from the block cache"
+    );
     pub static ref BEACON_STATE_READ_TIMES: Result<Histogram> = try_create_histogram(
         "store_beacon_state_read_seconds",
         "Total time required to read a BeaconState from the database"
@@ -81,6 +97,14 @@ lazy_static! {
     /*
      * Beacon Block
      */
+    pub static ref BEACON_BLOCK_GET_COUNT: Result<IntCounter> = try_create_int_counter(
+        "store_beacon_block_get_total",
+        "Total number of beacon blocks requested from the store (cache or DB)"
+    );
+    pub static ref BEACON_BLOCK_CACHE_HIT_COUNT: Result<IntCounter> = try_create_int_counter(
+        "store_beacon_block_cache_hit_total",
+        "Number of hits to the store's block cache"
+    );
     pub static ref BEACON_BLOCK_READ_TIMES: Result<Histogram> = try_create_histogram(
         "store_beacon_block_read_overhead_seconds",
         "Overhead on reading a beacon block from the DB (e.g., decoding)"

beacon_node/store/src/state_batch.rs

@@ -0,0 +1,47 @@
+use crate::{Error, HotStateSummary, Store};
+use types::{BeaconState, EthSpec, Hash256};
+
+/// A collection of states to be stored in the database.
+///
+/// Consumes minimal space in memory by not storing states between epoch boundaries.
+#[derive(Debug, Clone, Default)]
+pub struct StateBatch<E: EthSpec> {
+    items: Vec<BatchItem<E>>,
+}
+
+#[derive(Debug, Clone)]
+#[allow(clippy::large_enum_variant)]
+enum BatchItem<E: EthSpec> {
+    Full(Hash256, BeaconState<E>),
+    Summary(Hash256, HotStateSummary),
+}
+
+impl<E: EthSpec> StateBatch<E> {
+    /// Create a new empty batch.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Stage a `BeaconState` to be stored.
+    pub fn add_state(&mut self, state_root: Hash256, state: &BeaconState<E>) -> Result<(), Error> {
+        let item = if state.slot % E::slots_per_epoch() == 0 {
+            BatchItem::Full(state_root, state.clone())
+        } else {
+            BatchItem::Summary(state_root, HotStateSummary::new(&state_root, state)?)
+        };
+        self.items.push(item);
+        Ok(())
+    }
+
+    /// Write the batch to the database.
+    ///
+    /// May fail to write the full batch if any of the items error (i.e. not atomic!)
+    pub fn commit<S: Store<E>>(self, store: &S) -> Result<(), Error> {
+        self.items.into_iter().try_for_each(|item| match item {
+            BatchItem::Full(state_root, state) => store.put_state(&state_root, state),
+            BatchItem::Summary(state_root, summary) => {
+                store.put_state_summary(&state_root, summary)
+            }
+        })
+    }
+}