canic_core/model/memory/
pool.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_POOL_ID,
11    types::Cycles,
12};
13use candid::CandidType;
14use serde::{Deserialize, Serialize};
15use std::cell::RefCell;
16
17//
18// CANISTER_POOL
19//
20
21eager_static! {
22    static CANISTER_POOL: RefCell<BTreeMap<Principal, CanisterPoolEntry, VirtualMemory<DefaultMemoryImpl>>> =
23        RefCell::new(BTreeMap::init(
24            ic_memory!(CanisterPool, CANISTER_POOL_ID),
25        ));
26}
27
28///
29/// CanisterPoolView
30///
31
32pub type CanisterPoolView = Vec<(Principal, CanisterPoolEntry)>;
33
34///
35/// CanisterPoolStatus
36///
37
38#[derive(CandidType, Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
39pub enum CanisterPoolStatus {
40    PendingReset,
41    #[default]
42    Ready,
43    Failed {
44        reason: String,
45    },
46}
47
48impl CanisterPoolStatus {
49    #[must_use]
50    pub const fn is_pending_reset(&self) -> bool {
51        matches!(self, Self::PendingReset)
52    }
53
54    #[must_use]
55    pub const fn is_ready(&self) -> bool {
56        matches!(self, Self::Ready)
57    }
58
59    #[must_use]
60    pub const fn is_failed(&self) -> bool {
61        matches!(self, Self::Failed { .. })
62    }
63}
64
65///
66/// CanisterPoolEntry
67///
68
69#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
70pub struct CanisterPoolEntry {
71    pub created_at: u64,
72    pub cycles: Cycles,
73    #[serde(default)]
74    pub status: CanisterPoolStatus,
75    #[serde(default)]
76    pub role: Option<CanisterRole>,
77    #[serde(default)]
78    pub parent: Option<Principal>,
79    #[serde(default)]
80    pub module_hash: Option<Vec<u8>>,
81}
82
83impl_storable_unbounded!(CanisterPoolEntry);
84
85///
86/// CanisterPool
87///
88
89pub struct CanisterPool;
90
91impl CanisterPool {
92    /// Register a canister into the pool.
93    pub fn register(
94        pid: Principal,
95        cycles: Cycles,
96        status: CanisterPoolStatus,
97        role: Option<CanisterRole>,
98        parent: Option<Principal>,
99        module_hash: Option<Vec<u8>>,
100    ) {
101        let entry = CanisterPoolEntry {
102            created_at: now_secs(),
103            cycles,
104            status,
105            role,
106            parent,
107            module_hash,
108        };
109
110        CANISTER_POOL.with_borrow_mut(|map| {
111            map.insert(pid, entry);
112        });
113    }
114
115    #[must_use]
116    pub(crate) fn get(pid: Principal) -> Option<CanisterPoolEntry> {
117        CANISTER_POOL.with_borrow(|map| map.get(&pid))
118    }
119
120    #[must_use]
121    pub(crate) fn update(pid: Principal, entry: CanisterPoolEntry) -> bool {
122        CANISTER_POOL.with_borrow_mut(|map| {
123            if map.contains_key(&pid) {
124                map.insert(pid, entry);
125                true
126            } else {
127                false
128            }
129        })
130    }
131
132    /// Pop the oldest ready canister from the pool.
133    #[must_use]
134    pub(crate) fn pop_ready() -> Option<(Principal, CanisterPoolEntry)> {
135        CANISTER_POOL.with_borrow_mut(|map| {
136            let min_pid = map
137                .iter()
138                .filter(|entry| entry.value().status.is_ready())
139                .min_by_key(|entry| entry.value().created_at)
140                .map(|entry| *entry.key())?;
141            map.remove(&min_pid).map(|entry| (min_pid, entry))
142        })
143    }
144
145    /// Return true if the pool contains the given canister.
146    #[must_use]
147    pub(crate) fn contains(pid: &Principal) -> bool {
148        CANISTER_POOL.with_borrow(|map| map.contains_key(pid))
149    }
150
151    /// Remove a specific canister from the pool, returning its entry.
152    #[must_use]
153    pub(crate) fn take(pid: &Principal) -> Option<CanisterPoolEntry> {
154        CANISTER_POOL.with_borrow_mut(|map| map.remove(pid))
155    }
156
157    /// Remove a specific canister from the pool.
158    #[must_use]
159    #[cfg(test)]
160    pub(crate) fn remove(pid: &Principal) -> Option<CanisterPoolEntry> {
161        CANISTER_POOL.with_borrow_mut(|map| map.remove(pid))
162    }
163
164    /// Export the pool as a vector of (Principal, Entry).
165    #[must_use]
166    pub(crate) fn export() -> CanisterPoolView {
167        CANISTER_POOL.with_borrow(BTreeMap::to_vec)
168    }
169
170    /// Clear the pool (mainly for tests).
171    #[cfg(test)]
172    pub(crate) fn clear() {
173        CANISTER_POOL.with_borrow_mut(BTreeMap::clear);
174    }
175
176    /// Return the current pool size.
177    #[must_use]
178    pub(crate) fn len() -> u64 {
179        CANISTER_POOL.with_borrow(|map| map.len())
180    }
181
182    /// Return whether the pool is empty.
183    #[must_use]
184    #[cfg(test)]
185    pub(crate) fn is_empty() -> bool {
186        CANISTER_POOL.with_borrow(|map| map.is_empty())
187    }
188}
189
190///
191/// TESTS
192///
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use candid::Principal;
198
199    fn pid(n: u8) -> Principal {
200        Principal::self_authenticating(vec![n])
201    }
202
203    #[test]
204    fn register_and_export() {
205        CanisterPool::clear();
206
207        let p1 = pid(1);
208        let p2 = pid(2);
209
210        CanisterPool::register(
211            p1,
212            100u128.into(),
213            CanisterPoolStatus::Ready,
214            None,
215            None,
216            None,
217        );
218        CanisterPool::register(
219            p2,
220            200u128.into(),
221            CanisterPoolStatus::Ready,
222            None,
223            None,
224            None,
225        );
226
227        let view = CanisterPool::export();
228        assert_eq!(view.len(), 2);
229
230        let entry1 = view.iter().find(|(id, _)| *id == p1).unwrap();
231        assert_eq!(entry1.1.cycles, 100u128.into());
232
233        let entry2 = view.iter().find(|(id, _)| *id == p2).unwrap();
234        assert_eq!(entry2.1.cycles, 200u128.into());
235    }
236
237    #[test]
238    fn remove_specific_pid() {
239        CanisterPool::clear();
240
241        let p1 = pid(1);
242        let p2 = pid(2);
243
244        CanisterPool::register(
245            p1,
246            123u128.into(),
247            CanisterPoolStatus::Ready,
248            None,
249            None,
250            None,
251        );
252        CanisterPool::register(
253            p2,
254            456u128.into(),
255            CanisterPoolStatus::Ready,
256            None,
257            None,
258            None,
259        );
260
261        let removed = CanisterPool::remove(&p1).unwrap();
262        assert_eq!(removed.cycles, 123u128.into());
263
264        // only p2 should remain
265        let view = CanisterPool::export();
266        assert_eq!(view.len(), 1);
267        assert_eq!(view[0].0, p2);
268    }
269
270    #[test]
271    fn clear_resets_pool() {
272        CanisterPool::clear();
273
274        CanisterPool::register(
275            pid(1),
276            10u128.into(),
277            CanisterPoolStatus::Ready,
278            None,
279            None,
280            None,
281        );
282        assert!(!CanisterPool::is_empty());
283
284        CanisterPool::clear();
285        assert!(CanisterPool::is_empty());
286    }
287}
canic_core/model/memory/pool.rs

canic_core/model/memory/
pool.rs
