Merge branch 'alin/MR-427-state-manager-split' into 'master'

feat: [MR-427] Implement state splitting in State Manager

 

See merge request dfinity-lab/public/ic!12495

Author: alin-at-dfinity

rs/state_manager/BUILD.bazel

@@ -23,6 +23,7 @@ rust_library(
         "//rs/monitoring/logger",
         "//rs/monitoring/metrics",
         "//rs/protobuf",
+        "//rs/registry/routing_table",
         "//rs/registry/subnet_type",
         "//rs/replicated_state",
         "//rs/state_layout",

rs/state_manager/Cargo.toml

@@ -19,6 +19,7 @@ ic-interfaces-state-manager = { path = "../interfaces/state_manager" }
 ic-logger = { path = "../monitoring/logger" }
 ic-metrics = { path = "../monitoring/metrics" }
 ic-protobuf = { path = "../protobuf" }
+ic-registry-routing-table = { path = "../registry/routing_table" }
 ic-registry-subnet-type = { path = "../registry/subnet_type" }
 ic-replicated-state = { path = "../replicated_state" }
 ic-state-layout = { path = "../state_layout" }

rs/state_manager/src/lib.rs

@@ -3,6 +3,7 @@
 pub mod checkpoint;
 pub mod labeled_tree_visitor;
 pub mod manifest;
+pub mod split;
 pub mod state_sync;
 pub mod stream_encoding;
 pub mod tip;

rs/state_manager/src/split.rs

@@ -0,0 +1,147 @@
+//! Prunes a replicated state, as part of a subnet split.
+use crate::{
+    checkpoint::{load_checkpoint, make_checkpoint},
+    tip::spawn_tip_thread,
+    PageAllocatorFileDescriptorImpl, StateManagerMetrics, NUMBER_OF_CHECKPOINT_THREADS,
+};
+
+use ic_config::state_manager::Config;
+use ic_logger::ReplicaLogger;
+use ic_metrics::MetricsRegistry;
+use ic_registry_routing_table::{CanisterIdRanges, RoutingTable};
+use ic_registry_subnet_type::SubnetType;
+use ic_replicated_state::{page_map::PageAllocatorFileDescriptor, ReplicatedState};
+use ic_state_layout::{CheckpointLayout, ReadOnly, StateLayout};
+use ic_types::{malicious_flags::MaliciousFlags, PrincipalId, SubnetId};
+use scoped_threadpool::Pool;
+use std::{path::PathBuf, sync::Arc};
+
+#[cfg(test)]
+mod tests;
+
+/// Loads the latest checkpoint under the given root; splits off the state of
+/// `subnet_id`, hosting the provided canister ID ranges; and writes back the
+/// split state as a new checkpoint, under the same root.
+pub fn split(
+    root: PathBuf,
+    subnet_id: PrincipalId,
+    canister_id_ranges: CanisterIdRanges,
+    metrics_registry: &MetricsRegistry,
+    log: ReplicaLogger,
+) -> Result<(), String> {
+    // Load latest checkpoint under `root`.
+    let config = Config::new(root);
+    let state_layout =
+        StateLayout::try_new(log.clone(), config.state_root, metrics_registry).unwrap();
+
+    // A thread pool to use for reading and writing checkpoints.
+    let mut thread_pool = Pool::new(NUMBER_OF_CHECKPOINT_THREADS);
+
+    // Create the file descriptor factory that is used to create files for PageMaps.
+    let page_delta_path = state_layout.page_deltas();
+    let fd_factory: Arc<dyn PageAllocatorFileDescriptor> = Arc::new(
+        PageAllocatorFileDescriptorImpl::new(page_delta_path, config.file_backed_memory_allocator),
+    );
+
+    let metrics = StateManagerMetrics::new(metrics_registry);
+    let (cp, state) = read_checkpoint(
+        &state_layout,
+        &mut thread_pool,
+        fd_factory.clone(),
+        &metrics,
+    )?;
+
+    // Set up the split.
+    let subnet_id: SubnetId = subnet_id.into();
+    let mut routing_table = RoutingTable::new();
+    routing_table
+        .assign_ranges(canister_id_ranges, subnet_id)
+        .map_err(|e| format!("{:?}", e))?;
+
+    // Split the state.
+    let split_state = state.split(subnet_id, &routing_table);
+
+    // Write the split state as a new checkpoint.
+    write_checkpoint(
+        &split_state,
+        state_layout,
+        &cp,
+        &mut thread_pool,
+        fd_factory,
+        &metrics,
+        log,
+    )
+}
+
+/// Reads the `ReplicatedState` from the latest checkpoint under `state_layout`.
+fn read_checkpoint(
+    state_layout: &StateLayout,
+    thread_pool: &mut Pool,
+    fd_factory: Arc<dyn PageAllocatorFileDescriptor>,
+    metrics: &StateManagerMetrics,
+) -> Result<(CheckpointLayout<ReadOnly>, ReplicatedState), String> {
+    let height = *state_layout
+        .checkpoint_heights()
+        .map_err(|e| e.to_string())?
+        .last()
+        .ok_or(format!(
+            "No checkpoints found at {}",
+            state_layout.raw_path().display()
+        ))?;
+    let cp = state_layout.checkpoint(height).map_err(|e| e.to_string())?;
+
+    let state = load_checkpoint(
+        &cp,
+        SubnetType::Application,
+        &metrics.checkpoint_metrics,
+        Some(thread_pool),
+        fd_factory,
+    )
+    .map_err(|e| {
+        format!(
+            "Failed to load checkpoint at {}: {}",
+            cp.raw_path().display(),
+            e
+        )
+    })?;
+
+    Ok((cp, state))
+}
+
+/// Writes the given `ReplicatedState` into a new checkpoint under
+/// `state_layout`, based off of `old_cp`.
+fn write_checkpoint(
+    state: &ReplicatedState,
+    state_layout: StateLayout,
+    old_cp: &CheckpointLayout<ReadOnly>,
+    thread_pool: &mut Pool,
+    fd_factory: Arc<dyn PageAllocatorFileDescriptor>,
+    metrics: &StateManagerMetrics,
+    log: ReplicaLogger,
+) -> Result<(), String> {
+    let old_height = old_cp.height();
+
+    let mut tip_handler = state_layout.capture_tip_handler();
+    tip_handler
+        .reset_tip_to(&state_layout, old_cp, Some(thread_pool))
+        .map_err(|e| e.to_string())?;
+    let (_tip_thread, tip_channel) = spawn_tip_thread(
+        log,
+        tip_handler,
+        state_layout,
+        metrics.clone(),
+        MaliciousFlags::default(),
+    );
+
+    make_checkpoint(
+        state,
+        old_height.increment(),
+        &tip_channel,
+        &metrics.checkpoint_metrics,
+        thread_pool,
+        fd_factory,
+    )
+    .map_err(|e| format!("Failed to write checkpoint: {}", e))?;
+
+    Ok(())
+}