canic/memory/
registry.rs

1use crate::{
2    CRATE_NAME, Error,
3    cdk::structures::{
4        BTreeMap as StableBTreeMap, DefaultMemoryImpl,
5        memory::{MemoryId, VirtualMemory},
6    },
7    impl_storable_bounded, log,
8    memory::{
9        CANIC_MEMORY_MAX, CANIC_MEMORY_MIN, MEMORY_MANAGER, MEMORY_RANGES_ID, MEMORY_REGISTRY_ID,
10        MemoryError,
11    },
12    types::BoundedString256,
13    utils::time::now_secs,
14};
15use candid::CandidType;
16use serde::{Deserialize, Serialize};
17use std::cell::RefCell;
18use thiserror::Error as ThisError;
19
20//
21// MEMORY_REGISTRY
22//
23
24thread_local! {
25    static MEMORY_REGISTRY: RefCell<StableBTreeMap<u8, MemoryRegistryEntry, VirtualMemory<DefaultMemoryImpl>>> =
26        RefCell::new(StableBTreeMap::init(
27            MEMORY_MANAGER.with_borrow(|this| {
28                this.get(MemoryId::new(MEMORY_REGISTRY_ID))
29            }),
30        ));
31}
32
33//
34// MEMORY_RANGES
35//
36
37thread_local! {
38    static MEMORY_RANGES: RefCell<StableBTreeMap<String, MemoryRange, VirtualMemory<DefaultMemoryImpl>>> =
39        RefCell::new(StableBTreeMap::init(
40            MEMORY_MANAGER.with_borrow(|mgr| {
41                mgr.get(MemoryId::new(MEMORY_RANGES_ID))
42            }),
43        ));
44}
45
46//
47// PENDING_REGISTRATIONS
48//
49// Queue of memory registrations produced during TLS initialization
50// Each entry is (id, crate_name, label).
51// These are deferred until `flush_pending_registrations()` is called,
52// which validates and inserts them into the global MemoryRegistry.
53//
54
55thread_local! {
56    static PENDING_REGISTRATIONS: RefCell<Vec<(u8, &'static str, &'static str)>> = const {
57        RefCell::new(Vec::new())
58    };
59}
60
61pub fn defer_register(id: u8, crate_name: &'static str, label: &'static str) {
62    PENDING_REGISTRATIONS.with(|q| {
63        q.borrow_mut().push((id, crate_name, label));
64    });
65}
66
67//
68// PENDING_RANGES
69//
70
71thread_local! {
72    pub static PENDING_RANGES: RefCell<Vec<(&'static str, u8, u8)>> = const {
73        RefCell::new(Vec::new())
74    };
75}
76
77pub fn defer_reserve_range(crate_name: &'static str, start: u8, end: u8) {
78    PENDING_RANGES.with(|q| q.borrow_mut().push((crate_name, start, end)));
79}
80
81/// Initialize all registered memory segments.
82///
83/// This should be called once during `init` or `post_upgrade`
84/// to populate the global `MemoryRegistry`.
85///
86/// Panics if called more than once or if duplicate memory IDs exist
87pub fn init_memory() {
88    // reserve internal range
89    MemoryRegistry::reserve_range(CRATE_NAME, CANIC_MEMORY_MIN, CANIC_MEMORY_MAX).unwrap();
90
91    // FIRST: flush all pending ranges
92    PENDING_RANGES.with(|q| {
93        for (crate_name, start, end) in q.borrow_mut().drain(..) {
94            MemoryRegistry::reserve_range(crate_name, start, end).unwrap();
95        }
96    });
97
98    // THEN: flush all pending registrations
99    PENDING_REGISTRATIONS.with(|q| {
100        let mut regs = q.borrow_mut();
101        regs.sort_by_key(|(id, _, _)| *id);
102        for (id, crate_name, label) in regs.drain(..) {
103            MemoryRegistry::register(id, crate_name, label).unwrap();
104        }
105    });
106
107    // summary logs: one per range
108    MEMORY_RANGES.with_borrow(|ranges| {
109        MEMORY_REGISTRY.with_borrow(|reg| {
110            // get range entries
111            let mut entries: Vec<_> = ranges.iter().collect();
112            entries.sort_by_key(|entry| entry.value().start);
113
114            for entry in entries {
115                let crate_name = entry.key();
116                let range = entry.value();
117
118                let count = reg.iter().filter(|e| range.contains(*e.key())).count();
119
120                log!(
121                    Info,
122                    "💾 memory.range: {} [{}-{}] ({}/{} slots used)",
123                    crate_name,
124                    range.start,
125                    range.end,
126                    count,
127                    range.end - range.start + 1,
128                );
129            }
130        });
131    });
132}
133
134///
135/// MemoryRegistryError
136///
137
138#[derive(Debug, ThisError)]
139pub enum MemoryRegistryError {
140    #[error("ID {0} is already registered with type {1}, tried to register type {2}")]
141    AlreadyRegistered(u8, String, String),
142
143    #[error("crate `{0}` already has a reserved range")]
144    DuplicateRange(String),
145
146    #[error("crate `{0}` attempted to register ID {1}, but it is outside its allowed ranges")]
147    OutOfRange(String, u8),
148
149    #[error("crate `{0}` range {1}-{2} overlaps with crate `{3}` range {4}-{5}")]
150    Overlap(String, u8, u8, String, u8, u8),
151
152    #[error("crate `{0}` has not reserved any memory range")]
153    NoRange(String),
154}
155
156impl From<MemoryRegistryError> for Error {
157    fn from(err: MemoryRegistryError) -> Self {
158        MemoryError::from(err).into()
159    }
160}
161
162///
163/// MemoryRange
164///
165
166#[derive(Clone, Debug, Deserialize, Serialize)]
167pub struct MemoryRange {
168    pub crate_key: BoundedString256,
169    pub start: u8,
170    pub end: u8,
171    pub created_at: u64,
172}
173
174impl MemoryRange {
175    #[must_use]
176    pub fn new(crate_key: &str, start: u8, end: u8) -> Self {
177        Self {
178            crate_key: BoundedString256::new(crate_key),
179            start,
180            end,
181            created_at: now_secs(),
182        }
183    }
184
185    #[must_use]
186    pub fn contains(&self, id: u8) -> bool {
187        (self.start..=self.end).contains(&id)
188    }
189}
190
191impl_storable_bounded!(MemoryRange, 320, false);
192
193///
194/// MemoryRegistryEntry
195///
196
197#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
198pub struct MemoryRegistryEntry {
199    pub label: BoundedString256,
200    pub created_at: u64,
201}
202
203impl MemoryRegistryEntry {
204    #[must_use]
205    pub fn new(label: &str) -> Self {
206        Self {
207            label: BoundedString256::new(label),
208            created_at: now_secs(),
209        }
210    }
211}
212
213impl_storable_bounded!(MemoryRegistryEntry, 320, false);
214
215///
216/// MemoryRegistry
217///
218
219pub struct MemoryRegistry;
220
221pub type MemoryRegistryView = Vec<(u8, MemoryRegistryEntry)>;
222
223impl MemoryRegistry {
224    #[must_use]
225    pub fn is_empty() -> bool {
226        MEMORY_REGISTRY.with_borrow(|map| map.is_empty())
227    }
228
229    /// Register an ID, enforcing crate’s allowed range.
230    fn register(id: u8, crate_name: &str, label: &str) -> Result<(), Error> {
231        // convert to string once
232        let crate_key = crate_name.to_string();
233
234        // immutable borrow first: check ranges and existing registry entry
235        let range = MEMORY_RANGES.with_borrow(|ranges| ranges.get(&crate_key));
236        match range {
237            None => {
238                return Err(MemoryRegistryError::NoRange(crate_key))?;
239            }
240            Some(r) if !r.contains(id) => {
241                return Err(MemoryRegistryError::OutOfRange(crate_key, id))?;
242            }
243            Some(_) => {
244                // OK, continue
245            }
246        }
247
248        // check already registered
249        let existing = MEMORY_REGISTRY.with_borrow(|map| map.get(&id));
250        if let Some(existing) = existing {
251            if existing.label.as_ref() != label {
252                return Err(MemoryRegistryError::AlreadyRegistered(
253                    id,
254                    existing.label.to_string(),
255                    label.to_string(),
256                ))?;
257            }
258            return Ok(()); // idempotent case
259        }
260
261        // now borrow mutably for insertion
262        MEMORY_REGISTRY.with_borrow_mut(|map| {
263            map.insert(id, MemoryRegistryEntry::new(label));
264        });
265
266        Ok(())
267    }
268
269    /// Reserve a block of memory IDs for a crate.
270    fn reserve_range(crate_name: &str, start: u8, end: u8) -> Result<(), Error> {
271        if start > end {
272            return Err(MemoryRegistryError::OutOfRange(
273                crate_name.to_string(),
274                start,
275            ))?;
276        }
277
278        // convert to string once; use this everywhere.
279        let crate_key = crate_name.to_string();
280
281        // immutable borrow first
282        let conflict = MEMORY_RANGES.with_borrow(|ranges| {
283            if ranges.contains_key(&crate_key) {
284                return Some(MemoryRegistryError::DuplicateRange(crate_key.clone()));
285            }
286
287            for entry in ranges.iter() {
288                let other_crate = entry.key();
289                let other_range = entry.value();
290
291                if !(end < other_range.start || start > other_range.end) {
292                    return Some(MemoryRegistryError::Overlap(
293                        crate_key.clone(),
294                        start,
295                        end,
296                        other_crate.clone(),
297                        other_range.start,
298                        other_range.end,
299                    ));
300                }
301            }
302
303            None
304        });
305
306        if let Some(err) = conflict {
307            return Err(err)?;
308        }
309
310        // now borrow mutably once for insertion
311        MEMORY_RANGES.with_borrow_mut(|ranges| {
312            let range = MemoryRange::new(crate_name, start, end);
313            ranges.insert(crate_name.to_string(), range);
314        });
315
316        Ok(())
317    }
318
319    #[must_use]
320    pub fn get(id: u8) -> Option<MemoryRegistryEntry> {
321        MEMORY_REGISTRY.with_borrow(|map| map.get(&id))
322    }
323
324    #[must_use]
325    pub fn export() -> MemoryRegistryView {
326        MEMORY_REGISTRY.with_borrow(|map| {
327            map.iter()
328                .map(|entry| (*entry.key(), entry.value()))
329                .collect()
330        })
331    }
332
333    #[must_use]
334    pub fn export_ranges() -> Vec<(String, MemoryRange)> {
335        MEMORY_RANGES.with_borrow(|ranges| {
336            ranges
337                .iter()
338                .map(|e| (e.key().clone(), e.value()))
339                .collect()
340        })
341    }
342
343    pub fn clear() {
344        MEMORY_REGISTRY.with_borrow_mut(StableBTreeMap::clear);
345        MEMORY_RANGES.with_borrow_mut(StableBTreeMap::clear);
346    }
347}
canic/memory/registry.rs

canic/memory/
registry.rs
