Persistent committee caches and exit cache

Cargo.lock

@@ -160,6 +160,15 @@ version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c5d78ce20460b82d3fa150275ed9d55e21064fc7951177baacf86a145c4a4b1f"
 
+[[package]]
+name = "archery"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a8da9bc4c4053ee067669762bcaeea6e241841295a2b6c948312dad6ef4cc02"
+dependencies = [
+ "static_assertions",
+]
+
 [[package]]
 name = "arrayref"
 version = "0.3.6"
@@ -5032,6 +5041,15 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "rpds"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4ef5140bcb576bfd6d56cd2de709a7d17851ac1f3805e67fe9d99e42a11821f"
+dependencies = [
+ "archery",
+]
+
 [[package]]
 name = "rusqlite"
 version = "0.25.4"
@@ -6579,6 +6597,7 @@ dependencies = [
  "rand_xorshift",
  "rayon",
  "regex",
+ "rpds",
  "rusqlite",
  "safe_arith",
  "serde",

beacon_node/store/src/impls/beacon_state.rs

@@ -2,6 +2,7 @@ use crate::*;
 use ssz::{DecodeError, Encode};
 use ssz_derive::Encode;
 use std::convert::TryInto;
+use std::sync::Arc;
 use types::beacon_state::{CloneConfig, CommitteeCache, CACHED_EPOCHS};
 
 pub fn store_full_state<E: EthSpec>(
@@ -48,7 +49,7 @@ pub fn get_full_state<KV: KeyValueStore<E>, E: EthSpec>(
 #[derive(Encode)]
 pub struct StorageContainer<T: EthSpec> {
     state: BeaconState<T>,
-    committee_caches: Vec<CommitteeCache>,
+    committee_caches: Vec<Arc<CommitteeCache>>,
 }
 
 impl<T: EthSpec> StorageContainer<T> {

beacon_node/store/src/iter.rs

@@ -409,15 +409,15 @@ mod test {
         let mut hashes = (0..).map(Hash256::from_low_u64_be);
         let roots_a = state_a.block_roots_mut();
         for i in 0..roots_a.len() {
-            roots_a[i] = hashes.next().unwrap()
+            *roots_a.get_mut(i).unwrap() = hashes.next().unwrap()
         }
         let roots_b = state_b.block_roots_mut();
         for i in 0..roots_b.len() {
-            roots_b[i] = hashes.next().unwrap()
+            *roots_b.get_mut(i).unwrap() = hashes.next().unwrap()
         }
 
         let state_a_root = hashes.next().unwrap();
-        state_b.state_roots_mut()[0] = state_a_root;
+        *state_b.state_roots_mut().get_mut(0).unwrap() = state_a_root;
         store.put_state(&state_a_root, &state_a).unwrap();
 
         let iter = BlockRootsIterator::new(&store, &state_b);

beacon_node/store/src/lib.rs

@@ -25,6 +25,7 @@ pub mod metadata;
 pub mod metrics;
 mod partial_beacon_state;
 pub mod reconstruct;
+mod state_cache;
 
 pub mod iter;
 

beacon_node/store/src/state_cache.rs

@@ -0,0 +1,18 @@
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::mem::size_of;
+    use types::{
+        beacon_state::PubkeyCache, BeaconBlockHeader, BeaconState, BeaconStateAltair,
+        BeaconStateMerge, MainnetEthSpec,
+    };
+
+    #[test]
+    fn state_size() {
+        println!("{}", size_of::<BeaconStateAltair<MainnetEthSpec>>());
+        println!("{}", size_of::<BeaconStateMerge<MainnetEthSpec>>());
+        println!("{}", size_of::<BeaconState<MainnetEthSpec>>());
+        println!("{}", size_of::<PubkeyCache>());
+        assert!(false);
+    }
+}

consensus/types/Cargo.toml

@@ -46,6 +46,7 @@ itertools = "0.10.0"
 superstruct = "0.4.0"
 serde_json = "1.0.74"
 milhouse = { path = "../../../milhouse", optional = true }
+rpds = "0.11.0"
 
 [dev-dependencies]
 criterion = "0.3.3"

consensus/types/src/beacon_state.rs

@@ -7,7 +7,7 @@ use compare_fields_derive::CompareFields;
 use derivative::Derivative;
 use eth2_hashing::hash;
 use int_to_bytes::{int_to_bytes4, int_to_bytes8};
-use pubkey_cache::PubkeyCache;
+pub use pubkey_cache::PubkeyCache;
 use safe_arith::{ArithError, SafeArith};
 use serde_derive::{Deserialize, Serialize};
 use ssz::{ssz_encode, Decode, DecodeError, Encode};
@@ -314,7 +314,7 @@ where
     #[tree_hash(skip_hashing)]
     #[test_random(default)]
     #[derivative(Clone(clone_with = "clone_default"))]
-    pub committee_caches: [CommitteeCache; CACHED_EPOCHS],
+    pub committee_caches: [Arc<CommitteeCache>; CACHED_EPOCHS],
     #[serde(skip_serializing, skip_deserializing)]
     #[ssz(skip_serializing, skip_deserializing)]
     #[tree_hash(skip_hashing)]
@@ -347,6 +347,7 @@ impl<T: EthSpec> BeaconState<T> {
     ///
     /// Not a complete genesis state, see `initialize_beacon_state_from_eth1`.
     pub fn new(genesis_time: u64, eth1_data: Eth1Data, spec: &ChainSpec) -> Self {
+        let default_committee_cache = Arc::new(CommitteeCache::default());
         BeaconState::Base(BeaconStateBase {
             // Versioning
             genesis_time,
@@ -392,9 +393,9 @@ impl<T: EthSpec> BeaconState<T> {
             // Caching (not in spec)
             total_active_balance: None,
             committee_caches: [
-                CommitteeCache::default(),
+                default_committee_cache.clone(),
-                CommitteeCache::default(),
+                default_committee_cache.clone(),
-                CommitteeCache::default(),
+                default_committee_cache,
             ],
             pubkey_cache: PubkeyCache::default(),
             exit_cache: ExitCache::default(),
@@ -1464,7 +1465,7 @@ impl<T: EthSpec> BeaconState<T> {
         &self,
         epoch: Epoch,
         spec: &ChainSpec,
-    ) -> Result<CommitteeCache, Error> {
+    ) -> Result<Arc<CommitteeCache>, Error> {
         CommitteeCache::initialized(self, epoch, spec)
     }
 
@@ -1502,7 +1503,7 @@ impl<T: EthSpec> BeaconState<T> {
         *self.committee_cache_at_index_mut(curr)? = curr_cache;
 
         let next = Self::committee_cache_index(RelativeEpoch::Next);
-        *self.committee_cache_at_index_mut(next)? = CommitteeCache::default();
+        *self.committee_cache_at_index_mut(next)? = Arc::new(CommitteeCache::default());
         Ok(())
     }
 
@@ -1527,11 +1528,15 @@ impl<T: EthSpec> BeaconState<T> {
     fn committee_cache_at_index(&self, index: usize) -> Result<&CommitteeCache, Error> {
         self.committee_caches()
             .get(index)
+            .map(Arc::as_ref)
             .ok_or(Error::CommitteeCachesOutOfBounds(index))
     }
 
     /// Get a mutable reference to the committee cache at a given index.
-    fn committee_cache_at_index_mut(&mut self, index: usize) -> Result<&mut CommitteeCache, Error> {
+    fn committee_cache_at_index_mut(
+        &mut self,
+        index: usize,
+    ) -> Result<&mut Arc<CommitteeCache>, Error> {
         self.committee_caches_mut()
             .get_mut(index)
             .ok_or(Error::CommitteeCachesOutOfBounds(index))
@@ -1553,7 +1558,7 @@ impl<T: EthSpec> BeaconState<T> {
     /// Drops the cache, leaving it in an uninitialized state.
     pub fn drop_committee_cache(&mut self, relative_epoch: RelativeEpoch) -> Result<(), Error> {
         *self.committee_cache_at_index_mut(Self::committee_cache_index(relative_epoch))? =
-            CommitteeCache::default();
+            Arc::new(CommitteeCache::default());
         Ok(())
     }
 

consensus/types/src/beacon_state/committee_cache.rs

@@ -8,6 +8,7 @@ use serde_derive::{Deserialize, Serialize};
 use ssz::{four_byte_option_impl, Decode, DecodeError, Encode};
 use ssz_derive::{Decode, Encode};
 use std::ops::Range;
+use std::sync::Arc;
 use swap_or_not_shuffle::shuffle_list;
 
 mod tests;
@@ -37,7 +38,7 @@ impl CommitteeCache {
         state: &BeaconState<T>,
         epoch: Epoch,
         spec: &ChainSpec,
-    ) -> Result<CommitteeCache, Error> {
+    ) -> Result<Arc<CommitteeCache>, Error> {
         RelativeEpoch::from_epoch(state.current_epoch(), epoch)
             .map_err(|_| Error::EpochOutOfBounds)?;
 
@@ -77,13 +78,13 @@ impl CommitteeCache {
                 .ok_or(Error::ShuffleIndexOutOfBounds(v))? = NonZeroUsize::new(i + 1).into();
         }
 
-        Ok(CommitteeCache {
+        Ok(Arc::new(CommitteeCache {
             initialized_epoch: Some(epoch),
             shuffling,
             shuffling_positions,
             committees_per_slot,
             slots_per_epoch: T::slots_per_epoch(),
-        })
+        }))
     }
 
     /// Returns `true` if the cache has been initialized at the supplied `epoch`.

consensus/types/src/beacon_state/exit_cache.rs

@@ -1,13 +1,12 @@
 use super::{BeaconStateError, ChainSpec, Epoch, Validator};
+use rpds::HashTrieMapSync as HashTrieMap;
 use safe_arith::SafeArith;
-use serde_derive::{Deserialize, Serialize};
-use std::collections::HashMap;
 
 /// Map from exit epoch to the number of validators with that exit epoch.
-#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, PartialEq)]
 pub struct ExitCache {
     initialized: bool,
-    exit_epoch_counts: HashMap<Epoch, u64>,
+    exit_epoch_counts: HashTrieMap<Epoch, u64>,
 }
 
 impl ExitCache {
@@ -41,10 +40,12 @@ impl ExitCache {
     /// Record the exit epoch of a validator. Must be called only once per exiting validator.
     pub fn record_validator_exit(&mut self, exit_epoch: Epoch) -> Result<(), BeaconStateError> {
         self.check_initialized()?;
-        self.exit_epoch_counts
+
-            .entry(exit_epoch)
+        if let Some(count) = self.exit_epoch_counts.get_mut(&exit_epoch) {
-            .or_insert(0)
+            count.safe_add_assign(1)?;
-            .safe_add_assign(1)?;
+        } else {
+            self.exit_epoch_counts.insert_mut(exit_epoch, 1);
+        }
         Ok(())
     }