Implement tree states & hierarchical state DB

beacon_node/store/src/impls/beacon_state.rs

@@ -1,62 +1,83 @@
 use crate::*;
-use ssz::{DecodeError, Encode};
+use ssz::Encode;
 use ssz_derive::Encode;
-use std::convert::TryInto;
-use types::beacon_state::{CloneConfig, CommitteeCache, CACHED_EPOCHS};
+use std::io::{Read, Write};
+use std::sync::Arc;
+use types::{CompactBeaconState, PublicKeyBytes};
+use zstd::{Decoder, Encoder};
 
 pub fn store_full_state<E: EthSpec>(
     state_root: &Hash256,
     state: &BeaconState<E>,
     ops: &mut Vec<KeyValueStoreOp>,
+    config: &StoreConfig,
 ) -> Result<(), Error> {
     let bytes = {
         let _overhead_timer = metrics::start_timer(&metrics::BEACON_STATE_WRITE_OVERHEAD_TIMES);
         StorageContainer::new(state).as_ssz_bytes()
     };
-    metrics::inc_counter_by(&metrics::BEACON_STATE_WRITE_BYTES, bytes.len() as u64);
+    let mut compressed_value = Vec::with_capacity(config.estimate_compressed_size(bytes.len()));
+    let mut encoder = Encoder::new(&mut compressed_value, config.compression_level)
+        .map_err(Error::Compression)?;
+    encoder.write_all(&bytes).map_err(Error::Compression)?;
+    encoder.finish().map_err(Error::Compression)?;
+
+    metrics::inc_counter_by(
+        &metrics::BEACON_STATE_WRITE_BYTES,
+        compressed_value.len() as u64,
+    );
     metrics::inc_counter(&metrics::BEACON_STATE_WRITE_COUNT);
+
     let key = get_key_for_col(DBColumn::BeaconState.into(), state_root.as_bytes());
-    ops.push(KeyValueStoreOp::PutKeyValue(key, bytes));
+    ops.push(KeyValueStoreOp::PutKeyValue(key, compressed_value));
     Ok(())
 }
 
