feat: Add an optional argument to set memory allocation

nix/sources.json

@@ -32,9 +32,9 @@
         "type": "git"
     },
     "dfinity": {
-        "ref": "v0.5.6",
+        "ref": "master",
         "repo": "ssh://git@github.com/dfinity-lab/dfinity",
-        "rev": "c2445f9aa8cf16dd29aabf2c47c2a5f3c870ee76",
+        "rev": "915481642c6acd7bbfa7a4d674622bc5a2bc9fc2",
         "type": "git"
     },
     "ic-ref": {

src/agent/rust/src/agent/mod.rs

@@ -222,6 +222,7 @@ impl Agent {
             arg,
             nonce: &self.nonce_factory.generate(),
             compute_allocation: attributes.compute_allocation.0,
+            memory_allocation: attributes.memory_allocation.map(|x| x.into()),
         })
         .await
     }

src/agent/rust/src/agent/replica_api.rs

@@ -48,6 +48,8 @@ pub enum SubmitRequest<'a> {
         arg: &'a Blob,
         nonce: &'a Option<Blob>,
         compute_allocation: u8,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        memory_allocation: Option<u64>,
     },
     Call {
         canister_id: &'a CanisterId,

src/agent/rust/src/types/canister_attributes.rs

@@ -46,14 +46,86 @@ impl std::convert::TryFrom<u8> for ComputeAllocation {
     }
 }
 
