Memory Aware Caching in Lighthouse

Cargo.toml

@@ -188,7 +188,7 @@ malloc_utils = { path = "common/malloc_utils" }
 maplit = "1"
 merkle_proof = { path = "consensus/merkle_proof" }
 metrics = { path = "common/metrics" }
-milhouse = { version = "0.7", default-features = false }
+milhouse = { git = "https://github.com/sigp/milhouse", branch = "main", default-features = false }
 mockall = "0.13"
 mockall_double = "0.3"
 mockito = "1.5.0"

beacon_node/beacon_chain/src/chain_config.rs

@@ -10,7 +10,7 @@ pub const DEFAULT_RE_ORG_MAX_EPOCHS_SINCE_FINALIZATION: Epoch = Epoch::new(2);
 /// Default to 1/12th of the slot, which is 1 second on mainnet.
 pub const DEFAULT_RE_ORG_CUTOFF_DENOMINATOR: u32 = 12;
 pub const DEFAULT_FORK_CHOICE_BEFORE_PROPOSAL_TIMEOUT: u64 = 250;
-
+pub const DEFAULT_STATE_CACHE_MAX_BYTES: usize = 536870912;
 /// Default fraction of a slot lookahead for payload preparation (12/3 = 4 seconds on mainnet).
 pub const DEFAULT_PREPARE_PAYLOAD_LOOKAHEAD_FACTOR: u32 = 3;
 

beacon_node/src/cli.rs

@@ -816,6 +816,15 @@ pub fn cli_app() -> Command {
                 .action(ArgAction::Set)
                 .display_order(0)
         )
+        .arg(
+            Arg::new("state-cache-max-bytes")
+                .long("state-cache-max-bytes")
+                .value_name("BYTES")
+                .help("Specifies the maximum size of the state cache in bytes")
+                .default_value("536870912")
+                .action(ArgAction::Set)
+                .display_order(0)
+        )
         /*
          * Execution Layer Integration
          */

beacon_node/src/config.rs

@@ -377,6 +377,12 @@ pub fn get_config<E: EthSpec>(
             .map_err(|_| "state-cache-size is not a valid integer".to_string())?;
     }
 