-pub fn get_full_state<KV: KeyValueStore<E>, E: EthSpec>(
+pub fn get_full_state<KV: KeyValueStore<E>, E: EthSpec, F>(
     db: &KV,
     state_root: &Hash256,
+    immutable_validators: F,
+    config: &StoreConfig,
     spec: &ChainSpec,
-) -> Result<Option<BeaconState<E>>, Error> {
+) -> Result<Option<BeaconState<E>>, Error>
+where
+    F: Fn(usize) -> Option<Arc<PublicKeyBytes>>,
+{
     let total_timer = metrics::start_timer(&metrics::BEACON_STATE_READ_TIMES);
 
     match db.get_bytes(DBColumn::BeaconState.into(), state_root.as_bytes())? {
         Some(bytes) => {
+            let mut ssz_bytes = Vec::with_capacity(config.estimate_decompressed_size(bytes.len()));
+            let mut decoder = Decoder::new(&*bytes).map_err(Error::Compression)?;
+            decoder
+                .read_to_end(&mut ssz_bytes)
+                .map_err(Error::Compression)?;
+
             let overhead_timer = metrics::start_timer(&metrics::BEACON_STATE_READ_OVERHEAD_TIMES);
-            let container = StorageContainer::from_ssz_bytes(&bytes, spec)?;
+            let container = StorageContainer::from_ssz_bytes(&ssz_bytes, spec)?;
 
             metrics::stop_timer(overhead_timer);
             metrics::stop_timer(total_timer);
             metrics::inc_counter(&metrics::BEACON_STATE_READ_COUNT);
             metrics::inc_counter_by(&metrics::BEACON_STATE_READ_BYTES, bytes.len() as u64);
 
-            Ok(Some(container.try_into()?))
+            Ok(Some(container.into_beacon_state(immutable_validators)?))
         }
         None => Ok(None),
     }
 }
 
 /// A container for storing `BeaconState` components.
-// TODO: would be more space efficient with the caches stored separately and referenced by hash
 #[derive(Encode)]
 pub struct StorageContainer<T: EthSpec> {
-    state: BeaconState<T>,
-    committee_caches: Vec<CommitteeCache>,
+    state: CompactBeaconState<T>,
 }
 
 impl<T: EthSpec> StorageContainer<T> {
     /// Create a new instance for storing a `BeaconState`.
     pub fn new(state: &BeaconState<T>) -> Self {
         Self {
-            state: state.clone_with(CloneConfig::none()),
-            committee_caches: state.committee_caches().to_vec(),
+            state: state.clone().into_compact_state(),
         }
     }
 
@@ -66,36 +87,20 @@ impl<T: EthSpec> StorageContainer<T> {
         let mut builder = ssz::SszDecoderBuilder::new(bytes);
 
         builder.register_anonymous_variable_length_item()?;
-        builder.register_type::<Vec<CommitteeCache>>()?;
 
         let mut decoder = builder.build()?;
 
-        let state = decoder.decode_next_with(|bytes| BeaconState::from_ssz_bytes(bytes, spec))?;
-        let committee_caches = decoder.decode_next()?;
+        let state =
+            decoder.decode_next_with(|bytes| CompactBeaconState::from_ssz_bytes(bytes, spec))?;
 
-        Ok(Self {
-            state,
-            committee_caches,
-        })
+        Ok(Self { state })
     }
-}
-
-impl<T: EthSpec> TryInto<BeaconState<T>> for StorageContainer<T> {
-    type Error = Error;
-
-    fn try_into(mut self) -> Result<BeaconState<T>, Error> {
-        let mut state = self.state;
-
-        for i in (0..CACHED_EPOCHS).rev() {
-            if i >= self.committee_caches.len() {
-                return Err(Error::SszDecodeError(DecodeError::BytesInvalid(
-                    "Insufficient committees for BeaconState".to_string(),
-                )));
-            };
-
-            state.committee_caches_mut()[i] = self.committee_caches.remove(i);
-        }
 
+    fn into_beacon_state<F>(self, immutable_validators: F) -> Result<BeaconState<T>, Error>
+    where
+        F: Fn(usize) -> Option<Arc<PublicKeyBytes>>,
+    {
+        let state = self.state.try_into_full_state(immutable_validators)?;
         Ok(state)
     }
 }

beacon_node/store/src/impls/execution_payload.rs

@@ -9,8 +9,8 @@ macro_rules! impl_store_item {
                 DBColumn::ExecPayload
             }
 
-            fn as_store_bytes(&self) -> Vec<u8> {
-                self.as_ssz_bytes()
+            fn as_store_bytes(&self) -> Result<Vec<u8>, Error> {
+                Ok(self.as_ssz_bytes())
             }
 
             fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
@@ -31,8 +31,8 @@ impl<E: EthSpec> StoreItem for ExecutionPayload<E> {
         DBColumn::ExecPayload
     }
 
-    fn as_store_bytes(&self) -> Vec<u8> {
-        self.as_ssz_bytes()
+    fn as_store_bytes(&self) -> Result<Vec<u8>, Error> {
+        Ok(self.as_ssz_bytes())
     }
 
     fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {

beacon_node/store/src/impls/frozen_block_slot.rs

@@ -0,0 +1,19 @@
+use crate::{DBColumn, Error, StoreItem};
+use ssz::{Decode, Encode};
+use types::Slot;
+
+pub struct FrozenBlockSlot(pub Slot);
+
+impl StoreItem for FrozenBlockSlot {
+    fn db_column() -> DBColumn {
+        DBColumn::BeaconBlock
+    }
+
+    fn as_store_bytes(&self) -> Result<Vec<u8>, Error> {
+        Ok(self.0.as_ssz_bytes())
+    }
+
+    fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        Ok(FrozenBlockSlot(Slot::from_ssz_bytes(bytes)?))
+    }
+}