feat: Add wasm_memory_limit to canister settings (#548)

ulan · lwshang · web-flow · commit b9e49377a812 · 2024-04-19T15:43:26.000Z
* feat: Add `wasm_memory_limit` to canister settings Part of https://dfinity.atlassian.net/browse/SDK-1588 * use grouping --------- Co-authored-by: Linwei Shang <linwei.shang@dfinity.org>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   regardless the Agent level configuration from `AgentBuilder::with_verify_query_signatures`.
 * Function `Agent::fetch_api_boundary_nodes()` is split into two functions: `fetch_api_boundary_nodes_by_canister_id()` and `fetch_api_boundary_nodes_by_subnet_id()`.
 * `ReqwestTransport` and `HyperTransport` structures storing the trait object `route_provider: Box<dyn RouteProvider>` have been modified to allow for shared ownership via `Arc<dyn RouteProvider>`.
+* Added `wasm_memory_limit` to canister creation and canister setting update options.
  
 ## [0.34.0] - 2024-03-18
 
diff --git a/ic-utils/src/interfaces/management_canister.rs b/ic-utils/src/interfaces/management_canister.rs
@@ -157,6 +157,8 @@ pub struct DefiniteCanisterSettings {
     pub freezing_threshold: Nat,
     /// The upper limit of the canister's reserved cycles balance.
     pub reserved_cycles_limit: Option<Nat>,
+    /// A soft limit on the Wasm memory usage of the canister in bytes (up to 256TiB).
+    pub wasm_memory_limit: Option<Nat>,
 }
 
 impl std::fmt::Display for StatusCallResult {
diff --git a/ic-utils/src/interfaces/management_canister/attributes.rs b/ic-utils/src/interfaces/management_canister/attributes.rs
@@ -207,6 +207,53 @@ try_from_reserved_cycles_limit_decl!(i64);
 try_from_reserved_cycles_limit_decl!(i128);
 try_from_reserved_cycles_limit_decl!(u128);
 
+/// An error encountered when attempting to construct a [`WasmMemoryLimit`].
+#[derive(Error, Debug)]
+pub enum WasmMemoryLimitError {
+    /// The provided value was not in the range [0, 2^48] (i.e. 256 TiB).
+    #[error("Wasm memory limit must be between 0 and 2^48 (i.e 256TiB), inclusively. Got {0}.")]
+    InvalidMemoryLimit(i64),
+}
+
+/// A soft limit on the Wasm memory usage of the canister. Update calls,
+/// timers, heartbeats, install, and post-upgrade fail if the Wasm memory
+/// usage exceeds this limit. The main purpose of this field is to protect
+/// against the case when the canister reaches the hard 4GiB limit.
+/// Must be a number between 0 and 2^48^ (i.e 256TB), inclusively.
+#[derive(Copy, Clone, Debug)]
+pub struct WasmMemoryLimit(u64);
+
+impl std::convert::From<WasmMemoryLimit> for u64 {
+    fn from(wasm_memory_limit: WasmMemoryLimit) -> Self {
+        wasm_memory_limit.0
+    }
+}
+
+macro_rules! try_from_wasm_memory_limit_decl {
+    ( $t: ty ) => {
+        impl std::convert::TryFrom<$t> for WasmMemoryLimit {
+            type Error = WasmMemoryLimitError;
+
+            fn try_from(value: $t) -> Result<Self, Self::Error> {
+                if (value as i64) < 0 || (value as i64) > (1i64 << 48) {
+                    Err(Self::Error::InvalidMemoryLimit(value as i64))
+                } else {
+                    Ok(Self(value as u64))
+                }
+            }
+        }
+    };
+}
+
+try_from_wasm_memory_limit_decl!(u8);
+try_from_wasm_memory_limit_decl!(u16);
+try_from_wasm_memory_limit_decl!(u32);
+try_from_wasm_memory_limit_decl!(u64);
+try_from_wasm_memory_limit_decl!(i8);
+try_from_wasm_memory_limit_decl!(i16);
+try_from_wasm_memory_limit_decl!(i32);
+try_from_wasm_memory_limit_decl!(i64);
+
 #[test]
 #[allow(clippy::useless_conversion)]
 fn can_convert_compute_allocation() {
@@ -296,3 +343,32 @@ fn can_convert_reserved_cycles_limit() {
     let ft = ReservedCyclesLimit(100);
     let _ft_ft: ReservedCyclesLimit = ReservedCyclesLimit::try_from(ft).unwrap();
 }
+
+#[test]
+#[allow(clippy::useless_conversion)]
+fn can_convert_wasm_memory_limit() {
+    use std::convert::{TryFrom, TryInto};
+
+    // This is more of a compiler test than an actual test.
+    let _ma_u8: WasmMemoryLimit = 1u8.try_into().unwrap();
+    let _ma_u16: WasmMemoryLimit = 1u16.try_into().unwrap();
+    let _ma_u32: WasmMemoryLimit = 1u32.try_into().unwrap();
+    let _ma_u64: WasmMemoryLimit = 1u64.try_into().unwrap();
+    let _ma_i8: WasmMemoryLimit = 1i8.try_into().unwrap();
+    let _ma_i16: WasmMemoryLimit = 1i16.try_into().unwrap();
+    let _ma_i32: WasmMemoryLimit = 1i32.try_into().unwrap();
+    let _ma_i64: WasmMemoryLimit = 1i64.try_into().unwrap();
+
+    let ma = WasmMemoryLimit(100);
+    let _ma_ma: WasmMemoryLimit = WasmMemoryLimit::try_from(ma).unwrap();
+
+    assert!(matches!(
+        WasmMemoryLimit::try_from(-4).unwrap_err(),
+        WasmMemoryLimitError::InvalidMemoryLimit(-4)
+    ));
+
+    assert!(matches!(
+        WasmMemoryLimit::try_from(562949953421312_u64).unwrap_err(),
+        WasmMemoryLimitError::InvalidMemoryLimit(562949953421312)
+    ));
+}
diff --git a/ic-utils/src/interfaces/management_canister/builders.rs b/ic-utils/src/interfaces/management_canister/builders.rs
@@ -2,7 +2,7 @@
 
 #[doc(inline)]
 pub use super::attributes::{
-    ComputeAllocation, FreezingThreshold, MemoryAllocation, ReservedCyclesLimit,
+    ComputeAllocation, FreezingThreshold, MemoryAllocation, ReservedCyclesLimit, WasmMemoryLimit,
 };
 use super::{ChunkHash, ManagementCanister};
 use crate::{
@@ -62,6 +62,16 @@ pub struct CanisterSettings {
     /// If set to 0, disables the reservation mechanism for the canister.
     /// Doing so will cause the canister to trap when it tries to allocate storage, if the subnet's usage exceeds 450 GiB.
     pub reserved_cycles_limit: Option<Nat>,
+
+    /// A soft limit on the Wasm memory usage of the canister.
+    ///
+    /// Update calls, timers, heartbeats, install, and post-upgrade fail if the
+    /// Wasm memory usage exceeds this limit. The main purpose of this field is
+    /// to protect against the case when the canister reaches the hard 4GiB
+    /// limit.
+    ///
+    /// Must be a number between 0 and 2^48^ (i.e 256TB), inclusively.
+    pub wasm_memory_limit: Option<Nat>,
 }
 
 /// A builder for a `create_canister` call.
@@ -74,6 +84,7 @@ pub struct CreateCanisterBuilder<'agent, 'canister: 'agent> {
     memory_allocation: Option<Result<MemoryAllocation, AgentError>>,
     freezing_threshold: Option<Result<FreezingThreshold, AgentError>>,
     reserved_cycles_limit: Option<Result<ReservedCyclesLimit, AgentError>>,
+    wasm_memory_limit: Option<Result<WasmMemoryLimit, AgentError>>,
     is_provisional_create: bool,
     amount: Option<u128>,
     specified_id: Option<Principal>,
@@ -90,6 +101,7 @@ impl<'agent, 'canister: 'agent> CreateCanisterBuilder<'agent, 'canister> {
             memory_allocation: None,
             freezing_threshold: None,
             reserved_cycles_limit: None,
+            wasm_memory_limit: None,
             is_provisional_create: false,
             amount: None,
             specified_id: None,
@@ -278,6 +290,32 @@ impl<'agent, 'canister: 'agent> CreateCanisterBuilder<'agent, 'canister> {
         }
     }
 
+    /// Pass in a Wasm memory limit value for the canister.
+    pub fn with_wasm_memory_limit<C, E>(self, wasm_memory_limit: C) -> Self
+    where
+        E: std::fmt::Display,
+        C: TryInto<WasmMemoryLimit, Error = E>,
+    {
+        self.with_optional_wasm_memory_limit(Some(wasm_memory_limit))
+    }
+
+    /// Pass in a Wasm memory limit optional value for the canister. If this is [None],
+    /// it will revert the Wasm memory limit to default.
+    pub fn with_optional_wasm_memory_limit<E, C>(self, wasm_memory_limit: Option<C>) -> Self
+    where
+        E: std::fmt::Display,
+        C: TryInto<WasmMemoryLimit, Error = E>,
+    {
+        Self {
+            wasm_memory_limit: wasm_memory_limit.map(|limit| {
+                limit
+                    .try_into()
+                    .map_err(|e| AgentError::MessageError(format!("{}", e)))
+            }),
+            ..self
+        }
+    }
+
     /// Create an [AsyncCall] implementation that, when called, will create a
     /// canister.
     pub fn build(self) -> Result<impl 'agent + AsyncCall<(Principal,)>, AgentError> {
@@ -306,6 +344,11 @@ impl<'agent, 'canister: 'agent> CreateCanisterBuilder<'agent, 'canister> {
             Some(Ok(x)) => Some(Nat::from(u128::from(x))),
             None => None,
         };
+        let wasm_memory_limit = match self.wasm_memory_limit {
+            Some(Err(x)) => return Err(AgentError::MessageError(format!("{}", x))),
+            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
+            None => None,
+        };
 
         #[derive(Deserialize, CandidType)]
         struct Out {
@@ -327,6 +370,7 @@ impl<'agent, 'canister: 'agent> CreateCanisterBuilder<'agent, 'canister> {
                     memory_allocation,
                     freezing_threshold,
                     reserved_cycles_limit,
+                    wasm_memory_limit,
                 },
                 specified_id: self.specified_id,
             };
