State cache memory size WIP

Cargo.toml

@@ -149,7 +149,8 @@ libsecp256k1 = "0.7"
 log = "0.4"
 lru = "0.12"
 maplit = "1"
-milhouse = "0.3"
+# milhouse = "0.3"
+milhouse  = { git = "https://github.com/sigp/milhouse", branch = "mem-usage" }
 num_cpus = "1"
 parking_lot = "0.12"
 paste = "1"
@@ -274,3 +275,4 @@ incremental = false
 
 [patch.crates-io]
 quick-protobuf = { git = "https://github.com/sigp/quick-protobuf.git", rev = "681f413312404ab6e51f0b46f39b0075c6f4ebfd" }
+# metastruct = { git = "https://github.com/sigp/metastruct", rev = "v0.1.3" }

beacon_node/store/src/hot_cold_store.rs

@@ -437,7 +437,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
     ) -> Result<(), Error> {
         self.state_cache
             .lock()
-            .update_finalized_state(state_root, block_root, state)
+            .update_finalized_state(state_root, block_root, state, &self.log)
     }
 
     pub fn state_cache_len(&self) -> usize {

beacon_node/store/src/state_cache.rs

@@ -1,7 +1,9 @@
 use crate::Error;
 use lru::LruCache;
+use slog::{debug, Logger};
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::num::NonZeroUsize;
+use types::milhouse::mem::MemoryTracker;
 use types::{BeaconState, ChainSpec, Epoch, EthSpec, Hash256, Slot};
 
 /// Fraction of the LRU cache to leave intact during culling.
@@ -64,11 +66,54 @@ impl<E: EthSpec> StateCache<E> {
         self.states.cap().get()
     }
 
+    fn log_memory_stats(&self, log: &Logger) {
+        let mut mem_tracker = MemoryTracker::default();
+        let mut total_usage = 0;
+        if let Some(finalized_state) = &self.finalized_state {
+            let stats = mem_tracker.track_item(&finalized_state.state);
+            debug!(
+                log,
+                "Memory stats";
+                "slot" => finalized_state.state.slot(),
+                "block_slot" => finalized_state.state.latest_block_header().slot,
+                "total_kb" => stats.total_size / 1024,
+                "diff_kb" => stats.differential_size / 1024,
+            );
+            total_usage += stats.differential_size;
+        }
+        let mut states = self
+            .states
+            .iter()
+            .map(|(_, state)| state)
+            .collect::<Vec<_>>();
+        states.sort_by_key(|s| s.slot());
+
+        for state in states {
+            let stats = mem_tracker.track_item(state);
+            debug!(
+                log,
+                "Memory stats";
+                "slot" => state.slot(),
+                "block_slot" => state.latest_block_header().slot,
+                "total_kb" => stats.total_size / 1024,
+                "diff_kb" => stats.differential_size / 1024,
+            );
+            total_usage += stats.differential_size;
+        }
+        debug!(
+            log,
+            "Total memory stats";
+            "num_states" => self.states.len(),
+            "bytes_kb" => total_usage / 1024,
+        );
+    }
+
     pub fn update_finalized_state(
         &mut self,
         state_root: Hash256,
         block_root: Hash256,
         state: BeaconState<E>,
+        log: &Logger,
     ) -> Result<(), Error> {
         if state.slot() % E::slots_per_epoch() != 0 {
             return Err(Error::FinalizedStateUnaligned);
@@ -84,6 +129,10 @@ impl<E: EthSpec> StateCache<E> {
             return Err(Error::FinalizedStateDecreasingSlot);
         }
 
+        // Log memory states prior to pruning.
+        debug!(log, "Pre-pruning memory stats");
+        self.log_memory_stats(log);
+
         // Add to block map.
         self.block_map.insert(block_root, state.slot(), state_root);
 
@@ -97,6 +146,11 @@ impl<E: EthSpec> StateCache<E> {
 
         // Update finalized state.
         self.finalized_state = Some(FinalizedState { state_root, state });
+
+        // Log memory stats after pruning.
+        debug!(log, "Post-pruning memory stats");
+        self.log_memory_stats(log);
+
         Ok(())
     }
 
@@ -105,16 +159,19 @@ impl<E: EthSpec> StateCache<E> {
     /// This function should only be called on states that are likely not to already share tree
     /// nodes with the finalized state, e.g. states loaded from disk.
     ///
-    /// If the finalized state is not initialized this function is a no-op.
+    /// If the finalized state is not initialized this function is a no-op. Return `true` if the
+    /// state was rebased and `false`.
     pub fn rebase_on_finalized(
         &self,
         state: &mut BeaconState<E>,
         spec: &ChainSpec,
-    ) -> Result<(), Error> {
+    ) -> Result<bool, Error> {
         if let Some(finalized_state) = &self.finalized_state {
             state.rebase_on(&finalized_state.state, spec)?;
+            Ok(true)
+        } else {
+            Ok(false)
         }
-        Ok(())
     }
 
     /// Return a status indicating whether the state already existed in the cache.

consensus/types/src/beacon_state.rs

@@ -39,6 +39,7 @@ mod committee_cache;
 mod balance;
 mod exit_cache;
 mod iter;
+mod memsize;
 mod progressive_balances_cache;
 mod pubkey_cache;
 mod slashings_cache;

consensus/types/src/beacon_state/committee_cache.rs

@@ -425,3 +425,19 @@ impl Decode for NonZeroUsizeOption {
         four_byte_option_non_zero_usize::decode::from_ssz_bytes(bytes).map(Self)
     }
 }
+
+impl milhouse::mem::MemorySize for CommitteeCache {
+    fn self_pointer(&self) -> usize {
+        self as *const _ as usize
+    }
+
+    fn subtrees(&self) -> Vec<&dyn milhouse::mem::MemorySize> {
+        vec![]
+    }
+
+    fn intrinsic_size(&self) -> usize {
+        std::mem::size_of::<Self>()
+            + self.shuffling.len() * std::mem::size_of::<usize>()
+            + self.shuffling_positions.len() * std::mem::size_of::<usize>()
+    }
+}

consensus/types/src/beacon_state/memsize.rs

@@ -0,0 +1,111 @@
+use super::{
+    map_beacon_state_altair_tree_list_fields_immutable,
+    map_beacon_state_base_tree_list_fields_immutable,
+    map_beacon_state_bellatrix_tree_list_fields_immutable,
+    map_beacon_state_capella_tree_list_fields_immutable,
+    map_beacon_state_deneb_tree_list_fields_immutable,
+    map_beacon_state_electra_tree_list_fields_immutable,
+};
+use crate::{historical_summary::HistoricalSummary, *};
+use milhouse::{mem::MemorySize, List, Vector};
+
+impl<E: EthSpec> MemorySize for BeaconState<E> {
+    fn self_pointer(&self) -> usize {
+        self as *const _ as usize
+    }
+
+    fn subtrees<'a>(&'a self) -> Vec<&'a (dyn MemorySize + 'a)> {
+        let mut subtrees: Vec<&'a (dyn MemorySize + 'a)> = vec![];
+        match self {
+            Self::Base(self_inner) => {
+                map_beacon_state_base_tree_list_fields_immutable!(
+                    &'a _,
+                    self_inner,
+                    |_, self_field| {
+                        subtrees.push(self_field);
+                    }
+                );
+            }
+            Self::Altair(self_inner) => {
+                map_beacon_state_altair_tree_list_fields_immutable!(
+                    &'a _,
+                    self_inner,
+                    |_, self_field| {
+                        subtrees.push(self_field);
+                    }
+                );
+            }
+            Self::Bellatrix(self_inner) => {
+                map_beacon_state_bellatrix_tree_list_fields_immutable!(
+                    &'a _,
+                    self_inner,
+                    |_, self_field| {
+                        subtrees.push(self_field);
+                    }
+                );
+            }
+            Self::Capella(self_inner) => {
+                map_beacon_state_capella_tree_list_fields_immutable!(
+                    &'a _,
+                    self_inner,
+                    |_, self_field| {
+                        subtrees.push(self_field);
+                    }
+                );
+            }
+            Self::Deneb(self_inner) => {
+                map_beacon_state_deneb_tree_list_fields_immutable!(
+                    &'a _,
+                    self_inner,
+                    |_, self_field| {
+                        subtrees.push(self_field);
+                    }
+                );
+            }
+            Self::Electra(self_inner) => {
+                map_beacon_state_electra_tree_list_fields_immutable!(
+                    &'a _,
+                    self_inner,
+                    |_, self_field| {
+                        subtrees.push(self_field);
+                    }
+                );
+            }
+        }
+
+        if let Ok(current_sc) = self.current_sync_committee() {
+            subtrees.push(&**current_sc);
+        }
+        if let Ok(next_sc) = self.next_sync_committee() {
+            subtrees.push(&**next_sc);
+        }
+
+        for committee_cache in self.committee_caches() {
+            subtrees.push(&**committee_cache);
+        }
+
+        // FIXME(sproul): more caches
+
+        subtrees
+    }
+
+    fn intrinsic_size(&self) -> usize {
+        // This is a close-enough approximation for now.
+        std::mem::size_of::<Self>()
+    }
+}
+
+impl<E: EthSpec> MemorySize for SyncCommittee<E> {
+    fn self_pointer(&self) -> usize {
+        self as *const _ as usize
+    }
+
+    fn subtrees(&self) -> Vec<&dyn MemorySize> {
+        vec![]
+    }
+
+    #[allow(clippy::arithmetic_side_effects)]
+    fn intrinsic_size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.pubkeys.len() * std::mem::size_of::<PublicKeyBytes>()
+    }
+}