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

@@ -454,7 +454,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// Returns the validator index (if any) for the given public key.
     ///
     /// Information is retrieved from the present `beacon_state.validators`.
-    pub fn validator_index(&self, pubkey: &PublicKey) -> Option<usize> {
+    pub fn validator_index(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
         for (i, validator) in self.head().beacon_state.validators.iter().enumerate() {
             if validator.pubkey == *pubkey {
                 return Some(i);

beacon_node/rest_api/src/helpers.rs

@@ -1,6 +1,6 @@
 use crate::{ApiError, ApiResult};
 use beacon_chain::{BeaconChain, BeaconChainTypes};
-use bls::PublicKey;
+use bls::PublicKeyBytes;
 use eth2_libp2p::{PubsubMessage, Topic};
 use eth2_libp2p::{
     BEACON_ATTESTATION_TOPIC, BEACON_BLOCK_TOPIC, TOPIC_ENCODING_POSTFIX, TOPIC_PREFIX,
@@ -99,12 +99,12 @@ pub fn parse_root(string: &str) -> Result<Hash256, ApiError> {
 }
 
 /// Parse a PublicKey from a `0x` prefixed hex string
-pub fn parse_pubkey(string: &str) -> Result<PublicKey, ApiError> {
+pub fn parse_pubkey_bytes(string: &str) -> Result<PublicKeyBytes, ApiError> {
     const PREFIX: &str = "0x";
     if string.starts_with(PREFIX) {
         let pubkey_bytes = hex::decode(string.trim_start_matches(PREFIX))
             .map_err(|e| ApiError::BadRequest(format!("Invalid hex string: {:?}", e)))?;
-        let pubkey = PublicKey::from_bytes(pubkey_bytes.as_slice()).map_err(|e| {
+        let pubkey = PublicKeyBytes::from_bytes(pubkey_bytes.as_slice()).map_err(|e| {
             ApiError::BadRequest(format!("Unable to deserialize public key: {:?}.", e))
         })?;
         Ok(pubkey)

beacon_node/rest_api/src/lib.rs

@@ -37,7 +37,7 @@ use tokio::runtime::TaskExecutor;
 use tokio::sync::mpsc;
 use url_query::UrlQuery;
 
-pub use crate::helpers::parse_pubkey;
+pub use crate::helpers::parse_pubkey_bytes;
 pub use beacon::{BlockResponse, HeadResponse, StateResponse};
 pub use config::Config;
 pub use validator::{BulkValidatorDutiesRequest, ValidatorDuty};

beacon_node/rest_api/src/validator.rs

@@ -1,5 +1,5 @@
 use crate::helpers::{
-    check_content_type_for_json, parse_pubkey, publish_attestation_to_network,
+    check_content_type_for_json, parse_pubkey_bytes, publish_attestation_to_network,
     publish_beacon_block_to_network,
 };
 use crate::response_builder::ResponseBuilder;
@@ -7,7 +7,7 @@ use crate::{ApiError, ApiResult, BoxFut, NetworkChannel, UrlQuery};
 use beacon_chain::{
     AttestationProcessingOutcome, BeaconChain, BeaconChainTypes, BlockProcessingOutcome,
 };
-use bls::PublicKey;
+use bls::PublicKeyBytes;
 use futures::future::Future;
 use futures::stream::Stream;
 use hyper::{Body, Request};
@@ -21,7 +21,7 @@ use types::{Attestation, BeaconBlock, CommitteeIndex, Epoch, RelativeEpoch, Slot
 #[derive(PartialEq, Debug, Serialize, Deserialize, Clone)]
 pub struct ValidatorDuty {
     /// The validator's BLS public key, uniquely identifying them. _48-bytes, hex encoded with 0x prefix, case insensitive._
-    pub validator_pubkey: PublicKey,
+    pub validator_pubkey: PublicKeyBytes,
     /// The slot at which the validator must attest.
     pub attestation_slot: Option<Slot>,
     /// The index of the committee within `slot` of which the validator is a member.
@@ -35,7 +35,7 @@ pub struct ValidatorDuty {
 #[derive(PartialEq, Debug, Serialize, Deserialize, Clone, Encode, Decode)]
 pub struct BulkValidatorDutiesRequest {
     pub epoch: Epoch,
-    pub pubkeys: Vec<PublicKey>,
+    pub pubkeys: Vec<PublicKeyBytes>,
 }
 
 /// HTTP Handler to retrieve a the duties for a set of validators during a particular epoch. This
@@ -60,7 +60,11 @@ pub fn post_validator_duties<T: BeaconChainTypes>(
             })
         })
         .and_then(|bulk_request| {
-            return_validator_duties(beacon_chain, bulk_request.epoch, bulk_request.pubkeys)
+            return_validator_duties(
+                beacon_chain,
+                bulk_request.epoch,
+                bulk_request.pubkeys.into_iter().map(Into::into).collect(),
+            )
         })
         .and_then(|duties| response_builder?.body_no_ssz(&duties));
 
@@ -80,7 +84,7 @@ pub fn get_validator_duties<T: BeaconChainTypes>(
     let validator_pubkeys = query
         .all_of("validator_pubkeys")?
         .iter()
-        .map(|validator_pubkey_str| parse_pubkey(validator_pubkey_str))
+        .map(|validator_pubkey_str| parse_pubkey_bytes(validator_pubkey_str))
         .collect::<Result<_, _>>()?;
 
     let duties = return_validator_duties(beacon_chain, epoch, validator_pubkeys)?;
@@ -91,7 +95,7 @@ pub fn get_validator_duties<T: BeaconChainTypes>(
 fn return_validator_duties<T: BeaconChainTypes>(
     beacon_chain: Arc<BeaconChain<T>>,
     epoch: Epoch,
-    validator_pubkeys: Vec<PublicKey>,
+    validator_pubkeys: Vec<PublicKeyBytes>,
 ) -> Result<Vec<ValidatorDuty>, ApiError> {
     let slots_per_epoch = T::EthSpec::slots_per_epoch();
     let head_epoch = beacon_chain.head().beacon_state.current_epoch();

beacon_node/rest_api/tests/test.rs

@@ -6,6 +6,7 @@ use node_test_rig::{
     testing_client_config, ClientConfig, ClientGenesis, LocalBeaconNode,
 };
 use remote_beacon_node::{PublishStatus, ValidatorDuty};
+use std::convert::TryInto;
 use std::sync::Arc;
 use tree_hash::TreeHash;
 use types::{
@@ -182,7 +183,7 @@ fn validator_duties_bulk() {
         .beacon_state
         .validators
         .iter()
-        .map(|v| v.pubkey.clone())
+        .map(|v| (&v.pubkey).try_into().expect("pubkey should be valid"))
         .collect::<Vec<_>>();
 
     let duties = env
@@ -219,7 +220,7 @@ fn validator_duties() {
         .beacon_state
         .validators
         .iter()
-        .map(|v| v.pubkey.clone())
+        .map(|v| (&v.pubkey).try_into().expect("pubkey should be valid"))
         .collect::<Vec<_>>();
 
     let duties = env
@@ -270,10 +271,16 @@ fn check_duties<T: BeaconChainTypes>(
         .iter()
         .zip(duties.iter())
         .for_each(|(validator, duty)| {
-            assert_eq!(*validator, duty.validator_pubkey, "pubkey should match");
+            assert_eq!(
+                *validator,
+                (&duty.validator_pubkey)
+                    .try_into()
+                    .expect("should be valid pubkey"),
+                "pubkey should match"
+            );
 
             let validator_index = state
-                .get_validator_index(validator)
+                .get_validator_index(&validator.clone().into())
                 .expect("should have pubkey cache")
                 .expect("pubkey should exist");
 

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::*;