Tree states to support per-slot state diffs (#4652)

* Support per slot state diffs

* Store HierarchyConfig on disk. Support storing hdiffs at per slot level.

* Revert HierachyConfig change for testing.

* Add validity check for the hierarchy config when opening the DB.

* Update HDiff tests.

* Fix `get_cold_state` panic when the diff for the slot isn't stored.

* Use slots instead of epochs for storing snapshots in freezer DB.

* Add snapshot buffer to `diff_buffer_cache` instead of loading it from db every time.

* Add `hierarchy-exponents` cli flag to beacon node.

* Add test for `StorageStrategy::ReplayFrom` and ignore a flaky test.

* Drop hierarchy_config in tests for more frequent snapshot and fix an issue where hdiff wasn't stored unless it's a epoch boundary slot.

beacon_node/beacon_chain/tests/store_tests.rs

@@ -21,9 +21,11 @@ use std::collections::HashSet;
 use std::convert::TryInto;
 use std::sync::Arc;
 use std::time::Duration;
+use store::hdiff::HierarchyConfig;
 use store::{
+    config::StoreConfigError,
     iter::{BlockRootsIterator, StateRootsIterator},
-    HotColdDB, LevelDB, StoreConfig,
+    Error as StoreError, HotColdDB, LevelDB, StoreConfig,
 };
 use tempfile::{tempdir, TempDir};
 use types::test_utils::{SeedableRng, XorShiftRng};
@@ -49,13 +51,25 @@ fn get_store_with_spec(
     db_path: &TempDir,
     spec: ChainSpec,
 ) -> Arc<HotColdDB<E, LevelDB<E>, LevelDB<E>>> {
+    let config = StoreConfig {
+        // More frequent snapshots and hdiffs in tests for testing
+        hierarchy_config: HierarchyConfig {
+            exponents: vec![1, 3, 5],
+        },
+        ..Default::default()
+    };
+    try_get_store_with_spec_and_config(db_path, spec, config).expect("disk store should initialize")
+}
+
+fn try_get_store_with_spec_and_config(
+    db_path: &TempDir,
+    spec: ChainSpec,
+    config: StoreConfig,
+) -> Result<Arc<HotColdDB<E, LevelDB<E>, LevelDB<E>>>, StoreError> {
     let hot_path = db_path.path().join("hot_db");
     let cold_path = db_path.path().join("cold_db");
-    let config = StoreConfig::default();
     let log = test_logger();
-
     HotColdDB::open(&hot_path, &cold_path, |_, _, _| Ok(()), config, spec, log)
-        .expect("disk store should initialize")
 }
 
 fn get_harness(
@@ -2481,6 +2495,43 @@ async fn revert_minority_fork_on_resume() {
     assert_eq!(heads.len(), 1);
 }
 
+#[tokio::test]
+#[ignore]
+// FIXME(jimmy): Ignoring this now as the test is flaky :/ It intermittently fails with an IO error
+// "..cold_db/LOCK file held by another process".
+// There seems to be some race condition between dropping the lock file and and re-opening the db.
+// There's a higher chance this test would fail when the entire test suite is run. Maybe it isn't
+// fast enough at dropping the cold_db LOCK file before the test attempts to open it again.
+async fn should_not_initialize_incompatible_store_config() {
+    let validator_count = 16;
+    let spec = MinimalEthSpec::default_spec();
+    let db_path = tempdir().unwrap();
+    let store_config = StoreConfig::default();
+    let store = try_get_store_with_spec_and_config(&db_path, spec.clone(), store_config.clone())
+        .expect("disk store should initialize");
+    let harness = BeaconChainHarness::builder(MinimalEthSpec)
+        .spec(spec.clone())
+        .deterministic_keypairs(validator_count)
+        .fresh_disk_store(store)
+        .build();
+
+    // Resume from disk with a different store config.
+    drop(harness);
+    let different_store_config = StoreConfig {
+        linear_blocks: !store_config.linear_blocks,
+        ..store_config
+    };
+    let maybe_err =
+        try_get_store_with_spec_and_config(&db_path, spec, different_store_config).err();
+
+    assert!(matches!(
+        maybe_err,
+        Some(StoreError::ConfigError(
+            StoreConfigError::IncompatibleStoreConfig { .. }
+        ))
+    ));
+}
+
 // This test checks whether the schema downgrade from the latest version to some minimum supported
 // version is correct. This is the easiest schema test to write without historic versions of
 // Lighthouse on-hand, but has the disadvantage that the min version needs to be adjusted manually

beacon_node/src/cli.rs

@@ -533,6 +533,21 @@ pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
                        [default: 8192 (mainnet) or 64 (minimal)]")
                 .takes_value(true)
         )
+        .arg(
+            Arg::with_name("hierarchy-exponents")
+                .long("hierarchy-exponents")
+                .value_name("EXPONENTS")
+                .help("Specifies the frequency for storing full state snapshots and hierarchical \
+                        diffs in the freezer DB. Accepts a comma-separated list of ascending \
+                        exponents. Each exponent defines an interval for storing diffs to the layer \
+                        above. The last exponent defines the interval for full snapshots. \
+                        For example, a config of '4,8,12' would store a full snapshot every \
+                        4096 (2^12) slots, first-level diffs every 256 (2^8) slots, and second-level \
+                        diffs every 16 (2^4) slots. \
+                        Cannot be changed after initialization. \
+                        [default: 5,9,11,13,16,18,21]")
+                .takes_value(true)
+        )
         .arg(
             Arg::with_name("epochs-per-migration")
                 .long("epochs-per-migration")

beacon_node/src/config.rs

@@ -23,6 +23,7 @@ use std::net::{IpAddr, Ipv4Addr, ToSocketAddrs};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use std::time::Duration;
+use store::hdiff::HierarchyConfig;
 use types::{Checkpoint, Epoch, EthSpec, Hash256, PublicKeyBytes, GRAFFITI_BYTES_LEN};
 
 /// Gets the fully-initialized global client.
@@ -422,6 +423,24 @@ pub fn get_config<E: EthSpec>(
         client_config.store.epochs_per_state_diff = epochs_per_state_diff;
     }
 
+    if let Some(hierarchy_exponents) =
+        clap_utils::parse_optional::<String>(cli_args, "hierarchy-exponents")?
+    {
+        let exponents = hierarchy_exponents
+            .split(',')
+            .map(|s| {
+                s.parse()
+                    .map_err(|e| format!("invalid hierarchy-exponents: {e:?}"))
+            })
+            .collect::<Result<Vec<u8>, _>>()?;
+
+        if exponents.windows(2).any(|w| w[0] >= w[1]) {
+            return Err("hierarchy-exponents must be in ascending order".to_string());
+        }
+
+        client_config.store.hierarchy_config = HierarchyConfig { exponents };
+    }
+
     if let Some(epochs_per_migration) =
         clap_utils::parse_optional(cli_args, "epochs-per-migration")?
     {

beacon_node/store/src/config.rs

@@ -45,16 +45,25 @@ pub struct StoreConfig {
 
 /// Variant of `StoreConfig` that gets written to disk. Contains immutable configuration params.
 #[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
-// FIXME(sproul): schema migration, add hdiff
+// FIXME(sproul): schema migration
 pub struct OnDiskStoreConfig {
     pub linear_blocks: bool,
-    pub linear_restore_points: bool,
+    pub hierarchy_config: HierarchyConfig,
 }
 
 #[derive(Debug, Clone)]
 pub enum StoreConfigError {
-    MismatchedSlotsPerRestorePoint { config: u64, on_disk: u64 },
-    InvalidCompressionLevel { level: i32 },
+    MismatchedSlotsPerRestorePoint {
+        config: u64,
+        on_disk: u64,
+    },
+    InvalidCompressionLevel {
+        level: i32,
+    },
+    IncompatibleStoreConfig {
+        config: OnDiskStoreConfig,
+        on_disk: OnDiskStoreConfig,
+    },
 }
 
 impl Default for StoreConfig {
@@ -80,15 +89,21 @@ impl StoreConfig {
     pub fn as_disk_config(&self) -> OnDiskStoreConfig {
         OnDiskStoreConfig {
             linear_blocks: self.linear_blocks,
-            linear_restore_points: self.linear_restore_points,
+            hierarchy_config: self.hierarchy_config.clone(),
         }
     }
 
     pub fn check_compatibility(
         &self,
-        _on_disk_config: &OnDiskStoreConfig,
+        on_disk_config: &OnDiskStoreConfig,
     ) -> Result<(), StoreConfigError> {
-        // FIXME(sproul): TODO
+        let db_config = self.as_disk_config();
+        if db_config.ne(on_disk_config) {
+            return Err(StoreConfigError::IncompatibleStoreConfig {
+                config: db_config,
+                on_disk: on_disk_config.clone(),
+            });
+        }
         Ok(())
     }
 
@@ -146,3 +161,49 @@ impl StoreItem for OnDiskStoreConfig {
         Ok(Self::from_ssz_bytes(bytes)?)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn check_compatibility_ok() {
+        let store_config = StoreConfig {
+            linear_blocks: true,
+            ..Default::default()
+        };
+        let on_disk_config = OnDiskStoreConfig {
+            linear_blocks: true,
+            hierarchy_config: store_config.hierarchy_config.clone(),
+        };
+        assert!(store_config.check_compatibility(&on_disk_config).is_ok());
+    }
+
+    #[test]
+    fn check_compatibility_linear_blocks_mismatch() {
+        let store_config = StoreConfig {
+            linear_blocks: true,
+            ..Default::default()
+        };
+        let on_disk_config = OnDiskStoreConfig {
+            linear_blocks: false,
+            hierarchy_config: store_config.hierarchy_config.clone(),
+        };
+        assert!(store_config.check_compatibility(&on_disk_config).is_err());
+    }
+
+    #[test]
+    fn check_compatibility_hierarchy_config_incompatible() {
+        let store_config = StoreConfig {
+            linear_blocks: true,
+            ..Default::default()
+        };
+        let on_disk_config = OnDiskStoreConfig {
+            linear_blocks: true,
+            hierarchy_config: HierarchyConfig {
+                exponents: vec![5, 8, 11, 13, 16, 18, 21],
+            },
+        };
+        assert!(store_config.check_compatibility(&on_disk_config).is_err());
+    }
+}

beacon_node/store/src/errors.rs

@@ -41,7 +41,7 @@ pub enum Error {
     },
     MissingStateRoot(Slot),
     MissingState(Hash256),
-    MissingSnapshot(Epoch),
+    MissingSnapshot(Slot),
     MissingDiff(Epoch),
     NoBaseStateFound(Hash256),
     BlockReplayError(BlockReplayError),