canic_core/model/memory/
reserve.rs

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

canic_core/model/memory/
reserve.rs
