feat: add method to provision more empty canisters and refactor (#415)

* add method to provision more empty canisters and refactor

* fix test

src/canister/platform_orchestrator/can.did

@@ -92,6 +92,7 @@ service : (PlatformOrchestratorInitArgs) -> {
       PlatformOrchestratorGenericArgumentType,
     ) -> (PlatformOrchestratorGenericResultType);
   populate_known_principal_for_all_subnet : () -> ();
+  provision_empty_canisters_in_a_subnet : (principal, nat64) -> (Result_1);
   provision_subnet_orchestrator_canister : (principal) -> (Result_2);
   recharge_subnet_orchestrator : () -> (Result_1);
   register_new_subnet_orchestrator : (principal, bool) -> (Result_1);

src/canister/platform_orchestrator/src/api/canister_management/mod.rs

@@ -11,6 +11,7 @@ mod global_admin;
 mod known_principal;
 pub mod logging;
 mod populate_known_principal_for_all_subnet;
+pub mod provision_empty_canisters_in_a_subnet;
 pub mod provision_subnet_orchestrator;
 mod recharge_subnet_orchestrator;
 pub mod register_new_subnet_orhestrator;

src/canister/platform_orchestrator/src/api/canister_management/provision_empty_canisters_in_a_subnet.rs

@@ -0,0 +1,19 @@
+use candid::Principal;
+use ic_cdk_macros::update;
+
+use crate::{
+    guard::is_caller::is_caller_global_admin_or_controller,
+    utils::registered_subnet_orchestrator::RegisteredSubnetOrchestrator,
+};
+
+#[update(guard = "is_caller_global_admin_or_controller")]
+async fn provision_empty_canisters_in_a_subnet(
+    subnet_orchestrator_canister_id: Principal,
+    number_of_canisters: u64,
+) -> Result<(), String> {
+    let registered_subnet_orchestrator =
+        RegisteredSubnetOrchestrator::new(subnet_orchestrator_canister_id)?;
+    registered_subnet_orchestrator
+        .provision_empty_canisters(number_of_canisters)
+        .await
+}

src/canister/platform_orchestrator/src/utils/registered_subnet_orchestrator.rs

@@ -106,4 +106,14 @@ impl RegisteredSubnetOrchestrator {
 
         res
     }
+
+    pub async fn provision_empty_canisters(&self, number_of_canisters: u64) -> Result<(), String> {
+        ic_cdk::call::<_, ()>(
+            self.canister_id,
+            "provision_empty_canisters",
+            (number_of_canisters,),
+        )
+        .await
+        .map_err(|e| e.1)
+    }
 }

src/canister/user_index/can.did

