canic_core/model/memory/
reserve.rs

1use crate::{
2    cdk::{
3        structures::{BTreeMap, DefaultMemoryImpl, memory::VirtualMemory},
4        types::Principal,
5        utils::time::now_secs,
6    },
7    eager_static, ic_memory,
8    ids::CanisterRole,
9    memory::impl_storable_unbounded,
10    model::memory::id::root::CANISTER_RESERVE_ID,
11    types::Cycles,
12};
13use candid::CandidType;
14use serde::{Deserialize, Serialize};
15use std::cell::RefCell;
16
17//
18// CANISTER_RESERVE
19//
20
21eager_static! {
22    static CANISTER_RESERVE: RefCell<BTreeMap<Principal, CanisterReserveEntry, VirtualMemory<DefaultMemoryImpl>>> =
23        RefCell::new(BTreeMap::init(
24            ic_memory!(CanisterReserve, CANISTER_RESERVE_ID),
25        ));
26}
27
28///
29/// CanisterReserveView
30///
31
32pub type CanisterReserveView = Vec<(Principal, CanisterReserveEntry)>;
33
34///
35/// CanisterReserveEntry
36///
37
38#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
39pub struct CanisterReserveEntry {
40    pub created_at: u64,
41    pub cycles: Cycles,
42    #[serde(default)]
43    pub ty: Option<CanisterRole>,
44    #[serde(default)]
45    pub parent: Option<Principal>,
46    #[serde(default)]
47    pub module_hash: Option<Vec<u8>>,
48}
49
50impl_storable_unbounded!(CanisterReserveEntry);
51
52///
53/// CanisterReserve
54///
55
56pub(crate) struct CanisterReserve;
57
58impl CanisterReserve {
59    /// Register a canister into the reserve.
60    pub(crate) fn register(
61        pid: Principal,
62        cycles: Cycles,
63        ty: Option<CanisterRole>,
64        parent: Option<Principal>,
65        module_hash: Option<Vec<u8>>,
66    ) {
67        let entry = CanisterReserveEntry {
68            created_at: now_secs(),
69            cycles,
70            ty,
71            parent,
72            module_hash,
73        };
74
75        CANISTER_RESERVE.with_borrow_mut(|map| {
76            map.insert(pid, entry);
77        });
78    }
79
80    /// Pop the oldest canister from the reserve.
81    #[must_use]
82    pub(crate) fn pop_first() -> Option<(Principal, CanisterReserveEntry)> {
83        CANISTER_RESERVE.with_borrow_mut(|map| {
84            let min_pid = map
85                .iter()
86                .min_by_key(|entry| entry.value().created_at)
87                .map(|entry| *entry.key())?;
88            map.remove(&min_pid).map(|entry| (min_pid, entry))
89        })
90    }
91
92    /// Return true if the reserve contains the given canister.
93    #[must_use]
94    pub(crate) fn contains(pid: &Principal) -> bool {
95        CANISTER_RESERVE.with_borrow(|map| map.contains_key(pid))
96    }
97
98    /// Remove a specific canister from the reserve, returning its entry.
99    #[must_use]
100    pub(crate) fn take(pid: &Principal) -> Option<CanisterReserveEntry> {
101        CANISTER_RESERVE.with_borrow_mut(|map| map.remove(pid))
102    }
103
104    /// Remove a specific canister from the reserve.
105    #[must_use]
106    #[cfg(test)]
107    pub(crate) fn remove(pid: &Principal) -> Option<CanisterReserveEntry> {
108        CANISTER_RESERVE.with_borrow_mut(|map| map.remove(pid))
109    }
110
111    /// Export the reserve as a vector of (Principal, Entry).
112    #[must_use]
113    pub(crate) fn export() -> CanisterReserveView {
114        CANISTER_RESERVE.with_borrow(BTreeMap::to_vec)
115    }
116
117    /// Clear the reserve (mainly for tests).
118    #[cfg(test)]
119    pub(crate) fn clear() {
120        CANISTER_RESERVE.with_borrow_mut(BTreeMap::clear);
121    }
122
123    /// Return the current reserve size.
124    #[must_use]
125    pub(crate) fn len() -> u64 {
126        CANISTER_RESERVE.with_borrow(|map| map.len())
127    }
128
129    /// Return whether the reserve is empty.
130    #[must_use]
131    #[cfg(test)]
132    pub(crate) fn is_empty() -> bool {
133        CANISTER_RESERVE.with_borrow(|map| map.is_empty())
134    }
135}
136
137///
138/// TESTS
139///
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use candid::Principal;
145
146    fn pid(n: u8) -> Principal {
147        Principal::self_authenticating(vec![n])
148    }
149
150    #[test]
151    fn register_and_export() {
152        CanisterReserve::clear();
153
154        let p1 = pid(1);
155        let p2 = pid(2);
156
157        CanisterReserve::register(p1, 100u128.into(), None, None, None);
158        CanisterReserve::register(p2, 200u128.into(), None, None, None);
159
160        let view = CanisterReserve::export();
161        assert_eq!(view.len(), 2);
162
163        let entry1 = view.iter().find(|(id, _)| *id == p1).unwrap();
164        assert_eq!(entry1.1.cycles, 100u128.into());
165
166        let entry2 = view.iter().find(|(id, _)| *id == p2).unwrap();
167        assert_eq!(entry2.1.cycles, 200u128.into());
168    }
169
170    #[test]
171    fn remove_specific_pid() {
172        CanisterReserve::clear();
173
174        let p1 = pid(1);
175        let p2 = pid(2);
176
177        CanisterReserve::register(p1, 123u128.into(), None, None, None);
178        CanisterReserve::register(p2, 456u128.into(), None, None, None);
179
180        let removed = CanisterReserve::remove(&p1).unwrap();
181        assert_eq!(removed.cycles, 123u128.into());
182
183        // only p2 should remain
184        let view = CanisterReserve::export();
185        assert_eq!(view.len(), 1);
186        assert_eq!(view[0].0, p2);
187    }
188
189    #[test]
190    fn clear_resets_reserve() {
191        CanisterReserve::clear();
192
193        CanisterReserve::register(pid(1), 10u128.into(), None, None, None);
194        assert!(!CanisterReserve::is_empty());
195
196        CanisterReserve::clear();
197        assert!(CanisterReserve::is_empty());
198    }
199}
canic_core/model/memory/reserve.rs

canic_core/model/memory/
reserve.rs