@@ -343,6 +387,7 @@ impl<'agent, 'canister: 'agent> CreateCanisterBuilder<'agent, 'canister> {
                     memory_allocation,
                     freezing_threshold,
                     reserved_cycles_limit,
+                    wasm_memory_limit,
                 })
                 .with_effective_canister_id(self.effective_canister_id)
         };
@@ -808,6 +853,7 @@ pub struct UpdateCanisterBuilder<'agent, 'canister: 'agent> {
     memory_allocation: Option<Result<MemoryAllocation, AgentError>>,
     freezing_threshold: Option<Result<FreezingThreshold, AgentError>>,
     reserved_cycles_limit: Option<Result<ReservedCyclesLimit, AgentError>>,
+    wasm_memory_limit: Option<Result<WasmMemoryLimit, AgentError>>,
 }
 
 impl<'agent, 'canister: 'agent> UpdateCanisterBuilder<'agent, 'canister> {
@@ -821,6 +867,7 @@ impl<'agent, 'canister: 'agent> UpdateCanisterBuilder<'agent, 'canister> {
             memory_allocation: None,
             freezing_threshold: None,
             reserved_cycles_limit: None,
+            wasm_memory_limit: None,
         }
     }
 
@@ -963,6 +1010,32 @@ impl<'agent, 'canister: 'agent> UpdateCanisterBuilder<'agent, 'canister> {
         }
     }
 
