feat(ICP-Ledger): FI-1433: Implement V1 for ICP ledger - add ability …

…to read from memory manager in post_upgrade (#1020) Co-authored-by: IDX GitHub Automation <[email protected]> Co-authored-by: Mathias Björkqvist <[email protected]> Co-authored-by: mraszyk <[email protected]>
dfinity · Oct 4, 2024 · 6dcfafb · 6dcfafb
1 parent c8d0295
commit 6dcfafb
Show file tree

Hide file tree

Showing 8 changed files with 198 additions and 25 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rs/ledger_suite/icp/ledger/BUILD.bazel b/rs/ledger_suite/icp/ledger/BUILD.bazel
@@ -37,6 +37,7 @@ rust_library(
         "//rs/rust_canisters/dfn_core",
         "//rs/types/base_types",
         "@crate_index//:candid",
+        "@crate_index//:ic-stable-structures",
         "@crate_index//:intmap",
         "@crate_index//:lazy_static",
         "@crate_index//:num-traits",
@@ -78,7 +79,9 @@ LEDGER_CANISTER_DEPS = [
     "//rs/types/base_types",
     "@crate_index//:candid",
     "@crate_index//:ciborium",
+    "@crate_index//:ic-cdk",
     "@crate_index//:ic-metrics-encoder",
+    "@crate_index//:ic-stable-structures",
     "@crate_index//:num-traits",
     "@crate_index//:serde_bytes",
 ]
@@ -113,6 +116,17 @@ rust_canister(
     deps = LEDGER_CANISTER_DEPS,
 )
 
+rust_canister(
+    name = "ledger-canister-wasm-upgrade-to-memory-manager",
+    srcs = ["src/main.rs"],
+    compile_data = LEDGER_CANISTER_DATA,
+    crate_features = ["upgrade-to-memory-manager"],
+    data = LEDGER_CANISTER_DATA,
+    rustc_env = LEDGER_CANISTER_RUSTC_ENV,
+    service_file = "//rs/ledger_suite/icp:ledger.did",
+    deps = LEDGER_CANISTER_DEPS,
+)
+
 rust_test(
     name = "ledger_canister_unit_test",
     compile_data = LEDGER_CANISTER_DATA,
@@ -130,12 +144,14 @@ rust_ic_test(
     srcs = ["tests/tests.rs"],
     data = [
         ":ledger-canister-wasm",
+        ":ledger-canister-wasm-upgrade-to-memory-manager",
         "@mainnet_icp_ledger_canister//file",
     ],
     env = {
         "CARGO_MANIFEST_DIR": "rs/ledger_suite/icp/ledger",
         "ICP_LEDGER_DEPLOYED_VERSION_WASM_PATH": "$(rootpath @mainnet_icp_ledger_canister//file)",
         "LEDGER_CANISTER_WASM_PATH": "$(rootpath :ledger-canister-wasm)",
+        "LEDGER_CANISTER_UPGRADE_TO_MEMORY_MANAGER_WASM_PATH": "$(rootpath :ledger-canister-wasm-upgrade-to-memory-manager)",
     },
     flaky = True,
     deps = [

diff --git a/rs/ledger_suite/icp/ledger/Cargo.toml b/rs/ledger_suite/icp/ledger/Cargo.toml
@@ -21,12 +21,14 @@ dfn_http_metrics = { path = "../../../rust_canisters/dfn_http_metrics" }
 dfn_protobuf = { path = "../../../rust_canisters/dfn_protobuf" }
 ic-base-types = { path = "../../../types/base_types" }
 ic-canister-log = { path = "../../../rust_canisters/canister_log" }
+ic-cdk = { workspace = true }
 ic-limits = { path = "../../../limits" }
 ic-icrc1 = { path = "../../icrc1" }
 ic-ledger-canister-core = { path = "../../common/ledger_canister_core" }
 ic-ledger-core = { path = "../../common/ledger_core" }
 ic-ledger-hash-of = { path = "../../../../packages/ic-ledger-hash-of" }
 ic-metrics-encoder = "1"
+ic-stable-structures = { workspace = true }
 icp-ledger = { path = "../" }
 icrc-ledger-types = { path = "../../../../packages/icrc-ledger-types" }
 intmap = { version = "1.1.0", features = ["serde"] }
@@ -49,3 +51,4 @@ ic-test-utilities-load-wasm = { path = "../../../test_utilities/load_wasm" }
 
 [features]
 notify-method = []
+upgrade-to-memory-manager = []
diff --git a/rs/ledger_suite/icp/ledger/src/lib.rs b/rs/ledger_suite/icp/ledger/src/lib.rs
@@ -11,6 +11,8 @@ use ic_ledger_core::{
 };
 use ic_ledger_core::{block::BlockIndex, tokens::Tokens};
 use ic_ledger_hash_of::HashOf;
+use ic_stable_structures::memory_manager::{MemoryId, MemoryManager, VirtualMemory};
+use ic_stable_structures::DefaultMemoryImpl;
 use icp_ledger::{
     AccountIdentifier, Block, FeatureFlags, LedgerAllowances, LedgerBalances, Memo, Operation,
     PaymentError, Transaction, TransferError, TransferFee, UpgradeArgs, DEFAULT_TRANSFER_FEE,
@@ -20,6 +22,7 @@ use intmap::IntMap;
 use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
 use std::borrow::Cow;
+use std::cell::RefCell;
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::sync::RwLock;
 use std::time::Duration;
@@ -61,6 +64,18 @@ fn unknown_token() -> String {
     "???".to_string()
 }
 
+const UPGRADES_MEMORY_ID: MemoryId = MemoryId::new(0);
+
+thread_local! {
+    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> = RefCell::new(
+        MemoryManager::init(DefaultMemoryImpl::default())
+    );
+
+    // The memory where the ledger must write and read its state during an upgrade.
+    pub static UPGRADES_MEMORY: RefCell<VirtualMemory<DefaultMemoryImpl>> = MEMORY_MANAGER.with(|memory_manager|
+        RefCell::new(memory_manager.borrow().get(UPGRADES_MEMORY_ID)));
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 pub struct Ledger {
     pub balances: LedgerBalances,

diff --git a/rs/ledger_suite/icp/ledger/src/main.rs b/rs/ledger_suite/icp/ledger/src/main.rs
@@ -5,7 +5,7 @@ use dfn_core::BytesS;
 use dfn_core::{
     api::{caller, data_certificate, print, set_certified_data, time_nanos, trap_with},
     endpoint::reject_on_decode_error::{over, over_async, over_async_may_reject},
-    over_init, printer, setup, stable,
+    over_init, printer, setup,
 };
 use dfn_protobuf::protobuf;
 use ic_base_types::CanisterId;
@@ -25,6 +25,9 @@ use ic_ledger_core::{
     timestamp::TimeStamp,
     tokens::{Tokens, DECIMAL_PLACES},
 };
+use ic_stable_structures::reader::{BufferedReader, Reader};
+#[cfg(feature = "upgrade-to-memory-manager")]
+use ic_stable_structures::writer::{BufferedWriter, Writer};
 use icp_ledger::{
     max_blocks_per_request, protobuf, tokens_into_proto, AccountBalanceArgs, AccountIdBlob,
     AccountIdentifier, ArchiveInfo, ArchivedBlocksRange, ArchivedEncodedBlocksRange, Archives,
@@ -50,7 +53,7 @@ use icrc_ledger_types::{
     icrc1::transfer::TransferArg,
     icrc21::{errors::Icrc21Error, requests::ConsentMessageRequest, responses::ConsentInfo},
 };
-use ledger_canister::{Ledger, LEDGER, MAX_MESSAGE_SIZE_BYTES};
+use ledger_canister::{Ledger, LEDGER, MAX_MESSAGE_SIZE_BYTES, UPGRADES_MEMORY};
 use num_traits::cast::ToPrimitive;
 #[allow(unused_imports)]
 use on_wire::IntoWire;
@@ -745,11 +748,59 @@ fn main() {
     })
 }
 
+// We use 8MiB buffer
+const BUFFER_SIZE: usize = 8388608;
+
 fn post_upgrade(args: Option<LedgerCanisterPayload>) {
     let start = dfn_core::api::performance_counter(0);
-    let mut stable_reader = stable::StableReader::new();
+
+    // In order to read the first bytes we need to use ic_cdk.
+    // dfn_core assumes the first 4 bytes store stable memory length
+    // and return bytes starting from the 5th byte.
+    let mut magic_bytes_reader = ic_cdk::api::stable::StableReader::default();
+    const MAGIC_BYTES: &[u8; 3] = b"MGR";
+    let mut first_bytes = [0u8; 3];
+    let memory_manager_found = match magic_bytes_reader.read_exact(&mut first_bytes) {
+        Ok(_) => first_bytes == *MAGIC_BYTES,
+        Err(_) => false,
+    };
+
     let mut ledger = LEDGER.write().unwrap();
-    *ledger = ciborium::de::from_reader(&mut stable_reader).expect("Decoding stable memory failed");
+    let mut pre_upgrade_instructions_consumed = 0;
+    if !memory_manager_found {
+        // The ledger was written with dfn_core and has to be read with dfn_core in order
+        // to skip the first bytes that contain the length of the stable memory.
+        let mut stable_reader = dfn_core::stable::StableReader::new();
+        *ledger =
+            ciborium::de::from_reader(&mut stable_reader).expect("Decoding stable memory failed");
+        let mut pre_upgrade_instructions_counter_bytes = [0u8; 8];
+        pre_upgrade_instructions_consumed =
+            match stable_reader.read_exact(&mut pre_upgrade_instructions_counter_bytes) {
+                Ok(_) => u64::from_le_bytes(pre_upgrade_instructions_counter_bytes),
+                Err(_) => {
+                    // If upgrading from a version that didn't write the instructions counter to stable memory
+                    0u64
+                }
+            };
+    } else {
+        *ledger = UPGRADES_MEMORY.with_borrow(|bs| {
+            let reader = Reader::new(bs, 0);
+            let mut buffered_reader = BufferedReader::new(BUFFER_SIZE, reader);
+            let ledger_state = ciborium::de::from_reader(&mut buffered_reader).expect(
+                "Failed to read the Ledger state from memory manager managed stable memory",
+            );
+            let mut pre_upgrade_instructions_counter_bytes = [0u8; 8];
+            pre_upgrade_instructions_consumed =
+                match buffered_reader.read_exact(&mut pre_upgrade_instructions_counter_bytes) {
+                    Ok(_) => u64::from_le_bytes(pre_upgrade_instructions_counter_bytes),
+                    Err(_) => {
+                        // If upgrading from a version that didn't write the instructions counter to stable memory
+                        0u64
+                    }
+                };
+            ledger_state
+        });
+    }
 
     if let Some(args) = args {
         match args {
@@ -768,15 +819,6 @@ fn post_upgrade(args: Option<LedgerCanisterPayload>) {
             .map(|h| h.into_bytes())
             .unwrap_or([0u8; 32]),
     );
-    let mut pre_upgrade_instructions_counter_bytes = [0u8; 8];
-    let pre_upgrade_instructions_consumed =
-        match stable_reader.read_exact(&mut pre_upgrade_instructions_counter_bytes) {
-            Ok(_) => u64::from_le_bytes(pre_upgrade_instructions_counter_bytes),
-            Err(_) => {
-                // If upgrading from a version that didn't write the instructions counter to stable memory
-                0u64
-            }
-        };
     PRE_UPGRADE_INSTRUCTIONS_CONSUMED.with(|n| *n.borrow_mut() = pre_upgrade_instructions_consumed);
 
     let end = dfn_core::api::performance_counter(0);
@@ -790,6 +832,7 @@ fn post_upgrade_() {
     over_init(|CandidOne(args)| post_upgrade(args));
 }
 
+#[cfg(not(feature = "upgrade-to-memory-manager"))]
 #[export_name = "canister_pre_upgrade"]
 fn pre_upgrade() {
     let start = dfn_core::api::performance_counter(0);
@@ -801,7 +844,7 @@ fn pre_upgrade() {
         .read()
         // This should never happen, but it's better to be safe than sorry
         .unwrap_or_else(|poisoned| poisoned.into_inner());
-    let mut stable_writer = stable::StableWriter::new();
+    let mut stable_writer = dfn_core::stable::StableWriter::new();
     ciborium::ser::into_writer(&*ledger, &mut stable_writer)
         .expect("failed to write ledger state to stable memory");
     let end = dfn_core::api::performance_counter(0);
@@ -812,6 +855,32 @@ fn pre_upgrade() {
         .expect("failed to write instructions consumed to stable memory");
 }
 
+#[cfg(feature = "upgrade-to-memory-manager")]
+#[export_name = "canister_pre_upgrade"]
+fn pre_upgrade() {
+    let start = dfn_core::api::performance_counter(0);
+    setup::START.call_once(|| {
+        printer::hook();
+    });
+
+    let ledger = LEDGER
+        .read()
+        // This should never happen, but it's better to be safe than sorry
+        .unwrap_or_else(|poisoned| poisoned.into_inner());
+    UPGRADES_MEMORY.with_borrow_mut(|bs| {
+        let writer = Writer::new(bs, 0);
+        let mut buffered_writer = BufferedWriter::new(BUFFER_SIZE, writer);
+        ciborium::ser::into_writer(&*ledger, &mut buffered_writer)
+            .expect("Failed to write the Ledger state to memory manager managed stable memory");
+        let end = dfn_core::api::performance_counter(0);
+        let instructions_consumed = end - start;
+        let counter_bytes: [u8; 8] = instructions_consumed.to_le_bytes();
+        buffered_writer
+            .write_all(&counter_bytes)
+            .expect("failed to write instructions consumed to UPGRADES_MEMORY");
+    });
+}
+
 struct Access;
 
 impl LedgerAccess for Access {

diff --git a/rs/ledger_suite/icp/ledger/tests/tests.rs b/rs/ledger_suite/icp/ledger/tests/tests.rs
@@ -30,7 +30,7 @@ use icrc_ledger_types::icrc2::approve::ApproveArgs;
 use num_traits::cast::ToPrimitive;
 use on_wire::{FromWire, IntoWire};
 use serde_bytes::ByteBuf;
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
@@ -46,6 +46,14 @@ fn ledger_wasm() -> Vec<u8> {
     )
 }
 
+fn ledger_wasm_upgrade_to_memory_manager() -> Vec<u8> {
+    ic_test_utilities_load_wasm::load_wasm(
+        std::env::var("CARGO_MANIFEST_DIR").unwrap(),
+        "ledger-canister-upgrade-to-memory-manager",
+        &[],
+    )
+}
+
 fn encode_init_args(args: ic_icrc1_ledger_sm_tests::InitArgs) -> LedgerCanisterInitPayload {
     let initial_values = args
         .initial_balances
@@ -1147,27 +1155,36 @@ fn test_upgrade_serialization() {
     ic_icrc1_ledger_sm_tests::test_upgrade_serialization(
         ledger_wasm_mainnet,
         ledger_wasm_current,
-        None,
+        Some(ledger_wasm_upgrade_to_memory_manager()),
         init_args,
         upgrade_args,
         minter,
         false,
+        false,
     );
 }
 
 #[test]
-fn test_approval_upgrade() {
+fn test_upgrade_serialization_fixed_tx() {
     let ledger_wasm_mainnet =
         std::fs::read(std::env::var("ICP_LEDGER_DEPLOYED_VERSION_WASM_PATH").unwrap()).unwrap();
     let ledger_wasm_current = ledger_wasm();
+    let ledger_wasm_upgradetomemorymanager = ledger_wasm_upgrade_to_memory_manager();
 
     let p1 = PrincipalId::new_user_test_id(1);
     let p2 = PrincipalId::new_user_test_id(2);
     let p3 = PrincipalId::new_user_test_id(3);
+    let accounts = vec![
+        Account::from(p1.0),
+        Account::from(p2.0),
+        Account::from(p3.0),
+    ];
 
     let env = StateMachine::new();
     let mut initial_balances = HashMap::new();
-    initial_balances.insert(Account::from(p1.0).into(), Tokens::from_e8s(10_000_000));
+    for account in &accounts {
+        initial_balances.insert((*account).into(), Tokens::from_e8s(10_000_000));
+    }
 
     let payload = LedgerCanisterInitPayload::builder()
         .minting_account(MINTER.into())
@@ -1193,6 +1210,11 @@ fn test_approval_upgrade() {
     approve_args.expires_at = Some(expiration);
     send_approval(&env, canister_id, p1.0, &approve_args).expect("approval failed");
 
+    let mut balances = BTreeMap::new();
+    for account in &accounts {
+        balances.insert(account, balance_of(&env, canister_id, *account));
+    }
+
     let test_upgrade = |ledger_wasm: Vec<u8>| {
         env.upgrade_canister(
             canister_id,
@@ -1208,11 +1230,21 @@ fn test_approval_upgrade() {
         let allowance = get_allowance(&env, canister_id, p1.0, p3.0);
         assert_eq!(allowance.allowance.0.to_u64().unwrap(), 130_000);
         assert_eq!(allowance.expires_at, Some(expiration));
+
+        for account in &accounts {
+            assert_eq!(balances[account], balance_of(&env, canister_id, *account));
+        }
     };
 
-    // Test if the old serialized approvals are correctly deserialized
+    // Test if the old serialized approvals and balances are correctly deserialized
+    test_upgrade(ledger_wasm_current.clone());
+    // Test the new wasm serialization
     test_upgrade(ledger_wasm_current.clone());
-    // Test if new approvals serialization also works
+    // Test serializing to the memory manager
+    test_upgrade(ledger_wasm_upgradetomemorymanager.clone());
+    // Test upgrade to memory manager again
+    test_upgrade(ledger_wasm_upgradetomemorymanager);
+    // Test deserializing from memory manager
     test_upgrade(ledger_wasm_current);
     // Test if downgrade works
     test_upgrade(ledger_wasm_mainnet);
@@ -1582,7 +1614,9 @@ fn test_icrc21_standard() {
 }
 
 mod metrics {
-    use crate::{encode_init_args, encode_upgrade_args, ledger_wasm};
+    use crate::{
+        encode_init_args, encode_upgrade_args, ledger_wasm, ledger_wasm_upgrade_to_memory_manager,
+    };
     use ic_icrc1_ledger_sm_tests::metrics::LedgerSuiteType;
 
     #[test]
@@ -1614,7 +1648,7 @@ mod metrics {
     fn should_set_ledger_upgrade_instructions_consumed_metric() {
         ic_icrc1_ledger_sm_tests::metrics::assert_ledger_upgrade_instructions_consumed_metric_set(
             ledger_wasm(),
-            None,
+            Some(ledger_wasm_upgrade_to_memory_manager()),
             encode_init_args,
             encode_upgrade_args,
         );