diff --git a/src/dfx/src/commands/canister/install.rs b/src/dfx/src/commands/canister/install.rs index f98dd492fe..fde7a91237 100644 --- a/src/dfx/src/commands/canister/install.rs +++ b/src/dfx/src/commands/canister/install.rs @@ -1,3 +1,4 @@ +use crate::config::dfinity::ConfigInterface; use crate::lib::canister_info::CanisterInfo; use crate::lib::environment::Environment; use crate::lib::error::{DfxError, DfxResult}; @@ -71,7 +72,6 @@ pub fn construct() -> App<'static, 'static> { Arg::with_name("memory-allocation") .help(UserMessage::InstallMemoryAllocation.to_str()) .long("memory-allocation") - .default_value("8GB") .takes_value(true) .validator(memory_allocation_validator), ) @@ -96,6 +96,36 @@ fn memory_allocation_validator(memory_allocation: String) -> Result<(), String> Err("Must be a value between 0..256 TB inclusive.".to_string()) } +fn get_compute_allocation( + args: &ArgMatches<'_>, + config_interface: &ConfigInterface, + canister_name: &str, +) -> DfxResult> { + Ok(args + .value_of("compute-allocation") + .map(|v| v.to_string()) + .or(config_interface.get_compute_allocation(canister_name)?) + .map(|arg| { + ComputeAllocation::try_from(arg.parse::().unwrap()) + .expect("Compute Allocation must be a percentage.") + })) +} + +fn get_memory_allocation( + args: &ArgMatches<'_>, + config_interface: &ConfigInterface, + canister_name: &str, +) -> DfxResult> { + Ok(args + .value_of("memory-allocation") + .map(|v| v.to_string()) + .or(config_interface.get_memory_allocation(canister_name)?) + .map(|arg| { + MemoryAllocation::try_from(u64::try_from(arg.parse::().unwrap().size()).unwrap()) + .expect("Memory allocation must be between 0 and 2^48 (i.e 256TB), inclusively.") + })) +} + pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult { let config = env .get_config() @@ -107,16 +137,7 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult { .get_agent() .ok_or(DfxError::CommandMustBeRunInAProject)?; - let compute_allocation = args.value_of("compute-allocation").map(|arg| { - ComputeAllocation::try_from(arg.parse::().unwrap()) - .expect("Compute Allocation must be a percentage.") - }); - - let memory_allocation: Option = - args.value_of("memory-allocation").map(|arg| { - MemoryAllocation::try_from(u64::try_from(arg.parse::().unwrap().size()).unwrap()) - .expect("Memory allocation must be between 0 and 2^48 (i.e 256TB), inclusively.") - }); + let config_interface = config.get_config(); let mode = InstallMode::from_str(args.value_of("mode").unwrap())?; @@ -134,6 +155,9 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult { let arg_type: Option<&str> = args.value_of("type"); let install_args = blob_from_arguments(arguments, arg_type, &init_type)?; + let compute_allocation = get_compute_allocation(args, config_interface, canister_name)?; + let memory_allocation = get_memory_allocation(args, config_interface, canister_name)?; + runtime.block_on(install_canister( env, &agent, @@ -154,6 +178,11 @@ pub fn exec(env: &dyn Environment, args: &ArgMatches<'_>) -> DfxResult { let install_args = []; + let compute_allocation = + get_compute_allocation(args, config_interface, canister_name)?; + let memory_allocation = + get_memory_allocation(args, config_interface, canister_name)?; + runtime.block_on(install_canister( env, &agent, diff --git a/src/dfx/src/config/dfinity.rs b/src/dfx/src/config/dfinity.rs index 847ec7b819..822ab74828 100644 --- a/src/dfx/src/config/dfinity.rs +++ b/src/dfx/src/config/dfinity.rs @@ -261,6 +261,38 @@ impl ConfigInterface { Ok(canister_names) } + + pub fn get_compute_allocation(&self, canister_name: &str) -> DfxResult> { + self.get_initialization_value(canister_name, "compute_allocation") + } + + pub fn get_memory_allocation(&self, canister_name: &str) -> DfxResult> { + self.get_initialization_value(canister_name, "memory_allocation") + } + + fn get_initialization_value( + &self, + canister_name: &str, + field: &str, + ) -> DfxResult> { + let canister_map = (&self.canisters).as_ref().ok_or_else(|| { + DfxError::InvalidConfiguration("No canisters in the configuration file.".to_string()) + })?; + + let canister_config = canister_map + .get(canister_name) + .ok_or_else(|| DfxError::CannotFindCanisterName(canister_name.to_string()))?; + + canister_config + .extras + .get("initialization_values") + .and_then(|v| v.get(field)) + .map(String::deserialize) + .transpose() + .map_err(|_| { + DfxError::InvalidConfiguration(format!("Field {} is of the wrong type", field)) + }) + } } fn add_dependencies( @@ -600,4 +632,53 @@ mod tests { }) ); } + + #[test] + fn get_correct_initialization_values() { + let config = Config::from_str( + r#"{ + "canisters": { + "test_project": { + "initialization_values": { + "compute_allocation" : "100", + "memory_allocation": "8GB" + } + } + } + }"#, + ) + .unwrap(); + + let config_interface = config.get_config(); + let compute_allocation = config_interface + .get_compute_allocation("test_project") + .unwrap() + .unwrap(); + assert_eq!("100", compute_allocation); + + let memory_allocation = config_interface + .get_memory_allocation("test_project") + .unwrap() + .unwrap(); + assert_eq!("8GB", memory_allocation); + + let config_no_values = Config::from_str( + r#"{ + "canisters": { + "test_project_two": { + } + } + }"#, + ) + .unwrap(); + let config_interface = config_no_values.get_config(); + let compute_allocation = config_interface + .get_compute_allocation("test_project_two") + .unwrap(); + let memory_allocation = config_interface + .get_memory_allocation("test_project_two") + .unwrap(); + assert_eq!(None, compute_allocation); + assert_eq!(None, memory_allocation); + } } diff --git a/src/dfx/src/lib/operations/canister/deploy_canisters.rs b/src/dfx/src/lib/operations/canister/deploy_canisters.rs index b1f0ae138f..375804cd61 100644 --- a/src/dfx/src/lib/operations/canister/deploy_canisters.rs +++ b/src/dfx/src/lib/operations/canister/deploy_canisters.rs @@ -8,9 +8,11 @@ use crate::lib::models::canister_id_store::CanisterIdStore; use crate::lib::operations::canister::create_canister; use crate::lib::operations::canister::install_canister; use crate::util::{blob_from_arguments, get_candid_init_type}; +use humanize_rs::bytes::Bytes; use ic_agent::AgentError; -use ic_utils::interfaces::management_canister::InstallMode; +use ic_utils::interfaces::management_canister::{ComputeAllocation, InstallMode, MemoryAllocation}; use slog::{info, warn}; +use std::convert::TryFrom; use std::time::Duration; use tokio::runtime::Runtime; @@ -124,8 +126,23 @@ fn install_canisters( let init_type = maybe_path.and_then(|path| get_candid_init_type(&path)); let install_args = blob_from_arguments(argument, argument_type, &init_type)?; - let compute_allocation = None; - let memory_allocation = None; + let config_interface = config.get_config(); + let compute_allocation = + config_interface + .get_compute_allocation(canister_name)? + .map(|arg| { + ComputeAllocation::try_from(arg.parse::().unwrap()) + .expect("Compute Allocation must be a percentage.") + }); + let memory_allocation = config_interface + .get_memory_allocation(canister_name)? + .map(|arg| { + MemoryAllocation::try_from( + u64::try_from(arg.parse::().unwrap().size()).unwrap(), + ) + .expect("Memory allocation must be between 0 and 2^48 (i.e 256TB), inclusively.") + }); + let result = runtime.block_on(install_canister( env, &agent,