feat: Add an optional argument to set memory allocation

nix/sources.json

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

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

@@ -5,9 +5,11 @@ use crate::lib::error::{DfxError, DfxResult};
 use crate::lib::message::UserMessage;
 
 use clap::{App, Arg, ArgMatches, SubCommand};
-use ic_http_agent::{Agent, Blob, CanisterAttributes, ComputeAllocation, RequestId};
+use ic_http_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 +44,23 @@ 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)
+                .default_value("8GB")
+                .validator(memory_allocation_validator),
+        )
 }
 
 async fn install_canister(
     env: &dyn Environment,
     agent: &Agent,
     canister_info: &CanisterInfo,
     compute_allocation: ComputeAllocation,
+    memory_allocation: MemoryAllocation,
 ) -> DfxResult<RequestId> {
     let log = env.get_logger();
     let canister_id = canister_info.get_canister_id().ok_or_else(|| {
@@ -70,7 +82,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 +100,35 @@ fn compute_allocation_validator(compute_allocation: String) -> Result<(), String
     Err("Must be a percent between 0 and 100".to_string())
 }
 
+fn parse_memory_allocation(memory_allocation: String) -> Result<u64, String> {
+    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(|_| format!("Could not parse number: {}", raw_num))?;
+    let unit = unit.trim();
+    match unit {
+        "KB" => Ok(raw_num * 1024),
+        "MB" => Ok(raw_num * 1024 * 1024),
+        "GB" => Ok(raw_num * 1024 * 1024 * 1024),
+        _ => return Err(format!("Invalid unit for memory allocation, {}. Expected one of <KB|MB|GB>", unit)),
+    }
+}
+
+fn memory_allocation_validator(memory_allocation: String) -> Result<(), String> {
+    let num = parse_memory_allocation(memory_allocation)?;
+
+    if num <= (1 << 48) {
+        Ok(())
+    } else {
+        Err("Must be a number of bytes between 0 and 2^48".to_string())
+    }
+}
+
 pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
     let log = env.get_logger();
     let config = env
@@ -101,6 +145,16 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
         .try_into()
         .expect("Compute Allocation must be a percentage.");
 
+    let memory_allocation = MemoryAllocation::try_from(
+        parse_memory_allocation(
+            args.value_of("memory-allocation")
+                .unwrap_or("8GB")
+                .to_string(),
+        )
+        .unwrap(),
+    )
+    .expect("Memory Allocation must be a number between 0 and 2^48");
+
     let mut runtime = Runtime::new().expect("Unable to create a runtime");
 
     if let Some(canister_name) = args.value_of("canister_name") {
@@ -110,6 +164,7 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult {
             &agent,
             &canister_info,
             compute_allocation,
+            memory_allocation,
         ))?;
 
         if args.is_present("async") {
@@ -132,6 +187,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. This should be a number in the range [0..2^48]",
 
     // dfx canister mod
     ManageCanister => "Manages canisters deployed on a network client.",

src/ic_http_agent/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.0,
         })
         .await
     }

src/ic_http_agent/src/agent/replica_api.rs

@@ -48,6 +48,7 @@ pub enum SubmitRequest<'a> {
         arg: &'a Blob,
         nonce: &'a Option<Blob>,
         compute_allocation: u8,
+        memory_allocation: u64,
     },
     Call {
         canister_id: &'a CanisterId,

src/ic_http_agent/src/types/canister_attributes.rs

@@ -46,14 +46,52 @@ impl std::convert::TryFrom<u8> for ComputeAllocation {
     }
 }
 
+#[derive(Clone, Debug)]
+pub enum MemoryAllocationError {
+    ValueOutOfRange,
+}
+
+impl std::fmt::Display for MemoryAllocationError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+        match self {
+            MemoryAllocationError::ValueOutOfRange => {
+                f.write_str("Must be a number between 0 and 2^48.")
+            }
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct MemoryAllocation(pub(crate) 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)
+        } else {
+            Ok(Self(value))
+        }
+    }
+}
+
 pub struct CanisterAttributes {
     pub compute_allocation: ComputeAllocation,
+    pub memory_allocation: MemoryAllocation,
 }
 
 impl Default for CanisterAttributes {
     fn default() -> Self {
         CanisterAttributes {
             compute_allocation: ComputeAllocation(0),
+            memory_allocation: MemoryAllocation(8 * 1024 * 1024 * 1024),
         }
     }
 }

src/ic_http_agent/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};