feat(upgrader): new stable memory layout

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<DefaultMemoryImpl>;
 type StableMap<K, V> = StableBTreeMap<K, V, Memory>;
 type StableValue<T> = 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<MemoryManager<DefaultMemoryImpl>> =
-        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<StableValue<StorablePrincipal>> = RefCell::new(
+    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =
+        RefCell::new(MemoryManager::init_with_bucket_size(DefaultMemoryImpl::default(), STABLE_MEMORY_BUCKET_SIZE));
+    static STATE: RefCell<StableValue<State>> = 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<u8>);
+    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<RawBytes> =
+        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::<Principal>(&target_canister_bytes.0) {
+        let old_disaster_recovery: StableValue<DisasterRecovery> = 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<Timestamp, LogEntry, Memory> =
+            StableBTreeMap::init(old_memory_manager.get(MemoryId::new(OLD_MEMORY_ID_LOGS)));
+        let logs: BTreeMap<Timestamp, LogEntry> = 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! {

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,33 +10,21 @@ 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,
 };
 
 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<StableValue<DisasterRecovery>> = 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<DisasterRecoveryService> =
         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);
     }
 }
 

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<BTreeMap<Timestamp, LogEntry, Memory>> = RefCell::new(
-        BTreeMap::init(
+    static STORAGE: RefCell<StableBTreeMap<Timestamp, LogEntry, Memory>> = 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<Timestamp, LogEntry>) {
+    STORAGE.with(|storage| {
+        for (timestamp, log) in logs {
+            storage.borrow_mut().insert(timestamp, log);
+        }
+    });
+}
+
 lazy_static! {
     pub static ref LOGGER_SERVICE: Arc<LoggerService> = Arc::new(LoggerService::default());
 }

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);