fix: make asset canister reserve space when saving to stable storage (#4036)

https://dfinity.atlassian.net/browse/SDK-1827

Author: ericswanson-dfinity

CHANGELOG.md

@@ -69,6 +69,16 @@ Your principal for ICP wallets and decentralized exchanges: ueuar-wxbnk-bdcsr-dn
 
 ## Dependencies
 
+### Frontend canister
+
+### fix: 'unreachable' error when trying to upgrade an asset canister with over 1GB data
+
+The asset canister now estimates the size of the data to be serialized to stable memory,
+and reserves that much space for the ValueSerializer's buffer.
+
+- Module hash: bba3181888f3c59b4a5f608aedef05be6fa37276fb7dc394cbadf9cf6e10359b
+- https://github.com/dfinity/sdk/pull/4036
+
 ### Motoko
 
 Updated Motoko to [0.13.5](https://github.com/dfinity/motoko/releases/tag/0.13.5)

Cargo.lock

@@ -19,7 +19,7 @@ rust-version = "1.75.0"
 license = "Apache-2.0"
 
 [workspace.dependencies]
-candid = "0.10.4"
+candid = "0.10.11"
 candid_parser = "0.1.4"
 dfx-core = { path = "src/dfx-core", version = "0.1.0" }
 ic-agent = "0.39"

Cargo.toml

@@ -82,6 +82,22 @@ pub struct AssetEncoding {
 }
 
 impl AssetEncoding {
+    fn estimate_size(&self) -> usize {
+        let mut size = 0;
+        size += 8; // modified
+        size += self.total_length + self.content_chunks.len() * 4;
+        size += 5; // total_length
+        size += 1; //  certified
+        size += self.sha256.len();
+        size += 1 + self
+            .certificate_expression
+            .as_ref()
+            .map_or(0, |ce| 2 + ce.expression.len() + ce.expression_hash.len());
+        size += 1 + self.response_hashes.as_ref().map_or(0, |hashes| {
+            hashes.iter().fold(2, |acc, (_k, v)| acc + 2 + v.len())
+        });
+        size
+    }
     fn asset_hash_path_v2(&self, path: &AssetPath, status_code: u16) -> Option<HashTreePath> {
         self.certificate_expression.as_ref().and_then(|ce| {
             self.response_hashes.as_ref().and_then(|hashes| {
@@ -206,6 +222,25 @@ pub struct Configuration {
     pub max_bytes: Option<u64>,
 }
 
+impl Configuration {
+    fn estimate_size(&self) -> usize {
+        1 + self
+            .max_batches
+            .as_ref()
+            .map_or(0, |_| std::mem::size_of::<u64>())
+            + 1
+            + self
+                .max_chunks
+                .as_ref()
+                .map_or(0, |_| std::mem::size_of::<u64>())
+            + 1
+            + self
+                .max_bytes
+                .as_ref()
+                .map_or(0, |_| std::mem::size_of::<u64>())
+    }
+}
+
 #[derive(Default)]
 pub struct State {
     assets: HashMap<AssetKey, Asset>,
@@ -232,6 +267,16 @@ pub struct StableStatePermissions {
     manage_permissions: BTreeSet<Principal>,
 }
 
+impl StableStatePermissions {
+    fn estimate_size(&self) -> usize {
+        8 + self.commit.len() * std::mem::size_of::<Principal>()
+            + 8
+            + self.prepare.len() * std::mem::size_of::<Principal>()
+            + 8
+            + self.manage_permissions.len() * std::mem::size_of::<Principal>()
+    }
+}
+
 #[derive(Clone, Debug, CandidType, Deserialize)]
 pub struct StableState {
     authorized: Vec<Principal>, // ignored if permissions is Some(_)
@@ -242,7 +287,46 @@ pub struct StableState {
     configuration: Option<Configuration>,
 }
 
+impl StableState {
+    pub fn estimate_size(&self) -> usize {
+        let mut size = 0;
+        size += 2 + self.authorized.len() * std::mem::size_of::<Principal>();
+        size += 1 + self.permissions.as_ref().map_or(0, |p| p.estimate_size());
+        size += self.stable_assets.iter().fold(2, |acc, (name, asset)| {
+            acc + 2 + name.len() + asset.estimate_size()
+        });
+        size += 1 + self.next_batch_id.as_ref().map_or(0, |_| 8);
+        size += 1 + self.configuration.as_ref().map_or(0, |c| c.estimate_size());
+        size
+    }
+}
+
 impl Asset {
+    fn estimate_size(&self) -> usize {
+        let mut size = 0;
+        size += 1 + self.content_type.len();
+        size += self.encodings.iter().fold(1, |acc, (name, encoding)| {
+            acc + 1 + name.len() + encoding.estimate_size()
+        });
+        size += 1 + self
+            .max_age
+            .as_ref()
+            .map_or(0, |_| std::mem::size_of::<u64>());
+        size += 1 + self.headers.as_ref().map_or(0, |hm| {
+            hm.iter()
+                .fold(2, |acc, (k, v)| acc + 1 + k.len() + 2 + v.len())
+        });
+        size += 1 + self
+            .is_aliased
+            .as_ref()
+            .map_or(0, |_| std::mem::size_of::<bool>());
+        size += 1 + self
+            .allow_raw_access
+            .as_ref()
+            .map_or(0, |_| std::mem::size_of::<bool>());
+        size
+    }
+
     fn allow_raw_access(&self) -> bool {
         self.allow_raw_access.unwrap_or(true)
     }

src/canisters/frontend/ic-certified-assets/src/state_machine.rs

@@ -1,3 +1,5 @@
+use candid::ser::IDLBuilder;
+use ic_cdk::api::stable;
 use ic_cdk::{init, post_upgrade, pre_upgrade};
 use ic_certified_assets::types::AssetCanisterArgs;
 
@@ -8,10 +10,24 @@ fn init(args: Option<AssetCanisterArgs>) {
 
 #[pre_upgrade]
 fn pre_upgrade() {
-    ic_cdk::storage::stable_save((ic_certified_assets::pre_upgrade(),))
+    let stable_state = ic_certified_assets::pre_upgrade();
+    let value_serializer_estimate = stable_state.estimate_size();
+    stable_save_with_capacity((stable_state,), value_serializer_estimate)
         .expect("failed to save stable state");
 }
 
+// this is the same as ic_cdk::storage::stable_save,
+// but reserves the capacity for the value serializer
+fn stable_save_with_capacity<T>(t: T, value_capacity: usize) -> Result<(), candid::Error>
+where
+    T: candid::utils::ArgumentEncoder,
+{
+    let mut ser = IDLBuilder::new();
+    ser.try_reserve_value_serializer_capacity(value_capacity)?;
+    t.encode(&mut ser)?;
+    ser.serialize(stable::StableWriter::default())
+}
+
 #[post_upgrade]
 fn post_upgrade(args: Option<AssetCanisterArgs>) {
     let (stable_state,): (ic_certified_assets::StableState,) =