canic_core/model/memory/
reserve.rs

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

canic_core/model/memory/
reserve.rs
