prevent panics, add cumulative total instead of recent amount

dfinity · Sep 10, 2024 · 716fbca · 716fbca
1 parent 5af048d
commit 716fbca
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 21 deletions.
diff --git a/canfund-rs/src/manager/mod.rs b/canfund-rs/src/manager/mod.rs
@@ -271,15 +271,10 @@ impl FundManager {
                                 .await
                             {
                                 Ok(cycles_obtained) => {
-                                    manager
-                                        .borrow_mut()
-                                        .canisters
-                                        .get_mut(&canister_id)
-                                        .unwrap()
-                                        .set_deposited_cycles(CyclesBalance::new(
-                                            cycles_obtained,
-                                            time(),
-                                        ));
+                                    if let Some(record) = manager.borrow_mut().canisters.get_mut(&canister_id) {
+                                        record.add_deposited_cycles(cycles_obtained);
+                                    }
+
                                     print(format!(
                                         "Obtained {} cycles for canister {}",
                                         cycles_obtained,
@@ -323,15 +318,10 @@ impl FundManager {
                             ));
                         }
                         Ok(_) => {
-                            manager
-                                .borrow_mut()
-                                .canisters
-                                .get_mut(&canister_id)
-                                .unwrap()
-                                .set_deposited_cycles(CyclesBalance::new(
-                                    needed_cycles,
-                                    time(),
-                                ));
+                            if let Some(record) = manager.borrow_mut().canisters.get_mut(&canister_id) {
+                                record.add_deposited_cycles(needed_cycles);
+                            }
+
                             print(format!(
                                 "Funded canister {} with {} cycles",
                                 canister_id.to_text(),

diff --git a/canfund-rs/src/manager/record.rs b/canfund-rs/src/manager/record.rs
@@ -1,5 +1,7 @@
 use std::sync::Arc;
 
+use ic_cdk::api::time;
+
 use crate::operations::fetch::FetchCyclesBalance;
 
 #[derive(Clone)]
@@ -8,7 +10,7 @@ pub struct CanisterRecord {
     cycles: Option<CyclesBalance>,
     /// The canister cycles balance record when it was last funded.
     previous_cycles: Option<CyclesBalance>,
-    /// The number of cycles deposited in the last check.
+    /// The cumulative total of cycles deposited to the canister.
     deposited_cycles: Option<CyclesBalance>,
     /// The method to fetch the canister cycles balance.
     cycles_fetcher: Arc<dyn FetchCyclesBalance>,
@@ -40,8 +42,13 @@ impl CanisterRecord {
         &self.previous_cycles
     }
 
-    pub fn set_deposited_cycles(&mut self, deposited_cycles: CyclesBalance) {
-        self.deposited_cycles = Some(deposited_cycles);
+    pub fn add_deposited_cycles(&mut self, cycles: u128) {
+        if let Some(deposited_cycles) = self.deposited_cycles.as_mut() {
+            deposited_cycles.amount = deposited_cycles.amount.saturating_add(cycles);
+            deposited_cycles.timestamp = time();
+        } else {
+            self.deposited_cycles = Some(CyclesBalance::new(cycles, time()));
+        }
     }
 
     pub fn get_deposited_cycles(&self) -> &Option<CyclesBalance> {