Fix test OOM issues on tree-states-hot (#7176)

* Try rebasing

* Heaptrack

* Try to prevent and detect circular references

* Don't use StateRootsIterator in get_ancestor_state_root

* Skip one slot in get_ancestor_state_root

* Heaptrack clippy fix

* Fix store test

* Log diff base state root

* Jump by slots per historical vector

* Maybe ok fix for test

* get_ancestor_state_root

* lint

* Implement downgrade from v24

* Better on state write

* Log on ws test chain dump

* Fix WSS tests

* Tweak condition

* Revert "Heaptrack"

This reverts commit f69c38c.

* Reduce diff

* Use debug repr

* Spelling

---------

Co-authored-by: dapplion <[email protected]>

Author: michaelsproul

beacon_node/beacon_chain/src/builder.rs

@@ -380,7 +380,8 @@ where
         self.pending_io_batch.push(
             store
                 .init_anchor_info(
-                    genesis_beacon_block.message().parent_root(),
+                    genesis_beacon_block.parent_root(),
+                    genesis_beacon_block.slot(),
                     Slot::new(0),
                     retain_historic_states,
                 )
@@ -532,14 +533,21 @@ where
             .cold_db
             .do_atomically(block_root_batch)
             .map_err(|e| format!("Error writing frozen block roots: {e:?}"))?;
+        debug!(
+            from = %weak_subj_block.slot(),
+            to_excl = %weak_subj_state.slot(),
+            block_root = ?weak_subj_block_root,
+            "Stored frozen block roots at skipped slots"
+        );
 
         // Write the anchor to memory before calling `put_state` otherwise hot hdiff can't store
         // states that do not align with the start_slot grid
         let retain_historic_states = self.chain_config.reconstruct_historic_states;
         self.pending_io_batch.push(
             store
                 .init_anchor_info(
-                    weak_subj_block.message().parent_root(),
+                    weak_subj_block.parent_root(),
+                    weak_subj_block.slot(),
                     weak_subj_slot,
                     retain_historic_states,
                 )

beacon_node/beacon_chain/src/historical_blocks.rs

@@ -166,6 +166,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
             // Store block roots, including at all skip slots in the freezer DB.
             for slot in (block.slot().as_u64()..prev_block_slot.as_u64()).rev() {
+                debug!(%slot, ?block_root, "Storing frozen block to root mapping");
                 cold_batch.push(KeyValueStoreOp::PutKeyValue(
                     DBColumn::BeaconBlockRoots,
                     slot.to_be_bytes().to_vec(),

beacon_node/beacon_chain/src/migrate.rs

@@ -471,17 +471,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
             let state_summaries = store
                 .load_hot_state_summaries()?
                 .into_iter()
-                .map(|(state_root, summary)| {
-                    (
-                        state_root,
-                        DAGStateSummary {
-                            slot: summary.slot,
-                            latest_block_root: summary.latest_block_root,
-                            latest_block_slot: summary.latest_block_slot,
-                            previous_state_root: summary.previous_state_root,
-                        },
-                    )
-                })
+                .map(|(state_root, summary)| (state_root, summary.into()))
                 .collect::<Vec<(Hash256, DAGStateSummary)>>();
 
             // Sanity check, there is at least one summary with the new finalized block root

beacon_node/beacon_chain/src/schema_change/migration_schema_v24.rs

@@ -1,16 +1,17 @@
 use crate::{
     beacon_chain::BeaconChainTypes,
-    summaries_dag::{DAGStateSummaryV22, StateSummariesDAG},
+    summaries_dag::{DAGStateSummary, DAGStateSummaryV22, StateSummariesDAG},
 };
-use ssz::Decode;
+use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
 use std::{
     sync::Arc,
     time::{Duration, Instant},
 };
 use store::{
-    get_full_state_v22, hdiff::StorageStrategy, hot_cold_store::DiffBaseStateRoot, DBColumn, Error,
-    HotColdDB, HotStateSummary, KeyValueStore, KeyValueStoreOp, StoreItem,
+    get_full_state_v22, hdiff::StorageStrategy, hot_cold_store::DiffBaseStateRoot,
+    store_full_state_v22, DBColumn, Error, HotColdDB, HotStateSummary, KeyValueStore,
+    KeyValueStoreOp, StoreItem,
 };
 use tracing::{debug, info, warn};
 use types::{EthSpec, Hash256, Slot};
@@ -94,8 +95,8 @@ pub fn upgrade_to_v24<T: BeaconChainTypes>(
         slots_count = summaries_by_slot.len(),
         min_slot = ?summaries_by_slot.first_key_value().map(|(slot, _)| slot),
         max_slot = ?summaries_by_slot.last_key_value().map(|(slot, _)| slot),
-        state_summaries_dag_roots = ?state_summaries_dag_roots,
-        hot_hdiff_start_slot = %hot_hdiff_start_slot,
+        ?state_summaries_dag_roots,
+        %hot_hdiff_start_slot,
         split_state_root = ?split.state_root,
         "Starting hot states migration"
     );
@@ -178,21 +179,22 @@ pub fn upgrade_to_v24<T: BeaconChainTypes>(
                                 })?
                         };
 
-                        let diff_base_state_root =
-                            if let Some(diff_base_slot) = storage_strategy.diff_base_slot() {
-                                DiffBaseStateRoot::new(
+                        let diff_base_state_root = if let Some(diff_base_slot) =
+                            storage_strategy.diff_base_slot()
+                        {
+                            DiffBaseStateRoot::new(
                                     diff_base_slot,
                                     state_summaries_dag
                                         .ancestor_state_root_at_slot(state_root, diff_base_slot)
                                         .map_err(|e| {
                                             Error::MigrationError(format!(
-                                                "error computing ancestor_state_root_at_slot {e:?}"
+                                                "error computing ancestor_state_root_at_slot({state_root:?}, {diff_base_slot}) {e:?}"
                                             ))
                                         })?,
                                 )
-                            } else {
-                                DiffBaseStateRoot::zero()
-                            };
+                        } else {
+                            DiffBaseStateRoot::zero()
+                        };
 
                         let new_summary = HotStateSummary {
                             slot,
@@ -227,8 +229,8 @@ pub fn upgrade_to_v24<T: BeaconChainTypes>(
             if last_log_time.elapsed() > Duration::from_secs(5) {
                 last_log_time = Instant::now();
                 info!(
-                    diffs_written = diffs_written,
-                    summaries_written = summaries_written,
+                    diffs_written,
+                    summaries_written,
                     summaries_count = state_summaries_dag.summaries_count(),
                     "Hot states migration in progress"
                 );
@@ -239,8 +241,8 @@ pub fn upgrade_to_v24<T: BeaconChainTypes>(
     // TODO(hdiff): Should run hot DB compaction after deleting potentially a lot of states. Or should wait
     // for the next finality event?
     info!(
-        diffs_written = diffs_written,
-        summaries_written = summaries_written,
+        diffs_written,
+        summaries_written,
         summaries_count = state_summaries_dag.summaries_count(),
         "Hot states migration complete"
     );
@@ -249,10 +251,100 @@ pub fn upgrade_to_v24<T: BeaconChainTypes>(
 }
 
 pub fn downgrade_from_v24<T: BeaconChainTypes>(
-    _db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+    db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
 ) -> Result<Vec<KeyValueStoreOp>, Error> {
-    // TODO(hdiff): proper error
-    panic!("downgrade not supported");
+    let state_summaries = db
+        .load_hot_state_summaries()?
+        .into_iter()
+        .map(|(state_root, summary)| (state_root, summary.into()))
+        .collect::<Vec<(Hash256, DAGStateSummary)>>();
+
+    info!(
+        summaries_count = state_summaries.len(),
+        "DB downgrade of v24 state summaries started"
+    );
+
+    let state_summaries_dag = StateSummariesDAG::new(state_summaries)
+        .map_err(|e| Error::MigrationError(format!("Error on new StateSumariesDAG {e:?}")))?;
+
+    let mut migrate_ops = vec![];
+    let mut states_written = 0;
+    let mut summaries_written = 0;
+    let mut last_log_time = Instant::now();
+
+    // TODO(tree-states): What about the anchor_slot? Is it safe to run the prior version of
+    // Lighthouse with an a higher anchor_slot than expected?
+
+    for (state_root, summary) in state_summaries_dag
+        .summaries_by_slot_ascending()
+        .into_iter()
+        .flat_map(|(_, summaries)| summaries)
+    {
+        // If boundary state persist
+        if summary.slot % T::EthSpec::slots_per_epoch() == 0 {
+            let (state, _) = db
+                .load_hot_state(&state_root)?
+                .ok_or(Error::MissingState(state_root))?;
+
+            // Immediately commit the state. Otherwise we will OOM and it's stored in a different
+            // column. So if the migration crashes we just get extra harmless junk in the DB.
+            let mut state_write_ops = vec![];
+            store_full_state_v22(&state_root, &state, &mut state_write_ops)?;
+            db.hot_db.do_atomically(state_write_ops)?;
+            states_written += 1;
+        }
+
+        // Persist old summary
+        let epoch_boundary_state_slot = summary.slot - summary.slot % T::EthSpec::slots_per_epoch();
+        let old_summary = HotStateSummaryV22 {
+            slot: summary.slot,
+            latest_block_root: summary.latest_block_root,
+            epoch_boundary_state_root: state_summaries_dag
+                .ancestor_state_root_at_slot(state_root, epoch_boundary_state_slot)
+                .map_err(|e| {
+                    Error::MigrationError(format!(
+                        "error computing ancestor_state_root_at_slot({state_root:?}, {epoch_boundary_state_slot}) {e:?}"
+                    ))
+                })?,
+        };
+        migrate_ops.push(KeyValueStoreOp::PutKeyValue(
+            DBColumn::BeaconStateSummary,
+            state_root.as_slice().to_vec(),
+            old_summary.as_ssz_bytes(),
+        ));
+        summaries_written += 1;
+
+        // Delete existing data
+        for db_column in [
+            DBColumn::BeaconStateHotSummary,
+            DBColumn::BeaconStateHotDiff,
+            DBColumn::BeaconStateHotSnapshot,
+        ] {
+            migrate_ops.push(KeyValueStoreOp::DeleteKey(
+                db_column,
+                state_root.as_slice().to_vec(),
+            ));
+        }
+
+        if last_log_time.elapsed() > Duration::from_secs(5) {
+            last_log_time = Instant::now();
+            info!(
+                states_written,
+                summaries_written,
+                summaries_count = state_summaries_dag.summaries_count(),
+                "DB downgrade of v24 state summaries in progress"
+            );
+        }
+    }
+
+    info!(
+        states_written,
+        summaries_written,
+        summaries_count = state_summaries_dag.summaries_count(),
+        "DB downgrade of v24 state summaries completed"
+    );
+
+    Ok(migrate_ops)
 }
 
 fn new_dag<T: BeaconChainTypes>(

beacon_node/beacon_chain/src/summaries_dag.rs

@@ -3,6 +3,7 @@ use std::{
     cmp::Ordering,
     collections::{btree_map::Entry, BTreeMap, HashMap},
 };
+use store::HotStateSummary;
 use types::{Hash256, Slot};
 
 #[derive(Debug, Clone, Copy)]
@@ -64,6 +65,12 @@ pub enum Error {
         root_state_root: Hash256,
         root_state_slot: Slot,
     },
+    CircularAncestorChain {
+        state_root: Hash256,
+        previous_state_root: Hash256,
+        slot: Slot,
+        last_slot: Slot,
+    },
 }
 
 impl StateSummariesDAG {
@@ -335,10 +342,24 @@ impl StateSummariesDAG {
         }
 
         let mut ancestors = vec![];
+        let mut last_slot = None;
         loop {
             if let Some(summary) = self.state_summaries_by_state_root.get(&state_root) {
+                // Detect cycles, including the case where `previous_state_root == state_root`.
+                if let Some(last_slot) = last_slot {
+                    if summary.slot >= last_slot {
+                        return Err(Error::CircularAncestorChain {
+                            state_root,
+                            previous_state_root: summary.previous_state_root,
+                            slot: summary.slot,
+                            last_slot,
+                        });
+                    }
+                }
+
                 ancestors.push((state_root, summary.slot));
-                state_root = summary.previous_state_root
+                last_slot = Some(summary.slot);
+                state_root = summary.previous_state_root;
             } else {
                 return Ok(ancestors);
             }
@@ -360,6 +381,17 @@ impl StateSummariesDAG {
     }
 }
 
+impl From<HotStateSummary> for DAGStateSummary {
+    fn from(value: HotStateSummary) -> Self {
+        Self {
+            slot: value.slot,
+            latest_block_root: value.latest_block_root,
+            latest_block_slot: value.latest_block_slot,
+            previous_state_root: value.previous_state_root,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{DAGStateSummaryV22, Error, StateSummariesDAG};