+    if let Some(max_bytes) = cli_args.get_one::<String>("state-cache-max-bytes") {
+        client_config.store.max_state_cache_bytes = max_bytes
+            .parse()
+            .map_err(|_| "state-cache-max-bytes is not a valid integer".to_string())?;
+    }
+
     if let Some(historic_state_cache_size) =
         clap_utils::parse_optional(cli_args, "historic-state-cache-size")?
     {

beacon_node/store/Cargo.toml

@@ -20,6 +20,7 @@ leveldb = { version = "0.8.6", optional = true, default-features = false }
 logging = { workspace = true }
 lru = { workspace = true }
 metrics = { workspace = true }
+milhouse = { workspace = true }
 parking_lot = { workspace = true }
 redb = { version = "2.1.3", optional = true }
 safe_arith = { workspace = true }
@@ -44,3 +45,7 @@ tempfile = { workspace = true }
 [[bench]]
 name = "hdiff"
 harness = false
+
+[[bench]]
+name = "state_cache"
+harness = false

beacon_node/store/benches/state_cache.rs

@@ -0,0 +1,114 @@
+use bls::{FixedBytesExtended, Hash256, PublicKeyBytes};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use rand::Rng;
+use ssz::Decode;
+use std::num::NonZeroUsize;
+use store::state_cache::StateCache;
+
+use types::{
+    BeaconState, ChainSpec, Epoch, Eth1Data, EthSpec, MainnetEthSpec as E, Slot, Validator,
+};
+
+fn build_state(
+    spec: &ChainSpec,
+    slot: Slot,
+    validator_count: usize,
+    rng: &mut impl Rng,
+) -> BeaconState<E> {
+    let genesis_time = 0;
+    let eth1_data = Eth1Data::default();
+    let mut state = BeaconState::<E>::new(genesis_time, eth1_data, spec);
+
+    for _ in 0..validator_count {
+        append_validator(&mut state, rng);
+    }
+
+    *state.slot_mut() = slot;
+    state.latest_block_header_mut().slot = slot;
+    state.apply_pending_mutations().unwrap();
+    state
+}
+
+fn append_validator(state: &mut BeaconState<E>, mut rng: &mut impl Rng) {
+    state
+        .balances_mut()
+        .push(32_000_000_000 + rng.random_range(1..=1_000_000_000))
+        .unwrap();
+    if let Ok(inactivity_scores) = state.inactivity_scores_mut() {
+        inactivity_scores.push(0).unwrap();
+    }
+    state
+        .validators_mut()
+        .push(rand_validator(&mut rng))
+        .unwrap();
+}
+
+fn rand_validator(rng: &mut impl Rng) -> Validator {
+    let mut pubkey = [0u8; 48];
+    rng.fill_bytes(&mut pubkey);
+    let withdrawal_credentials: [u8; 32] = rng.random();
+
+    Validator {
+        pubkey: PublicKeyBytes::from_ssz_bytes(&pubkey).unwrap(),
+        withdrawal_credentials: withdrawal_credentials.into(),
+        slashed: false,
+        effective_balance: 32_000_000_000,
+        activation_eligibility_epoch: Epoch::max_value(),
+        activation_epoch: Epoch::max_value(),
+        exit_epoch: Epoch::max_value(),
+        withdrawable_epoch: Epoch::max_value(),
+    }
+}
+
+pub fn all_benches(c: &mut Criterion) {
+    let spec = E::default_spec();
+    let mut rng = rand::rng();
+    let num_states = 20;
+    let validator_count = 1024;
+
+    let states: Vec<(Hash256, Hash256, BeaconState<E>)> = (0..num_states)
+        .map(|i| {
+            let slot = Slot::new(i as u64);
+            let state = build_state(&spec, slot, validator_count, &mut rng);
+            let root = Hash256::from_low_u64_le(i as u64 + 1);
+            (root, root, state)
+        }).collect();
+
+    let capacity = NonZeroUsize::new(num_states).unwrap();
+    let headroom = NonZeroUsize::new(1).unwrap();
+    let hdiff_capacity = NonZeroUsize::new(1).unwrap();
+
+        c.bench_function("state_cache_insert_without_memory_limit", |b| {
+        b.iter_batched(
+            || StateCache::new(capacity, headroom, hdiff_capacity, usize::MAX),
+            |mut cache| {
+                for (state_root, block_root, state) in &states {
+                    cache.put_state(*state_root, *block_root, state).unwrap();
+                }
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    let low_max_bytes = 1_000_000;
+    c.bench_function("state_cache_insert_with_memory_limit", |b| {
+        b.iter_batched(
+            || StateCache::new(capacity, headroom, hdiff_capacity, low_max_bytes),
+            |mut cache| {
+                for (state_root, block_root, state) in &states {
+                    cache.put_state(*state_root, *block_root, state).unwrap();
+                }
+                assert!(cache.cached_bytes() <= cache.max_cached_bytes());
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+}
+
+criterion_group!{
+    name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = all_benches
+}
+criterion_main!(benches);

beacon_node/store/src/config.rs

@@ -21,6 +21,7 @@ pub const DEFAULT_SLOTS_PER_RESTORE_POINT: u64 = 8192;
 pub const DEFAULT_EPOCHS_PER_STATE_DIFF: u64 = 8;
 pub const DEFAULT_BLOCK_CACHE_SIZE: NonZeroUsize = new_non_zero_usize(64);
 pub const DEFAULT_STATE_CACHE_SIZE: NonZeroUsize = new_non_zero_usize(128);
+pub const DEFAULT_STATE_CACHE_MAX_BYTES: usize = 512 * 1024 * 1024;
 pub const DEFAULT_STATE_CACHE_HEADROOM: NonZeroUsize = new_non_zero_usize(1);
 pub const DEFAULT_COMPRESSION_LEVEL: i32 = 1;
 pub const DEFAULT_HISTORIC_STATE_CACHE_SIZE: NonZeroUsize = new_non_zero_usize(1);
@@ -37,6 +38,8 @@ pub struct StoreConfig {
     pub block_cache_size: NonZeroUsize,
     /// Maximum number of states to store in the in-memory state cache.
     pub state_cache_size: NonZeroUsize,
+    /// Maximum number of bytes to store in the in-memory state cache.
+    pub max_state_cache_bytes: usize,
     /// Minimum number of states to cull from the state cache upon fullness.
     pub state_cache_headroom: NonZeroUsize,
     /// Compression level for blocks, state diffs and other compressed values.
@@ -107,6 +110,7 @@ impl Default for StoreConfig {
         Self {
             block_cache_size: DEFAULT_BLOCK_CACHE_SIZE,
             state_cache_size: DEFAULT_STATE_CACHE_SIZE,
+            max_state_cache_bytes: DEFAULT_STATE_CACHE_MAX_BYTES,
             state_cache_headroom: DEFAULT_STATE_CACHE_HEADROOM,
             historic_state_cache_size: DEFAULT_HISTORIC_STATE_CACHE_SIZE,
             cold_hdiff_buffer_cache_size: DEFAULT_COLD_HDIFF_BUFFER_CACHE_SIZE,

beacon_node/store/src/hot_cold_store.rs

@@ -234,6 +234,7 @@ impl<E: EthSpec> HotColdDB<E, MemoryStore<E>, MemoryStore<E>> {
                 config.state_cache_size,
                 config.state_cache_headroom,
                 config.hot_hdiff_buffer_cache_size,
+                config.max_state_cache_bytes,
             )),
             historic_state_cache: Mutex::new(HistoricStateCache::new(
                 config.cold_hdiff_buffer_cache_size,
@@ -286,6 +287,7 @@ impl<E: EthSpec> HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>> {
                 config.state_cache_size,
                 config.state_cache_headroom,
                 config.hot_hdiff_buffer_cache_size,
+                config.max_state_cache_bytes,
             )),
             historic_state_cache: Mutex::new(HistoricStateCache::new(
                 config.cold_hdiff_buffer_cache_size,
@@ -501,6 +503,10 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             &metrics::STORE_BEACON_STATE_CACHE_SIZE,
             state_cache.len() as i64,
         );
+        metrics::set_gauge(
+            &metrics::STORE_BEACON_STATE_CACHE_MEMORY_SIZE,
+            state_cache.cached_bytes() as i64,
+        );
         metrics::set_gauge_vec(
             &metrics::STORE_BEACON_HDIFF_BUFFER_CACHE_SIZE,
             HOT_METRIC,

beacon_node/store/src/lib.rs

@@ -19,6 +19,7 @@ pub mod historic_state_cache;
 pub mod hot_cold_store;
 mod impls;
 mod memory_store;
+pub mod memsize;
 pub mod metadata;
 pub mod metrics;
 pub mod partial_beacon_state;

beacon_node/store/src/memsize.rs

@@ -0,0 +1,44 @@
+use crate::metrics::BEACON_STATE_MEMORY_SIZE_CALCULATION_TIME;
+use milhouse::mem::{MemorySize, MemoryTracker};
+use std::time::Instant;
+use types::{BeaconState, EthSpec};
+
+/// BeaconState Wrapper for memory tracking.
+pub struct BeaconStateWrapper<'a, E: EthSpec>(pub &'a BeaconState<E>);
+
+impl<'a, E: EthSpec> MemorySize for BeaconStateWrapper<'a, E> {
+    fn self_pointer(&self) -> usize {
+        self.0 as *const _ as usize
+    }
+
+    fn subtrees(&self) -> Vec<&dyn MemorySize> {
+        vec![]
+    }
+
+    fn intrinsic_size(&self) -> usize {
+        std::mem::size_of::<Self>()
+    }
+}
+
+/// Extension trait for approximate memory consumption of a `BeaconState`.
+pub trait BeaconStateMemorySize {
+    fn memory_size(&self) -> usize;
+}
+
+impl<E: EthSpec> BeaconStateMemorySize for BeaconState<E> {
+    fn memory_size(&self) -> usize {
+        let wrapper = BeaconStateWrapper(self);
+        // Timer for MemorySize
+        let timer = Instant::now();
+        // Use MemoryTracker on the wrapper
+        let mut tracker = MemoryTracker::default();
+        let stats = tracker.track_item(&wrapper);
+
+        let elapsed_time = timer.elapsed();
+        metrics::observe(
+            &BEACON_STATE_MEMORY_SIZE_CALCULATION_TIME,
+            elapsed_time.as_secs_f64(),
+        );
+        stats.differential_size
+    }
+}

beacon_node/store/src/metrics.rs

@@ -269,6 +269,19 @@ pub static STORE_BEACON_STATE_CACHE_SIZE: LazyLock<Result<IntGauge>> = LazyLock:
         "Current count of items in beacon store state cache",
     )
 });
+pub static STORE_BEACON_STATE_CACHE_MEMORY_SIZE: LazyLock<Result<IntGauge>> = LazyLock::new(|| {
+    try_create_int_gauge(
+        "store_beacon_state_cache_memory_size",
+        "Memory consumed by items in the beacon store state cache",
+    )
+});
+pub static BEACON_STATE_MEMORY_SIZE_CALCULATION_TIME: LazyLock<Result<Histogram>> =
+    LazyLock::new(|| {
+        try_create_histogram(
+            "beacon_state_memory_size_calculation_time",
+            "Time taken to calculate the memory size of a beacon state.",
+        )
+    });
 pub static STORE_BEACON_HISTORIC_STATE_CACHE_SIZE: LazyLock<Result<IntGauge>> =
     LazyLock::new(|| {
         try_create_int_gauge(