Faster BeaconState enc/dec (#671)

* Add state enc/dec benches

* Add example for flamegraph

* Use `PublicKeyBytes` for `Validator`

* Ripple PublicKeyBytes change through codebase

* Add benches, optimizations to store BeaconState

* Store BeaconState in StorageContainer too

* Optimize StorageContainer with std::mem magic

* Fix rest_api tests

beacon_node/store/Cargo.toml

@@ -4,9 +4,15 @@ version = "0.1.0"
 authors = ["Paul Hauner <paul@paulhauner.com>"]
 edition = "2018"
 
+[[bench]]
+name = "benches"
+harness = false
+
 [dev-dependencies]
 tempfile = "3.1.0"
 sloggers = "0.3.2"
+criterion = "0.3.0"
+rayon = "1.2.0"
 
 [dependencies]
 db-key = "0.0.5"

beacon_node/store/benches/benches.rs

@@ -0,0 +1,114 @@
+use criterion::Criterion;
+use criterion::{black_box, criterion_group, criterion_main, Benchmark};
+use rayon::prelude::*;
+use ssz::{Decode, Encode};
+use std::convert::TryInto;
+use store::BeaconStateStorageContainer;
+use types::{
+    test_utils::generate_deterministic_keypair, BeaconState, Epoch, Eth1Data, EthSpec, Hash256,
+    MainnetEthSpec, Validator,
+};
+
+fn get_state<E: EthSpec>(validator_count: usize) -> BeaconState<E> {
+    let spec = &E::default_spec();
+    let eth1_data = Eth1Data {
+        deposit_root: Hash256::zero(),
+        deposit_count: 0,
+        block_hash: Hash256::zero(),
+    };
+
+    let mut state = BeaconState::new(0, eth1_data, spec);
+
+    for i in 0..validator_count {
+        state.balances.push(i as u64).expect("should add balance");
+    }
+
+    state.validators = (0..validator_count)
+        .into_iter()
+        .collect::<Vec<_>>()
+        .par_iter()
+        .map(|&i| Validator {
+            pubkey: generate_deterministic_keypair(i).pk.into(),
+            withdrawal_credentials: Hash256::from_low_u64_le(i as u64),
+            effective_balance: spec.max_effective_balance,
+            slashed: false,
+            activation_eligibility_epoch: Epoch::new(0),
+            activation_epoch: Epoch::new(0),
+            exit_epoch: Epoch::from(u64::max_value()),
+            withdrawable_epoch: Epoch::from(u64::max_value()),
+        })
+        .collect::<Vec<_>>()
+        .into();
+
+    state.build_all_caches(spec).expect("should build caches");
+
+    state
+}
+
+fn all_benches(c: &mut Criterion) {
+    let validator_count = 16_384;
+    let state = get_state::<MainnetEthSpec>(validator_count);
+    let storage_container = BeaconStateStorageContainer::new(&state);
+    let state_bytes = storage_container.as_ssz_bytes();
+
+    let inner_state = state.clone();
+    c.bench(
+        &format!("{}_validators", validator_count),
+        Benchmark::new("encode/beacon_state", move |b| {
+            b.iter_batched_ref(
+                || inner_state.clone(),
+                |state| black_box(BeaconStateStorageContainer::new(state).as_ssz_bytes()),
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+
+    let inner_state = state.clone();
+    c.bench(
+        &format!("{}_validators", validator_count),
+        Benchmark::new("encode/beacon_state/tree_hash_cache", move |b| {
+            b.iter_batched_ref(
+                || inner_state.tree_hash_cache.clone(),
+                |tree_hash_cache| black_box(tree_hash_cache.as_ssz_bytes()),
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+
+    let inner_state = state.clone();
+    c.bench(
+        &format!("{}_validators", validator_count),
+        Benchmark::new("encode/beacon_state/committee_cache[0]", move |b| {
+            b.iter_batched_ref(
+                || inner_state.committee_caches[0].clone(),
+                |committee_cache| black_box(committee_cache.as_ssz_bytes()),
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+
+    c.bench(
+        &format!("{}_validators", validator_count),
+        Benchmark::new("decode/beacon_state", move |b| {
+            b.iter_batched_ref(
+                || state_bytes.clone(),
+                |bytes| {
+                    let state: BeaconState<MainnetEthSpec> =
+                        BeaconStateStorageContainer::from_ssz_bytes(&bytes)
+                            .expect("should decode")
+                            .try_into()
+                            .expect("should convert into state");
+                    black_box(state)
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        })
+        .sample_size(10),
+    );
+}
+
+criterion_group!(benches, all_benches,);
+criterion_main!(benches);

beacon_node/store/examples/ssz_encode_state.rs

@@ -0,0 +1,63 @@
+//! These examples only really exist so we can use them for flamegraph. If they get annoying to
+//! maintain, feel free to delete.
+
+use rayon::prelude::*;
+use ssz::{Decode, Encode};
+use std::convert::TryInto;
+use store::BeaconStateStorageContainer;
+use types::{
+    test_utils::generate_deterministic_keypair, BeaconState, Epoch, Eth1Data, EthSpec, Hash256,
+    MainnetEthSpec, Validator,
+};
+
+type E = MainnetEthSpec;
+
+fn get_state<E: EthSpec>(validator_count: usize) -> BeaconState<E> {
+    let spec = &E::default_spec();
+    let eth1_data = Eth1Data {
+        deposit_root: Hash256::zero(),
+        deposit_count: 0,
+        block_hash: Hash256::zero(),
+    };
+
+    let mut state = BeaconState::new(0, eth1_data, spec);
+
+    for i in 0..validator_count {
+        state.balances.push(i as u64).expect("should add balance");
+    }
+
+    state.validators = (0..validator_count)
+        .into_iter()
+        .collect::<Vec<_>>()
+        .par_iter()
+        .map(|&i| Validator {
+            pubkey: generate_deterministic_keypair(i).pk.into(),
+            withdrawal_credentials: Hash256::from_low_u64_le(i as u64),
+            effective_balance: spec.max_effective_balance,
+            slashed: false,
+            activation_eligibility_epoch: Epoch::new(0),
+            activation_epoch: Epoch::new(0),
+            exit_epoch: Epoch::from(u64::max_value()),
+            withdrawable_epoch: Epoch::from(u64::max_value()),
+        })
+        .collect::<Vec<_>>()
+        .into();
+
+    state.build_all_caches(spec).expect("should build caches");
+
+    state
+}
+
+fn main() {
+    let validator_count = 1_024;
+    let state = get_state::<E>(validator_count);
+    let storage_container = BeaconStateStorageContainer::new(&state);
+
+    for _ in 0..1024 {
+        let container_bytes = storage_container.as_ssz_bytes();
+        let _: BeaconState<E> = BeaconStateStorageContainer::from_ssz_bytes(&container_bytes)
+            .expect("should decode")
+            .try_into()
+            .expect("should convert into state");
+    }
+}

beacon_node/store/src/impls/beacon_state.rs

@@ -44,48 +44,51 @@ pub fn get_full_state<S: Store, E: EthSpec>(
 /// A container for storing `BeaconState` components.
 // TODO: would be more space efficient with the caches stored separately and referenced by hash
 #[derive(Encode, Decode)]
-struct StorageContainer {
-    state_bytes: Vec<u8>,
-    committee_caches_bytes: Vec<Vec<u8>>,
-    tree_hash_cache_bytes: Vec<u8>,
+pub struct StorageContainer<T: EthSpec> {
+    state: BeaconState<T>,
+    committee_caches: Vec<CommitteeCache>,
+    tree_hash_cache: BeaconTreeHashCache,
 }
 
-impl StorageContainer {
+impl<T: EthSpec> StorageContainer<T> {
     /// Create a new instance for storing a `BeaconState`.
-    pub fn new<T: EthSpec>(state: &BeaconState<T>) -> Self {
-        let mut committee_caches_bytes = vec![];
+    pub fn new(state: &BeaconState<T>) -> Self {
+        let mut state = state.clone();
 
-        for cache in state.committee_caches[..].iter() {
-            committee_caches_bytes.push(cache.as_ssz_bytes());
+        let mut committee_caches = vec![CommitteeCache::default(); CACHED_EPOCHS];
+
+        for i in 0..CACHED_EPOCHS {
+            std::mem::swap(&mut state.committee_caches[i], &mut committee_caches[i]);
         }
 
-        let tree_hash_cache_bytes = state.tree_hash_cache.as_ssz_bytes();
+        let tree_hash_cache =
+            std::mem::replace(&mut state.tree_hash_cache, BeaconTreeHashCache::default());
 
         Self {
-            state_bytes: state.as_ssz_bytes(),
-            committee_caches_bytes,
-            tree_hash_cache_bytes,
+            state,
+            committee_caches,
+            tree_hash_cache,
         }
     }
 }
 
-impl<T: EthSpec> TryInto<BeaconState<T>> for StorageContainer {
+impl<T: EthSpec> TryInto<BeaconState<T>> for StorageContainer<T> {
     type Error = Error;
 
-    fn try_into(self) -> Result<BeaconState<T>, Error> {
-        let mut state: BeaconState<T> = BeaconState::from_ssz_bytes(&self.state_bytes)?;
+    fn try_into(mut self) -> Result<BeaconState<T>, Error> {
+        let mut state = self.state;
 
-        for i in 0..CACHED_EPOCHS {
-            let bytes = &self.committee_caches_bytes.get(i).ok_or_else(|| {
-                Error::SszDecodeError(DecodeError::BytesInvalid(
+        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[i] = CommitteeCache::from_ssz_bytes(bytes)?;
+            state.committee_caches[i] = self.committee_caches.remove(i);
         }
 
-        state.tree_hash_cache = BeaconTreeHashCache::from_ssz_bytes(&self.tree_hash_cache_bytes)?;
+        state.tree_hash_cache = self.tree_hash_cache;
 
         Ok(state)
     }

beacon_node/store/src/lib.rs

@@ -33,6 +33,7 @@ pub use self::memory_store::MemoryStore;
 pub use self::migrate::Migrate;
 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 types::*;