+    /// Pass in a Wasm memory limit value for the canister.
+    pub fn with_wasm_memory_limit<C, E>(self, wasm_memory_limit: C) -> Self
+    where
+        E: std::fmt::Display,
+        C: TryInto<WasmMemoryLimit, Error = E>,
+    {
+        self.with_optional_wasm_memory_limit(Some(wasm_memory_limit))
+    }
+
+    /// Pass in a Wasm memory limit optional value for the canister. If this is [None],
+    /// leaves the Wasm memory limit unchanged.
+    pub fn with_optional_wasm_memory_limit<E, C>(self, wasm_memory_limit: Option<C>) -> Self
+    where
+        E: std::fmt::Display,
+        C: TryInto<WasmMemoryLimit, Error = E>,
+    {
+        Self {
+            wasm_memory_limit: wasm_memory_limit.map(|limit| {
+                limit
+                    .try_into()
+                    .map_err(|e| AgentError::MessageError(format!("{}", e)))
+            }),
+            ..self
+        }
+    }
+
     /// Create an [AsyncCall] implementation that, when called, will update a
     /// canisters settings.
     pub fn build(self) -> Result<impl 'agent + AsyncCall<()>, AgentError> {
@@ -997,6 +1070,11 @@ impl<'agent, 'canister: 'agent> UpdateCanisterBuilder<'agent, 'canister> {
             Some(Ok(x)) => Some(Nat::from(u128::from(x))),
             None => None,
         };
