From 62b3a59300e9ba4c953f340b7ef4203406ba3e27 Mon Sep 17 00:00:00 2001 From: mraszyk <31483726+mraszyk@users.noreply.github.com> Date: Tue, 14 Jan 2025 11:25:08 +0100 Subject: [PATCH] feat(upgrader): new stable memory layout (#474) This PR introduces a new stable memory layout for the upgrader: - the bucket size is reduced from the default 8MiB to 1MiB; - the target canister id is bundled with the disaster recovery state (and a new stable memory version field) into a single stable memory region. The above two changes reduce the default upgrader stable memory footprint from 3 x 8MiB down to 2 x 1MiB. --- core/upgrader/impl/src/lib.rs | 156 +++++++++++++++--- .../impl/src/services/disaster_recovery.rs | 20 +-- core/upgrader/impl/src/services/logger.rs | 17 +- .../src/upgrader_migration_tests.rs | 4 +- 4 files changed, 151 insertions(+), 46 deletions(-) diff --git a/core/upgrader/impl/src/lib.rs b/core/upgrader/impl/src/lib.rs index c31bd72db..5403e2b56 100644 --- a/core/upgrader/impl/src/lib.rs +++ b/core/upgrader/impl/src/lib.rs @@ -1,16 +1,23 @@ +use crate::model::{DisasterRecovery, LogEntry}; +use crate::services::insert_logs; use crate::upgrade::{ CheckController, Upgrade, Upgrader, WithAuthorization, WithBackground, WithLogs, WithStart, WithStop, }; use candid::Principal; -use ic_cdk::{api::management_canister::main::CanisterInstallMode, init, update}; +use ic_cdk::api::stable::{stable_size, stable_write}; +use ic_cdk::{ + api::management_canister::main::CanisterInstallMode, init, post_upgrade, trap, update, +}; use ic_stable_structures::{ memory_manager::{MemoryId, MemoryManager, VirtualMemory}, - DefaultMemoryImpl, StableBTreeMap, + storable::Bound, + DefaultMemoryImpl, StableBTreeMap, Storable, }; use lazy_static::lazy_static; use orbit_essentials::storable; -use std::{cell::RefCell, sync::Arc}; +use orbit_essentials::types::Timestamp; +use std::{borrow::Cow, cell::RefCell, collections::BTreeMap, sync::Arc}; use upgrade::{UpgradeError, UpgradeParams}; use upgrader_api::{InitArg, TriggerUpgradeError}; @@ -30,41 +37,142 @@ type Memory = VirtualMemory; type StableMap = StableBTreeMap; type StableValue = StableMap<(), T>; -const MEMORY_ID_TARGET_CANISTER_ID: u8 = 0; -const MEMORY_ID_DISASTER_RECOVERY: u8 = 1; -const MEMORY_ID_LOGS: u8 = 4; +/// Represents one mebibyte. +pub const MIB: u32 = 1 << 20; -thread_local! { - static MEMORY_MANAGER: RefCell> = - RefCell::new(MemoryManager::init(DefaultMemoryImpl::default())); -} +/// Canisters use 64KiB pages for Wasm memory, more details in the PR that introduced this constant: +/// - https://github.com/WebAssembly/design/pull/442#issuecomment-153203031 +pub const WASM_PAGE_SIZE: u32 = 65536; -#[storable] -pub struct StorablePrincipal(Principal); +/// The size of the stable memory bucket in WASM pages. +/// +/// We use a bucket size of 1MiB to ensure that the default memory allocated to the canister is as small as possible, +/// this is due to the fact that this cansiter uses several MemoryIds to manage the stable memory similarly to to how +/// a database arranges data per table. +/// +/// Currently a bucket size of 1MiB limits the canister to 32GiB of stable memory, which is more than enough for the +/// current use case, however, if the canister needs more memory in the future, `ic-stable-structures` will need to be +/// updated to support storing more buckets in a backwards compatible way. +pub const STABLE_MEMORY_BUCKET_SIZE: u16 = (MIB / WASM_PAGE_SIZE) as u16; + +/// Current version of stable memory layout. +pub const STABLE_MEMORY_VERSION: u32 = 1; + +const MEMORY_ID_STATE: u8 = 0; +const MEMORY_ID_LOGS: u8 = 1; thread_local! { - static TARGET_CANISTER_ID: RefCell> = RefCell::new( + static MEMORY_MANAGER: RefCell> = + RefCell::new(MemoryManager::init_with_bucket_size(DefaultMemoryImpl::default(), STABLE_MEMORY_BUCKET_SIZE)); + static STATE: RefCell> = RefCell::new( StableValue::init( - MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_TARGET_CANISTER_ID))), + MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_STATE))), ) ); } +#[storable] +struct State { + target_canister: Principal, + disaster_recovery: DisasterRecovery, + stable_memory_version: u32, +} + +impl Default for State { + fn default() -> Self { + Self { + target_canister: Principal::anonymous(), + disaster_recovery: Default::default(), + stable_memory_version: STABLE_MEMORY_VERSION, + } + } +} + +fn get_state() -> State { + STATE.with(|storage| storage.borrow().get(&()).unwrap_or_default()) +} + +fn set_state(state: State) { + STATE.with(|storage| storage.borrow_mut().insert((), state)); +} + pub fn get_target_canister() -> Principal { - TARGET_CANISTER_ID.with(|id| { - id.borrow() - .get(&()) - .map(|id| id.0) - .unwrap_or(Principal::anonymous()) - }) + get_state().target_canister +} + +fn set_target_canister(target_canister: Principal) { + let mut state = get_state(); + state.target_canister = target_canister; + set_state(state); +} + +pub fn get_disaster_recovery() -> DisasterRecovery { + get_state().disaster_recovery +} + +pub fn set_disaster_recovery(value: DisasterRecovery) { + let mut state = get_state(); + state.disaster_recovery = value; + set_state(state); } #[init] fn init_fn(InitArg { target_canister }: InitArg) { - TARGET_CANISTER_ID.with(|id| { - let mut id = id.borrow_mut(); - id.insert((), StorablePrincipal(target_canister)); - }); + set_target_canister(target_canister); +} + +#[post_upgrade] +fn post_upgrade() { + pub struct RawBytes(pub Vec); + impl Storable for RawBytes { + fn to_bytes(&self) -> Cow<[u8]> { + trap("RawBytes should never be serialized") + } + + fn from_bytes(bytes: Cow<[u8]>) -> Self { + Self(bytes.to_vec()) + } + + const BOUND: Bound = Bound::Unbounded; + } + + const OLD_MEMORY_ID_TARGET_CANISTER_ID: u8 = 0; + const OLD_MEMORY_ID_DISASTER_RECOVERY: u8 = 1; + const OLD_MEMORY_ID_LOGS: u8 = 4; + + let old_memory_manager = MemoryManager::init(DefaultMemoryImpl::default()); + + // determine stable memory layout by trying to parse the target canister from memory with OLD_MEMORY_ID_TARGET_CANISTER_ID + let old_target_canister_bytes: StableValue = + StableValue::init(old_memory_manager.get(MemoryId::new(OLD_MEMORY_ID_TARGET_CANISTER_ID))); + let target_canister_bytes = old_target_canister_bytes + .get(&()) + .unwrap_or_else(|| trap("Could not determine stable memory layout.")); + // if a principal can be parsed out of memory with OLD_MEMORY_ID_TARGET_CANISTER_ID + // then we need to perform stable memory migration + if let Ok(target_canister) = serde_cbor::from_slice::(&target_canister_bytes.0) { + let old_disaster_recovery: StableValue = StableValue::init( + old_memory_manager.get(MemoryId::new(OLD_MEMORY_ID_DISASTER_RECOVERY)), + ); + let disaster_recovery: DisasterRecovery = + old_disaster_recovery.get(&()).unwrap_or_default(); + + let old_logs: StableBTreeMap = + StableBTreeMap::init(old_memory_manager.get(MemoryId::new(OLD_MEMORY_ID_LOGS))); + let logs: BTreeMap = old_logs.iter().collect(); + + // clear the stable memory + let stable_memory_size_bytes = stable_size() * (WASM_PAGE_SIZE as u64); + stable_write(0, &vec![0; stable_memory_size_bytes as usize]); + + let state = State { + target_canister, + disaster_recovery, + stable_memory_version: STABLE_MEMORY_VERSION, + }; + set_state(state); + insert_logs(logs); + } } lazy_static! { diff --git a/core/upgrader/impl/src/services/disaster_recovery.rs b/core/upgrader/impl/src/services/disaster_recovery.rs index dd3ee7073..6a97f763e 100644 --- a/core/upgrader/impl/src/services/disaster_recovery.rs +++ b/core/upgrader/impl/src/services/disaster_recovery.rs @@ -1,7 +1,7 @@ use super::{InstallCanister, LoggerService, INSTALL_CANISTER}; use crate::{ errors::UpgraderApiError, - get_target_canister, + get_disaster_recovery, get_target_canister, model::{ Account, AdminUser, Asset, DisasterRecovery, DisasterRecoveryCommittee, DisasterRecoveryInProgressLog, DisasterRecoveryResultLog, DisasterRecoveryStartLog, @@ -10,16 +10,14 @@ use crate::{ SetAccountsLog, SetCommitteeLog, StationRecoveryRequest, }, services::LOGGER_SERVICE, + set_disaster_recovery, upgrader_ic_cdk::{api::time, spawn}, - StableValue, MEMORY_ID_DISASTER_RECOVERY, MEMORY_MANAGER, }; use candid::Principal; -use ic_stable_structures::memory_manager::MemoryId; use lazy_static::lazy_static; use orbit_essentials::{api::ServiceResult, utils::sha256_hash}; use std::{ - cell::RefCell, collections::{HashMap, HashSet}, sync::Arc, }; @@ -27,16 +25,6 @@ use std::{ pub const DISASTER_RECOVERY_REQUEST_EXPIRATION_NS: u64 = 60 * 60 * 24 * 7 * 1_000_000_000; // 1 week pub const DISASTER_RECOVERY_IN_PROGESS_EXPIRATION_NS: u64 = 60 * 60 * 1_000_000_000; // 1 hour -thread_local! { - - static STORAGE: RefCell> = RefCell::new( - StableValue::init( - MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_DISASTER_RECOVERY))), - ) - ); - -} - lazy_static! { pub static ref DISASTER_RECOVERY_SERVICE: Arc = Arc::new(DisasterRecoveryService { @@ -81,11 +69,11 @@ pub struct DisasterRecoveryStorage {} impl DisasterRecoveryStorage { pub fn get(&self) -> DisasterRecovery { - STORAGE.with(|storage| storage.borrow().get(&()).unwrap_or_default()) + get_disaster_recovery() } fn set(&self, value: DisasterRecovery) { - STORAGE.with(|storage| storage.borrow_mut().insert((), value)); + set_disaster_recovery(value); } } diff --git a/core/upgrader/impl/src/services/logger.rs b/core/upgrader/impl/src/services/logger.rs index faf6f9c4d..a026378a5 100644 --- a/core/upgrader/impl/src/services/logger.rs +++ b/core/upgrader/impl/src/services/logger.rs @@ -1,6 +1,6 @@ -use std::{cell::RefCell, sync::Arc}; +use std::{cell::RefCell, collections::BTreeMap, sync::Arc}; -use ic_stable_structures::{memory_manager::MemoryId, BTreeMap}; +use ic_stable_structures::{memory_manager::MemoryId, StableBTreeMap}; use lazy_static::lazy_static; use orbit_essentials::types::Timestamp; @@ -14,13 +14,22 @@ pub const DEFAULT_GET_LOGS_LIMIT: u64 = 10; pub const MAX_LOG_ENTRIES: u64 = 25000; thread_local! { - static STORAGE: RefCell> = RefCell::new( - BTreeMap::init( + static STORAGE: RefCell> = RefCell::new( + StableBTreeMap::init( MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_LOGS))), ) ); } +// only use this function for stable memory migration! +pub fn insert_logs(logs: BTreeMap) { + STORAGE.with(|storage| { + for (timestamp, log) in logs { + storage.borrow_mut().insert(timestamp, log); + } + }); +} + lazy_static! { pub static ref LOGGER_SERVICE: Arc = Arc::new(LoggerService::default()); } diff --git a/tests/integration/src/upgrader_migration_tests.rs b/tests/integration/src/upgrader_migration_tests.rs index 34488580a..12937e7cf 100644 --- a/tests/integration/src/upgrader_migration_tests.rs +++ b/tests/integration/src/upgrader_migration_tests.rs @@ -51,8 +51,8 @@ fn upgrade_from_v0(env: &PocketIc, upgrader_id: Principal, station_id: Principal fn upgrade_from_latest(env: &PocketIc, upgrader_id: Principal, station_id: Principal) { let mut canister_memory = env.get_stable_memory(upgrader_id); - // Assert that stable memory size is 5 buckets of 8MiB each + stable structures header (64KiB) for the latest layout. - assert_eq!(canister_memory.len(), 5 * (8 << 20) + (64 << 10)); + // Assert that stable memory size is 21 buckets of 1MiB each + stable structures header (64KiB) for the latest layout. + assert_eq!(canister_memory.len(), 21 * (1 << 20) + (64 << 10)); // This is used to store the stable memory of the canister for future use canister_memory = compress_to_gzip(&canister_memory);