Update tests

Author: dapplion

beacon_node/beacon_chain/tests/store_tests.rs

@@ -31,7 +31,6 @@ use store::{
     BlobInfo, DBColumn, HotColdDB, StoreConfig,
 };
 use tempfile::{tempdir, TempDir};
-use tokio::time::sleep;
 use types::test_utils::{SeedableRng, XorShiftRng};
 use types::*;
 
@@ -119,6 +118,17 @@ fn get_harness_generic(
     harness
 }
 
+fn count_states_descendant_of_block(
+    store: &HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>>,
+    block_root: Hash256,
+) -> usize {
+    let summaries = store.load_hot_state_summaries().unwrap();
+    summaries
+        .iter()
+        .filter(|(_, s)| s.latest_block_root == block_root)
+        .count()
+}
+
 #[tokio::test]
 async fn light_client_bootstrap_test() {
     let spec = test_spec::<E>();
@@ -2136,64 +2146,6 @@ async fn pruning_test(
     check_no_blocks_exist(&harness, stray_blocks.values());
 }
 
-#[tokio::test]
-async fn garbage_collect_temp_states_from_failed_block_on_startup() {
-    let db_path = tempdir().unwrap();
-
-    // Wrap these functions to ensure the variables are dropped before we try to open another
-    // instance of the store.
-    let mut store = {
-        let store = get_store(&db_path);
-        let harness = get_harness(store.clone(), LOW_VALIDATOR_COUNT);
-
-        let slots_per_epoch = E::slots_per_epoch();
-
-        let genesis_state = harness.get_current_state();
-        let block_slot = Slot::new(2 * slots_per_epoch);
-        let ((signed_block, _), state) = harness.make_block(genesis_state, block_slot).await;
-
-        let (mut block, _) = (*signed_block).clone().deconstruct();
-
-        // Mutate the block to make it invalid, and re-sign it.
-        *block.state_root_mut() = Hash256::repeat_byte(0xff);
-        let proposer_index = block.proposer_index() as usize;
-        let block = Arc::new(block.sign(
-            &harness.validator_keypairs[proposer_index].sk,
-            &state.fork(),
-            state.genesis_validators_root(),
-            &harness.spec,
-        ));
-
-        // The block should be rejected, but should store a bunch of temporary states.
-        harness.set_current_slot(block_slot);
-        harness
-            .process_block_result((block, None))
-            .await
-            .unwrap_err();
-
-        assert_eq!(
-            store.iter_temporary_state_roots().count(),
-            block_slot.as_usize() - 1
-        );
-        store
-    };
-
-    // Wait until all the references to the store have been dropped, this helps ensure we can
-    // re-open the store later.
-    loop {
-        store = if let Err(store_arc) = Arc::try_unwrap(store) {
-            sleep(Duration::from_millis(500)).await;
-            store_arc
-        } else {
-            break;
-        }
-    }
-
-    // On startup, the store should garbage collect all the temporary states.
-    let store = get_store(&db_path);
-    assert_eq!(store.iter_temporary_state_roots().count(), 0);
-}
-
 #[tokio::test]
 async fn garbage_collect_temp_states_from_failed_block_on_finalization() {
     let db_path = tempdir().unwrap();
@@ -2208,6 +2160,7 @@ async fn garbage_collect_temp_states_from_failed_block_on_finalization() {
     let ((signed_block, _), state) = harness.make_block(genesis_state, block_slot).await;
 
     let (mut block, _) = (*signed_block).clone().deconstruct();
+    let bad_block_parent_root = block.parent_root();
 
     // Mutate the block to make it invalid, and re-sign it.
     *block.state_root_mut() = Hash256::repeat_byte(0xff);
@@ -2226,9 +2179,11 @@ async fn garbage_collect_temp_states_from_failed_block_on_finalization() {
         .await
         .unwrap_err();
 
+    // The bad block parent root is the genesis block root. There's `block_slot - 1` temporary
+    // states to remove + the genesis state = block_slot.
     assert_eq!(
-        store.iter_temporary_state_roots().count(),
-        block_slot.as_usize() - 1
+        count_states_descendant_of_block(&store, bad_block_parent_root),
+        block_slot.as_usize(),
     );
 
     // Finalize the chain without the block, which should result in pruning of all temporary states.
@@ -2245,8 +2200,12 @@ async fn garbage_collect_temp_states_from_failed_block_on_finalization() {
     // Check that the finalization migration ran.
     assert_ne!(store.get_split_slot(), 0);
 
-    // Check that temporary states have been pruned.
-    assert_eq!(store.iter_temporary_state_roots().count(), 0);
+    // Check that temporary states have been pruned. The genesis block is not a descendant of the
+    // latest finalized checkpoint, so all its states have been pruned from the hot DB, = 0.
+    assert_eq!(
+        count_states_descendant_of_block(&store, bad_block_parent_root),
+        0
+    );
 }
 
 #[tokio::test]

beacon_node/store/src/hot_cold_store.rs

@@ -14,8 +14,8 @@ use crate::metadata::{
 };
 use crate::state_cache::{PutStateOutcome, StateCache};
 use crate::{
-    get_data_column_key, metrics, parse_data_column_key, BlobSidecarListFromRoot, ColumnKeyIter,
-    DBColumn, DatabaseBlock, Error, ItemStore, KeyValueStore, KeyValueStoreOp, StoreItem, StoreOp,
+    get_data_column_key, metrics, parse_data_column_key, BlobSidecarListFromRoot, DBColumn,
+    DatabaseBlock, Error, ItemStore, KeyValueStoreOp, StoreItem, StoreOp,
 };
 use itertools::{process_results, Itertools};
 use lru::LruCache;
@@ -395,6 +395,12 @@ impl<E: EthSpec> HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>> {
         }
         db.store_config()?;
 
+        // TODO(tree-states): Here we can choose to prune advanced states to reclaim disk space. As
+        // it's a foreground task there's no risk of race condition that can corrupt the DB.
+        // Advanced states for invalid blocks that were never written to the DB, or descendants of
+        // heads can be safely pruned at the expense of potentially having to recompute them in the
+        // future. However this would require a new dedicated pruning routine.
+
         // If configured, run a foreground compaction pass.
         if db.config.compact_on_init {
             info!(db.log, "Running foreground compaction");
@@ -404,12 +410,6 @@ impl<E: EthSpec> HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>> {
 
         Ok(db)
     }
-
-    /// Return an iterator over the state roots of all temporary states.
-    pub fn iter_temporary_state_roots(&self) -> ColumnKeyIter<Hash256> {
-        self.hot_db
-            .iter_column_keys::<Hash256>(DBColumn::BeaconStateTemporary)
-    }
 }
 
 impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold> {