+        let wasm_memory_limit = match self.wasm_memory_limit {
+            Some(Err(x)) => return Err(AgentError::MessageError(format!("{}", x))),
+            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
+            None => None,
+        };
 
         Ok(self
             .canister
@@ -1009,6 +1087,7 @@ impl<'agent, 'canister: 'agent> UpdateCanisterBuilder<'agent, 'canister> {
                     memory_allocation,
                     freezing_threshold,
                     reserved_cycles_limit,
+                    wasm_memory_limit,
                 },
             })
             .with_effective_canister_id(self.canister_id)
diff --git a/ic-utils/src/interfaces/wallet.rs b/ic-utils/src/interfaces/wallet.rs
@@ -658,6 +658,7 @@ impl<'agent> WalletCanister<'agent> {
             memory_allocation: memory_allocation.map(u64::from).map(Nat::from),
             freezing_threshold: freezing_threshold.map(u64::from).map(Nat::from),
             reserved_cycles_limit: None,
+            wasm_memory_limit: None,
         };
 
         self.update("wallet_create_canister")
@@ -687,6 +688,7 @@ impl<'agent> WalletCanister<'agent> {
             memory_allocation: memory_allocation.map(u64::from).map(Nat::from),
             freezing_threshold: freezing_threshold.map(u64::from).map(Nat::from),
             reserved_cycles_limit: None,
+            wasm_memory_limit: None,
         };
 
         self.update("wallet_create_canister128")
@@ -700,6 +702,10 @@ impl<'agent> WalletCanister<'agent> {
     /// This method does not have a `reserved_cycles_limit` parameter,
     /// as the wallet does not support the setting.  If you need to create a canister
     /// with a `reserved_cycles_limit` set, use the management canister.
+    ///
+    /// This method does not have a `wasm_memory_limit` parameter,
+    /// as the wallet does not support the setting.  If you need to create a canister
+    /// with a `wasm_memory_limit` set, use the management canister.
     pub async fn wallet_create_canister(
         &self,
         cycles: u128,
@@ -810,6 +816,7 @@ impl<'agent> WalletCanister<'agent> {
             memory_allocation: memory_allocation.map(u64::from).map(Nat::from),
             freezing_threshold: freezing_threshold.map(u64::from).map(Nat::from),
             reserved_cycles_limit: None,
+            wasm_memory_limit: None,
         };
 
         self.update("wallet_create_wallet")
@@ -839,6 +846,7 @@ impl<'agent> WalletCanister<'agent> {
             memory_allocation: memory_allocation.map(u64::from).map(Nat::from),
             freezing_threshold: freezing_threshold.map(u64::from).map(Nat::from),
             reserved_cycles_limit: None,
+            wasm_memory_limit: None,
         };
 
         self.update("wallet_create_wallet128")
diff --git a/ref-tests/tests/ic-ref.rs b/ref-tests/tests/ic-ref.rs
diff --git a/ref-tests/tests/integration.rs b/ref-tests/tests/integration.rs

Original file line number	Diff line number	Diff line change
`@@ -157,6 +157,8 @@ pub struct DefiniteCanisterSettings {`
`157`	`157`	`pub freezing_threshold: Nat,`
`158`	`158`	`/// The upper limit of the canister's reserved cycles balance.`
`159`	`159`	`pub reserved_cycles_limit: Option<Nat>,`
	`160`	`+ /// A soft limit on the Wasm memory usage of the canister in bytes (up to 256TiB).`
	`161`	`+ pub wasm_memory_limit: Option<Nat>,`
`160`	`162`	`}`
`161`	`163`
`162`	`164`	`impl std::fmt::Display for StatusCallResult {`