@@ -152,6 +152,7 @@ service : (UserIndexInitArgs) -> {
   issue_rewards_for_referral : (principal, principal, principal) -> (Result);
   make_individual_canister_logs_private : (principal) -> (Result_3);
   make_individual_canister_logs_public : (principal) -> (Result_3);
+  provision_empty_canisters : (nat64) -> ();
   recharge_individual_user_canister : () -> (Result_3);
   reclaim_cycles_from_individual_canisters : () -> ();
   request_cycles : (nat) -> (Result_3);

src/canister/user_index/src/api/canister_management/mod.rs

@@ -42,6 +42,7 @@ pub mod get_subnet_available_capacity;
 pub mod get_subnet_backup_capacity;
 pub mod make_individual_canister_logs_private;
 pub mod make_individual_canister_logs_public;
+pub mod provision_empty_canisters;
 pub mod recharge_individual_user_canister;
 pub mod recycle_canisters;
 pub mod request_cycles;

src/canister/user_index/src/api/canister_management/provision_empty_canisters.rs

@@ -0,0 +1,12 @@
+use ic_cdk_macros::update;
+use shared_utils::common::utils::permissions::is_caller_controller;
+
+use crate::util::canister_management::provision_number_of_empty_canisters;
+
+#[update(guard = "is_caller_controller")]
+async fn provision_empty_canisters(number_of_canisters: u64) {
+    ic_cdk::spawn(provision_number_of_empty_canisters(
+        number_of_canisters,
+        || false,
+    ));
+}

src/canister/user_index/src/api/cycle_management/mod.rs

@@ -1,3 +1,3 @@
 pub mod get_user_index_canister_cycle_balance;
 pub mod reclaim_cycles_from_individual_canisters;
-pub mod return_cycles_to_platform_orchestrator_canister;
+pub mod return_cycles_to_platform_orchestrator_canister;

src/canister/user_index/src/api/user_record/get_requester_principals_canister_id_create_if_not_exists_and_optionally_allow_referrer.rs

@@ -1,12 +1,17 @@
 use crate::{
     util::canister_management::{
         check_and_request_cycles_from_platform_orchestrator, create_empty_user_canister,
-        install_canister_wasm,
+        install_canister_wasm, provision_number_of_empty_canisters, recharge_canister,
     },
     CANISTER_DATA,
 };
 use candid::Principal;
-use ic_cdk::api::call;
+use ic_cdk::api::{
+    call,
+    management_canister::main::{
+        canister_info, canister_status, CanisterIdRecord, CanisterInfoRequest,
+    },
+};
 use ic_cdk_macros::update;
 use shared_utils::{
     canister_specific::individual_user_template::types::session::SessionType,
@@ -21,7 +26,7 @@ use shared_utils::{
         get_backup_individual_user_canister_batch_size,
         get_backup_individual_user_canister_threshold,
         get_individual_user_canister_subnet_batch_size,
-        get_individual_user_canister_subnet_threshold,
+        get_individual_user_canister_subnet_threshold, INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT,
         INDIVIDUAL_USER_CANISTER_SUBNET_MAX_CAPACITY,
     },
 };
@@ -225,6 +230,8 @@ async fn provision_new_available_canisters(individual_user_template_canister_was
 
             async move {
                 let _ = check_and_request_cycles_from_platform_orchestrator().await;
+                //recharge backup canister if required
+                recharge_empty_canister_if_required(canister_id).await;
                 let future = install_canister_wasm(
                     canister_id,
                     None,
@@ -258,32 +265,29 @@ async fn provision_new_available_canisters(individual_user_template_canister_was
     .await;
 }
 
-async fn provision_new_backup_canisters(canister_count: u64) {
-    let create_canister_futures = (0..canister_count).map(|_| {
-        let future = create_empty_user_canister();
-        future
-    });
-
-    let result_callback = |canister_id: Principal| {
-        CANISTER_DATA.with_borrow_mut(|canister_data| {
-            canister_data.backup_canister_pool.insert(canister_id)
-        });
-    };
+async fn recharge_empty_canister_if_required(canister_id: Principal) {
+    let canister_status_res = canister_status(CanisterIdRecord { canister_id }).await;
+    match canister_status_res {
+        Ok((canister_status,))
+            if canister_status.cycles < INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT =>
+        {
+            let _ = recharge_canister(&canister_id, INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT).await;
+        }
+        Err(_e) => {
+            let _ = recharge_canister(&canister_id, INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT).await;
+        }
+        _ => {}
+    }
+}
 
+async fn provision_new_backup_canisters(canister_count: u64) {
     let breaking_condition = || {
         CANISTER_DATA.with_borrow(|canister_data| {
             canister_data.backup_canister_pool.len() as u64
                 > get_backup_individual_user_canister_batch_size()
         })
     };
-
-    run_task_concurrently(
-        create_canister_futures.into_iter(),
-        10,
-        result_callback,
-        breaking_condition,
-    )
-    .await;
+    provision_number_of_empty_canisters(canister_count, breaking_condition).await;
 }
 
 #[cfg(test)]

src/canister/user_index/src/util/canister_management.rs

@@ -20,9 +20,10 @@ use shared_utils::{
     canister_specific::{
         individual_user_template::types::arg::IndividualUserTemplateInitArgs, platform_orchestrator,
     },
-    common::types::known_principal::KnownPrincipalType,
+    common::{types::known_principal::KnownPrincipalType, utils::task::run_task_concurrently},
     constant::{
-        INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT, SUBNET_ORCHESTRATOR_CANISTER_CYCLES_THRESHOLD,
+        EMPTY_CANISTER_RECHARGE_AMOUNT, INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT,
+        SUBNET_ORCHESTRATOR_CANISTER_CYCLES_THRESHOLD,
     },
     cycles::calculate_threshold_and_recharge_cycles_for_canister,
 };
@@ -72,16 +73,39 @@ pub async fn create_empty_user_canister() -> Principal {
     let _ = check_and_request_cycles_from_platform_orchestrator().await;
 
     // * provisioned canister
-    let canister_id: Principal =
-        main::create_canister(arg, INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT)
-            .await
-            .unwrap()
-            .0
-            .canister_id;
+    let canister_id: Principal = main::create_canister(arg, EMPTY_CANISTER_RECHARGE_AMOUNT)
+        .await
+        .unwrap()
+        .0
+        .canister_id;
 
     canister_id
 }
 
+pub async fn provision_number_of_empty_canisters(
+    number_of_canisters: u64,
+    breaking_condition: impl Fn() -> bool,
+) {
+    let create_canister_futures = (0..number_of_canisters).map(|_| {
+        let future = create_empty_user_canister();
+        future
+    });
+
+    let result_callback = |canister_id: Principal| {
+        CANISTER_DATA.with_borrow_mut(|canister_data| {
+            canister_data.backup_canister_pool.insert(canister_id)
+        });
+    };
+
+    run_task_concurrently(
+        create_canister_futures.into_iter(),
+        10,
+        result_callback,
+        breaking_condition,
+    )
+    .await;
+}
+
 pub async fn install_canister_wasm(
     canister_id: Principal,
     profile_owner: Option<Principal>,

src/lib/integration_tests/tests/platform_orchestrator/main.rs

@@ -1,4 +1,5 @@
 pub mod known_principal;
+pub mod provision_empty_canisters_in_a_subnet_test;
 pub mod provision_subnet_orchestrator_test;
 pub mod recharge_subnet_orchestrator_test;
 pub mod register_and_deregister_new_subnet_orchestrator_test;

src/lib/integration_tests/tests/platform_orchestrator/provision_empty_canisters_in_a_subnet_test.rs

@@ -0,0 +1,119 @@
+use candid::{encode_args, encode_one, Principal};
+use pocket_ic::WasmResult;
+use shared_utils::{
+    common::types::known_principal::KnownPrincipalType,
+    constant::TEST_BACKUP_INDIVIDUAL_USER_CANISTER_BATCH_SIZE,
+};
+use test_utils::setup::{
+    env::pocket_ic_env::get_new_pocket_ic_env, test_constants::get_mock_user_charlie_principal_id,
+};
+
+#[test]
+fn provision_empty_canisters_in_a_subnet_test() {
+    let (pocket_ic, known_principal) = get_new_pocket_ic_env();
+    let platform_canister_id = known_principal
+        .get(&KnownPrincipalType::CanisterIdPlatformOrchestrator)
+        .cloned()
+        .unwrap();
+
+    let super_admin = known_principal
+        .get(&KnownPrincipalType::UserIdGlobalSuperAdmin)
+        .cloned()
+        .unwrap();
+
+    let application_subnets = pocket_ic.topology().get_app_subnets();
+
+    let charlie_global_admin = get_mock_user_charlie_principal_id();
+
+    pocket_ic
+        .update_call(
+            platform_canister_id,
+            super_admin,
+            "add_principal_as_global_admin",
+            candid::encode_one(charlie_global_admin).unwrap(),
+        )
+        .unwrap();
+
+    let subnet_orchestrator_canister_id: Principal = pocket_ic
+        .update_call(
+            platform_canister_id,
+            charlie_global_admin,
+            "provision_subnet_orchestrator_canister",
+            candid::encode_one(application_subnets[0]).unwrap(),
+        )
+        .map(|res| {
+            let canister_id_result: Result<Principal, String> = match res {
+                WasmResult::Reply(payload) => candid::decode_one(&payload).unwrap(),
+                _ => panic!("Canister call failed"),
+            };
+            canister_id_result.unwrap()
+        })
+        .unwrap();
+
+    for i in 0..110 {
+        pocket_ic.tick();
+    }
+
+    let empty_canisters_cnt = pocket_ic
+        .query_call(
+            subnet_orchestrator_canister_id,
+            super_admin,
+            "get_subnet_backup_capacity",
+            encode_one(()).unwrap(),
+        )
+        .map(|reply_payload| {
+            let subnet_capacity: u64 = match reply_payload {
+                WasmResult::Reply(payload) => candid::decode_one(&payload).unwrap(),
+                _ => panic!("\n🛑 get_subnet_backup_capacity failed\n"),
+            };
+            subnet_capacity
+        })
+        .unwrap();
+
+    assert_eq!(
+        empty_canisters_cnt,
+        TEST_BACKUP_INDIVIDUAL_USER_CANISTER_BATCH_SIZE
+    );
+
+    pocket_ic
+        .update_call(
+            platform_canister_id,
+            charlie_global_admin,
+            "provision_empty_canisters_in_a_subnet",
+            encode_args((subnet_orchestrator_canister_id, 100_u64)).unwrap(),
+        )
+        .map(|reply_payload| {
+            let res: Result<(), String> = match reply_payload {
+                WasmResult::Reply(payload) => candid::decode_one(&payload).unwrap(),
+                _ => panic!("\n🛑 provision_empty_canisters_in_a_subnet failed\n"),
+            };
+            res
+        })
+        .unwrap()
+        .unwrap();
+
+    for _ in 0..110 {
+        pocket_ic.tick();
+    }
+
+    let empty_canisters_cnt = pocket_ic
+        .query_call(
+            subnet_orchestrator_canister_id,
+            super_admin,
+            "get_subnet_backup_capacity",
+            encode_one(()).unwrap(),
+        )
+        .map(|reply_payload| {
+            let subnet_capacity: u64 = match reply_payload {
+                WasmResult::Reply(payload) => candid::decode_one(&payload).unwrap(),
+                _ => panic!("\n🛑 get_subnet_backup_capacity failed\n"),
+            };
+            subnet_capacity
+        })
+        .unwrap();
+
+    assert_eq!(
+        empty_canisters_cnt,
+        TEST_BACKUP_INDIVIDUAL_USER_CANISTER_BATCH_SIZE + 100
+    );
+}

src/lib/shared_utils/src/constant.rs

@@ -3,6 +3,7 @@ use std::cell::RefCell;
 use crate::common::types::known_principal::{KnownPrincipalMap, KnownPrincipalType};
 use candid::Principal;
 pub const INDIVIDUAL_USER_CANISTER_RECHARGE_AMOUNT: u128 = 1_000_000_000_000; // 1T Cycles
+pub const EMPTY_CANISTER_RECHARGE_AMOUNT: u128 = 300_000_000_000; //0.3T Cycles
 pub const CYCLES_THRESHOLD_TO_INITIATE_RECHARGE: u128 = 100_000_000_000; // 0.1T Cycles
 
 pub const SUBNET_ORCHESTRATOR_CANISTER_INITIAL_CYCLES: u128 = 2_500_000_000_000_000; //2.5kT Cycles
@@ -24,8 +25,8 @@ const INDIVIDUAL_USER_CANISTER_SUBNET_THRESHOLD: u64 = 1_000;
 
 pub const TEST_BACKUP_INDIVIDUAL_USER_CANISTER_BATCH_SIZE: u64 = 10;
 pub const TEST_INDIVIDUAL_USER_CANISTER_SUBNET_BATCH_SIZE: u64 = 120;
-pub const TEST_BACKUP_INDIVIDUAL_USER_CANISTER_THRESHOLD: u64 = 1;
-pub const TEST_INDIVIDUAL_USER_CANISTER_SUBNET_THRESHOLD: u64 = 5;
+pub const TEST_BACKUP_INDIVIDUAL_USER_CANISTER_THRESHOLD: u64 = 5;
+pub const TEST_INDIVIDUAL_USER_CANISTER_SUBNET_THRESHOLD: u64 = 10;
 
 pub const INDIVIDUAL_USER_CANISTER_SUBNET_MAX_CAPACITY: u64 = 50_000;