+#[derive(Clone, Debug)]
+pub enum MemoryAllocationError {
+    ValueOutOfRange(u64),
+    NotANumber(String),
+    InvalidUnit(String),
+}
+
+impl std::fmt::Display for MemoryAllocationError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+        match self {
+            MemoryAllocationError::ValueOutOfRange(num) => {
+                write!(f, "Must be a number between 0 and 2^48. Got {}", num)
+            }
+            MemoryAllocationError::NotANumber(input) => {
+                write!(f, "Expecting a number for memory allocation, got {}", input)
+            }
+            MemoryAllocationError::InvalidUnit(unit) => write!(
+                f,
+                "Invalid unit for memory allocation {}. Expected one of <KB|MB|GB>.",
+                unit
+            ),
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct MemoryAllocation(u64);
+
+impl std::convert::From<MemoryAllocation> for u64 {
+    fn from(memory_allocation: MemoryAllocation) -> Self {
+        memory_allocation.0
+    }
+}
+
+impl std::convert::TryFrom<u64> for MemoryAllocation {
+    type Error = MemoryAllocationError;
+
+    fn try_from(value: u64) -> Result<Self, Self::Error> {
+        if value > (1 << 48) {
+            Err(MemoryAllocationError::ValueOutOfRange(value))
+        } else {
+            Ok(Self(value))
+        }
+    }
+}
+
+impl std::convert::TryFrom<String> for MemoryAllocation {
+    type Error = MemoryAllocationError;
+
+    fn try_from(memory_allocation: String) -> Result<Self, Self::Error> {
+        let split_point = memory_allocation.find(|c: char| !c.is_numeric());
+        let memory_allocation = memory_allocation.trim();
+        let (raw_num, unit) = split_point.map_or_else(
+            || (memory_allocation, ""),
+            |p| memory_allocation.split_at(p),
+        );
+        let raw_num = raw_num
+            .parse::<u64>()
+            .map_err(|_| MemoryAllocationError::NotANumber(raw_num.to_string()))?;
+        let unit = unit.trim();
+        let num = match unit {
+            "KB" => raw_num * 1024,
+            "MB" => raw_num * 1024 * 1024,
+            "GB" => raw_num * 1024 * 1024 * 1024,
+            _ => return Err(MemoryAllocationError::InvalidUnit(unit.to_string())),
+        };
+        MemoryAllocation::try_from(num)
+    }
+}
+
 pub struct CanisterAttributes {
     pub compute_allocation: ComputeAllocation,
+    pub memory_allocation: Option<MemoryAllocation>,
 }
 
 impl Default for CanisterAttributes {
     fn default() -> Self {
         CanisterAttributes {
             compute_allocation: ComputeAllocation(0),
+            memory_allocation: None,
         }
     }
 }

src/agent/rust/src/types/mod.rs

@@ -9,7 +9,10 @@ pub(crate) mod public {
     use super::*;
 
     pub use blob::Blob;
-    pub use canister_attributes::{CanisterAttributes, ComputeAllocation, ComputeAllocationError};
+    pub use canister_attributes::{
+        CanisterAttributes, ComputeAllocation, ComputeAllocationError, MemoryAllocation,
+        MemoryAllocationError,
+    };
     pub use canister_id::{CanisterId, TextualCanisterIdError};
     pub use principal::Principal;
     pub use request_id::{to_request_id, RequestId};

src/dfx/src/commands/canister/install.rs

@@ -5,9 +5,9 @@ use crate::lib::error::{DfxError, DfxResult};
 use crate::lib::message::UserMessage;
 
 use clap::{App, Arg, ArgMatches, SubCommand};
-use ic_agent::{Agent, Blob, CanisterAttributes, ComputeAllocation, RequestId};
+use ic_agent::{Agent, Blob, CanisterAttributes, ComputeAllocation, MemoryAllocation, RequestId};
 use slog::info;
-use std::convert::TryInto;
+use std::convert::{TryFrom, TryInto};
 use tokio::runtime::Runtime;
 
 pub fn construct() -> App<'static, 'static> {
@@ -42,13 +42,22 @@ pub fn construct() -> App<'static, 'static> {
                 .default_value("0")
                 .validator(compute_allocation_validator),
         )
+        .arg(
+            Arg::with_name("memory-allocation")
+                .help(UserMessage::InstallMemoryAllocation.to_str())
+                .long("memory-allocation")
+                .short("m")
+                .takes_value(true)
+                .validator(memory_allocation_validator),
+        )
 }
 
 async fn install_canister(
     env: &dyn Environment,
     agent: &Agent,
     canister_info: &CanisterInfo,
     compute_allocation: ComputeAllocation,
+    memory_allocation: Option<MemoryAllocation>,
 ) -> DfxResult<RequestId> {
     let log = env.get_logger();
     let canister_id = canister_info.get_canister_id().ok_or_else(|| {
@@ -70,7 +79,10 @@ async fn install_canister(
             &canister_id,
             &Blob::from(wasm),
             &Blob::empty(),
-            &CanisterAttributes { compute_allocation },
+            &CanisterAttributes {
+                compute_allocation,
+                memory_allocation,
+            },
         )
         .await
         .map_err(DfxError::from)
@@ -85,6 +97,13 @@ fn compute_allocation_validator(compute_allocation: String) -> Result<(), String
     Err("Must be a percent between 0 and 100".to_string())
 }
 
+fn memory_allocation_validator(memory_allocation: String) -> Result<(), String> {
+    match MemoryAllocation::try_from(memory_allocation) {
+        Ok(_) => Ok(()),
+        Err(err) => Err(format!("{}", err))
+    }
+}
+
 pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
     let log = env.get_logger();
     let config = env
@@ -101,6 +120,14 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
         .try_into()
         .expect("Compute Allocation must be a percentage.");
 
+    let memory_allocation = match args.value_of("memory-allocation") {
+        Some(mem_alloc_arg) => Some(
+            MemoryAllocation::try_from(mem_alloc_arg.to_string())
+            .expect("Memory Allocation must be a number between 0 and 2^48")
+        ),
+        None => None,
+    };
+
     let mut runtime = Runtime::new().expect("Unable to create a runtime");
 
     if let Some(canister_name) = args.value_of("canister_name") {
@@ -110,6 +137,7 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
             &agent,
             &canister_info,
             compute_allocation,
+            memory_allocation,
         ))?;
 
         if args.is_present("async") {
@@ -132,6 +160,7 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
                     &agent,
                     &canister_info,
                     compute_allocation,
+                    memory_allocation,
                 ))?;
 
                 if args.is_present("async") {

src/dfx/src/lib/message.rs

@@ -45,6 +45,7 @@ user_message!(
     InstallAll => "Install all canisters configured in dfx.json.",
     InstallCanisterName => "Specifies the canister name. Either this or the --all flag are required.",
     InstallComputeAllocation => "Specifies the canister's compute allocation. This should be a percent in the range [0..100]",
+    InstallMemoryAllocation => "Specifies the canister's memory allocation in bytes. The value should be in the range [0..2^48]. The input should be in the form <num><unit> where unit is one of {KB|MB|GB}.",
 
     // dfx canister mod
     ManageCanister => "Manages canisters deployed on a network client.",