fix: Incremental manifest computation with LSMT

Author: schneiderstefan

rs/state_manager/src/lib.rs

@@ -2299,8 +2299,14 @@ impl StateManagerImpl {
                 })
                 .and_then(|(base_manifest, base_height)| {
                     if let Ok(checkpoint_layout) = self.state_layout.checkpoint(base_height) {
+                        // If `lsmt_storage` is enabled, then `dirty_pages` is not needed, as each file is either completely
+                        // new, or identical (same inode) to before.
+                        let dirty_pages = match self.lsmt_storage {
+                            FlagStatus::Enabled => Vec::new(),
+                            FlagStatus::Disabled => get_dirty_pages(state),
+                        };
                         Some(PreviousCheckpointInfo {
-                            dirty_pages: get_dirty_pages(state),
+                            dirty_pages,
                             base_manifest,
                             base_height,
                             checkpoint_layout,
@@ -2418,6 +2424,7 @@ impl StateManagerImpl {
                         target_height: height,
                         dirty_memory_pages: dirty_pages,
                         base_checkpoint: checkpoint_layout,
+                        lsmt_storage: self.lsmt_storage,
                     }
                 },
             )

rs/state_manager/src/manifest.rs

@@ -16,6 +16,7 @@ use crate::{
 };
 use bit_vec::BitVec;
 use hash::{chunk_hasher, file_hasher, manifest_hasher, ManifestHash};
+use ic_config::flag_status::FlagStatus;
 use ic_crypto_sha2::Sha256;
 use ic_logger::{error, fatal, replica_logger::no_op_logger, ReplicaLogger};
 use ic_metrics::MetricsRegistry;
@@ -235,6 +236,7 @@ pub struct ManifestDelta {
     /// state at `base_height`.
     pub(crate) dirty_memory_pages: DirtyPages,
     pub(crate) base_checkpoint: CheckpointLayout<ReadOnly>,
+    pub(crate) lsmt_storage: FlagStatus,
 }
 
 /// Groups small files into larger chunks.
@@ -790,26 +792,37 @@ fn dirty_pages_to_dirty_chunks(
     );
 
     let mut dirty_chunks: BTreeMap<PathBuf, BitVec> = Default::default();
-    for dirty_page in &manifest_delta.dirty_memory_pages {
-        if dirty_page.height != manifest_delta.base_height {
-            continue;
-        }
 
-        let path = dirty_page.page_type.path(checkpoint);
+    // If `lsmt_storage` is enabled, we shouldn't have populated `dirty_memory_pages` in the first place.
+    debug_assert!(
+        manifest_delta.lsmt_storage == FlagStatus::Disabled
+            || manifest_delta.dirty_memory_pages.is_empty()
+    );
+
+    // Any information on dirty pages is not relevant to what files might have changed with `lsmt_storage`
+    // enabled.
+    if manifest_delta.lsmt_storage == FlagStatus::Disabled {
+        for dirty_page in &manifest_delta.dirty_memory_pages {
+            if dirty_page.height != manifest_delta.base_height {
+                continue;
+            }
 
-        if let Ok(path) = path {
-            let relative_path = path
-                .strip_prefix(checkpoint.raw_path())
-                .expect("failed to strip path prefix");
+            let path = dirty_page.page_type.path(checkpoint);
 
-            if let Some(chunks_bitmap) = dirty_chunks_of_file(
-                relative_path,
-                &dirty_page.page_delta_indices,
-                files,
-                max_chunk_size,
-                &manifest_delta.base_manifest,
-            ) {
-                dirty_chunks.insert(relative_path.to_path_buf(), chunks_bitmap);
+            if let Ok(path) = path {
+                let relative_path = path
+                    .strip_prefix(checkpoint.raw_path())
+                    .expect("failed to strip path prefix");
+
+                if let Some(chunks_bitmap) = dirty_chunks_of_file(
+                    relative_path,
+                    &dirty_page.page_delta_indices,
+                    files,
+                    max_chunk_size,
+                    &manifest_delta.base_manifest,
+                ) {
+                    dirty_chunks.insert(relative_path.to_path_buf(), chunks_bitmap);
+                }
             }
         }
     }
@@ -847,6 +860,7 @@ fn dirty_pages_to_dirty_chunks(
             let num_chunks = count_chunks(*size_bytes, max_chunk_size);
             let chunks_bitmap = BitVec::from_elem(num_chunks, false);
             let _prev_chunk = dirty_chunks.insert(path.clone(), chunks_bitmap);
+            // Check that for hardlinked files there are no dirty pages.
             debug_assert!(_prev_chunk.is_none());
         }
     }

rs/state_manager/src/manifest/tests/computation.rs

@@ -8,6 +8,7 @@ use crate::manifest::{
     MAX_FILE_SIZE_TO_GROUP,
 };
 
+use ic_config::flag_status::FlagStatus;
 use ic_crypto_sha2::Sha256;
 use ic_logger::replica_logger::no_op_logger;
 use ic_metrics::MetricsRegistry;
@@ -1089,6 +1090,7 @@ fn test_dirty_pages_to_dirty_chunks_accounts_for_hardlinks() {
                 Height::new(0),
             )
             .unwrap(),
+            lsmt_storage: FlagStatus::Enabled,
         },
         &CheckpointLayout::new_untracked(checkpoint1.to_path_buf(), Height::new(1)).unwrap(),
         &[

rs/state_manager/tests/common/mod.rs

@@ -1,6 +1,9 @@
 use assert_matches::assert_matches;
 use ic_base_types::NumSeconds;
-use ic_config::state_manager::Config;
+use ic_config::{
+    flag_status::FlagStatus,
+    state_manager::{lsmt_storage_default, Config},
+};
 use ic_interfaces::{
     certification::{CertificationPermanentError, Verifier, VerifierError},
     validation::ValidationResult,
@@ -652,12 +655,18 @@ where
     });
 }
 
-pub fn state_manager_restart_test_with_metrics<Test>(test: Test)
+pub fn state_manager_restart_test_with_lsmt<Test>(lsmt_storage: FlagStatus, test: Test)
 where
     Test: FnOnce(
         &MetricsRegistry,
         StateManagerImpl,
-        Box<dyn Fn(StateManagerImpl, Option<Height>) -> (MetricsRegistry, StateManagerImpl)>,
+        Box<
+            dyn Fn(
+                StateManagerImpl,
+                Option<Height>,
+                FlagStatus,
+            ) -> (MetricsRegistry, StateManagerImpl),
+        >,
     ),
 {
     let tmp = tmpdir("sm");
@@ -666,9 +675,12 @@ where
     let verifier: Arc<dyn Verifier> = Arc::new(FakeVerifier::new());
 
     with_test_replica_logger(|log| {
-        let make_state_manager = move |starting_height| {
+        let make_state_manager = move |starting_height, lsmt_storage| {
             let metrics_registry = MetricsRegistry::new();
 
+            let mut config = config.clone();
+            config.lsmt_storage = lsmt_storage;
+
             let state_manager = StateManagerImpl::new(
                 Arc::clone(&verifier),
                 own_subnet,
@@ -683,17 +695,36 @@ where
             (metrics_registry, state_manager)
         };
 
-        let (metrics_registry, state_manager) = make_state_manager(None);
+        let (metrics_registry, state_manager) = make_state_manager(None, lsmt_storage);
 
-        let restart_fn = Box::new(move |state_manager, starting_height| {
+        let restart_fn = Box::new(move |state_manager, starting_height, lsmt_storage| {
             drop(state_manager);
-            make_state_manager(starting_height)
+            make_state_manager(starting_height, lsmt_storage)
         });
 
         test(&metrics_registry, state_manager, restart_fn);
     });
 }
 
+pub fn state_manager_restart_test_with_metrics<Test>(test: Test)
+where
+    Test: FnOnce(
+        &MetricsRegistry,
+        StateManagerImpl,
+        Box<dyn Fn(StateManagerImpl, Option<Height>) -> (MetricsRegistry, StateManagerImpl)>,
+    ),
+{
+    state_manager_restart_test_with_lsmt(
+        lsmt_storage_default(),
+        |metrics, state_manager, restart_fn| {
+            let restart_fn_simplified = Box::new(move |state_manager, starting_height| {
+                restart_fn(state_manager, starting_height, lsmt_storage_default())
+            });
+            test(metrics, state_manager, restart_fn_simplified);
+        },
+    );
+}
+
 pub fn state_manager_restart_test<Test>(test: Test)
 where
     Test:

rs/state_manager/tests/state_manager.rs

@@ -4626,6 +4626,124 @@ fn checkpoints_are_readonly() {
     });
 }
 
+#[test]
+fn can_upgrade_to_lsmt() {
+    fn vmemory0_exists(state_manager: &StateManagerImpl, height: Height) -> bool {
+        state_manager
+            .state_layout()
+            .checkpoint(height)
+            .unwrap()
+            .canister(&canister_test_id(1))
+            .unwrap()
+            .vmemory_0()
+            .exists()
+    }
+
+    fn num_overlays(state_manager: &StateManagerImpl, height: Height) -> usize {
+        state_manager
+            .state_layout()
+            .checkpoint(height)
+            .unwrap()
+            .canister(&canister_test_id(1))
+            .unwrap()
+            .vmemory_0_overlays()
+            .unwrap()
+            .len()
+    }
+
+    state_manager_restart_test_with_lsmt(
+        FlagStatus::Disabled,
+        |metrics, state_manager, restart_fn| {
+            let (_height, mut state) = state_manager.take_tip();
+
+            insert_dummy_canister(&mut state, canister_test_id(1));
+            let canister_state = state.canister_state_mut(&canister_test_id(1)).unwrap();
+            let execution_state = canister_state.execution_state.as_mut().unwrap();
+
+            const NEW_WASM_PAGE: u64 = 100;
+
+            fn verify_page_map(page_map: &PageMap, value: u8) {
+                assert_eq!(page_map.get_page(PageIndex::new(0)), &[1u8; PAGE_SIZE]);
+                for i in 1..NEW_WASM_PAGE {
+                    assert_eq!(page_map.get_page(PageIndex::new(i)), &[0u8; PAGE_SIZE]);
+                }
+                assert_eq!(
+                    page_map.get_page(PageIndex::new(NEW_WASM_PAGE)),
+                    &[value; PAGE_SIZE]
+                );
+            }
+
+            execution_state.wasm_memory.page_map.update(&[
+                (PageIndex::new(0), &[1u8; PAGE_SIZE]),
+                (PageIndex::new(NEW_WASM_PAGE), &[2u8; PAGE_SIZE]),
+            ]);
+
+            verify_page_map(&execution_state.wasm_memory.page_map, 2u8);
+
+            state_manager.commit_and_certify(state, height(1), CertificationScope::Full);
+            wait_for_checkpoint(&state_manager, height(1));
+
+            assert!(vmemory0_exists(&state_manager, height(1)));
+            assert_eq!(num_overlays(&state_manager, height(1)), 0);
+
+            let (_height, mut state) = state_manager.take_tip();
+            let canister_state = state.canister_state_mut(&canister_test_id(1)).unwrap();
+            let execution_state = canister_state.execution_state.as_mut().unwrap();
+            verify_page_map(&execution_state.wasm_memory.page_map, 2u8);
+
+            execution_state.wasm_memory.page_map.update(&[
+                (PageIndex::new(0), &[1u8; PAGE_SIZE]),
+                (PageIndex::new(NEW_WASM_PAGE), &[3u8; PAGE_SIZE]),
+            ]);
+
+            state_manager.commit_and_certify(state, height(2), CertificationScope::Full);
+            wait_for_checkpoint(&state_manager, height(2));
+
+            let (_height, mut state) = state_manager.take_tip();
+            let canister_state = state.canister_state_mut(&canister_test_id(1)).unwrap();
+            let execution_state = canister_state.execution_state.as_mut().unwrap();
+            verify_page_map(&execution_state.wasm_memory.page_map, 3u8);
+
+            assert!(vmemory0_exists(&state_manager, height(2)));
+            assert_eq!(num_overlays(&state_manager, height(2)), 0);
+
+            assert_error_counters(metrics);
+
+            // restart the state_manager
+            let (metrics, state_manager) = restart_fn(state_manager, None, FlagStatus::Enabled);
+
+            let (_height, mut state) = state_manager.take_tip();
+            let canister_state = state.canister_state_mut(&canister_test_id(1)).unwrap();
+            let execution_state = canister_state.execution_state.as_mut().unwrap();
+            verify_page_map(&execution_state.wasm_memory.page_map, 3u8);
+
+            assert!(vmemory0_exists(&state_manager, height(2)));
+            assert_eq!(num_overlays(&state_manager, height(2)), 0);
+
+            execution_state.wasm_memory.page_map.update(&[
+                (PageIndex::new(0), &[1u8; PAGE_SIZE]),
+                (PageIndex::new(NEW_WASM_PAGE), &[4u8; PAGE_SIZE]),
+            ]);
+
+            state_manager.commit_and_certify(state, height(3), CertificationScope::Full);
+            wait_for_checkpoint(&state_manager, height(3));
+
+            let (_height, mut state) = state_manager.take_tip();
+            let canister_state = state.canister_state_mut(&canister_test_id(1)).unwrap();
+            let execution_state = canister_state.execution_state.as_mut().unwrap();
+            verify_page_map(&execution_state.wasm_memory.page_map, 4u8);
+
+            assert!(vmemory0_exists(&state_manager, height(3)));
+            assert_eq!(num_overlays(&state_manager, height(3)), 1);
+
+            state_manager.commit_and_certify(state, height(4), CertificationScope::Full);
+            wait_for_checkpoint(&state_manager, height(4));
+
+            assert_error_counters(&metrics);
+        },
+    );
+}
+
 proptest! {
     #[test]
     fn stream_store_encode_decode(stream in arb_stream(0, 10, 0, 10), size_limit in 0..20usize) {