canic_template_runtime/storage/template/
chunked.rs

1use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, memory::VirtualMemory};
2use canic_memory::{eager_static, ic_memory, impl_storable_unbounded};
3use canic_template_types::ids::{TemplateChunkKey, TemplateReleaseKey};
4use serde::{Deserialize, Serialize};
5use std::{cell::RefCell, collections::BTreeMap as StdBTreeMap};
6
7const TEMPLATE_CHUNK_SETS_ID: u8 = 11;
8const TEMPLATE_CHUNKS_ID: u8 = 12;
9
10eager_static! {
11    static TEMPLATE_CHUNK_SETS: RefCell<
12        BTreeMap<TemplateReleaseKey, TemplateChunkSetRecord, VirtualMemory<DefaultMemoryImpl>>
13    > = RefCell::new(
14        BTreeMap::init(ic_memory!(TemplateChunkSetStateStore, TEMPLATE_CHUNK_SETS_ID)),
15    );
16}
17
18eager_static! {
19    static TEMPLATE_CHUNKS: RefCell<
20        BTreeMap<TemplateChunkKey, TemplateChunkRecord, VirtualMemory<DefaultMemoryImpl>>
21    > = RefCell::new(
22        BTreeMap::init(ic_memory!(TemplateChunkStore, TEMPLATE_CHUNKS_ID)),
23    );
24}
25
26///
27/// TemplateChunkSetRecord
28///
29
30#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
31pub struct TemplateChunkSetRecord {
32    pub payload_hash: Vec<u8>,
33    pub payload_size_bytes: u64,
34    pub chunk_count: u32,
35    pub chunk_hashes: Vec<Vec<u8>>,
36    pub created_at: u64,
37}
38
39impl_storable_unbounded!(TemplateChunkSetRecord);
40
41///
42/// TemplateChunkRecord
43///
44
45#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
46pub struct TemplateChunkRecord {
47    pub bytes: Vec<u8>,
48}
49
50impl_storable_unbounded!(TemplateChunkRecord);
51
52///
53/// TemplateChunkSetStateStore
54///
55
56pub struct TemplateChunkSetStateStore;
57
58impl TemplateChunkSetStateStore {
59    // Insert or replace one template chunk-set metadata record.
60    pub fn upsert(release: TemplateReleaseKey, record: TemplateChunkSetRecord) {
61        TEMPLATE_CHUNK_SETS.with_borrow_mut(|map| {
62            map.insert(release, record);
63        });
64    }
65
66    // Fetch one template chunk-set metadata record, if present.
67    #[must_use]
68    pub fn get(release: &TemplateReleaseKey) -> Option<TemplateChunkSetRecord> {
69        TEMPLATE_CHUNK_SETS.with_borrow(|map| map.get(release))
70    }
71
72    // Export the full chunk-set metadata snapshot for ops-owned accounting.
73    #[must_use]
74    pub fn export() -> Vec<(TemplateReleaseKey, TemplateChunkSetRecord)> {
75        TEMPLATE_CHUNK_SETS.with_borrow(|map| {
76            map.iter()
77                .map(|entry| (entry.key().clone(), entry.value()))
78                .collect()
79        })
80    }
81
82    // Clear the chunk-set metadata store.
83    pub fn clear() {
84        TEMPLATE_CHUNK_SETS.with_borrow_mut(BTreeMap::clear);
85    }
86
87    // Clear the chunk-set metadata store for isolated unit tests.
88    pub fn clear_for_test() {
89        Self::clear();
90    }
91}
92
93///
94/// TemplateChunkStore
95///
96
97pub struct TemplateChunkStore;
98
99impl TemplateChunkStore {
100    // Insert or replace one template chunk.
101    pub fn upsert(chunk_key: TemplateChunkKey, record: TemplateChunkRecord) {
102        TEMPLATE_CHUNKS.with_borrow_mut(|map| {
103            map.insert(chunk_key, record);
104        });
105    }
106
107    // Fetch one template chunk, if present.
108    #[must_use]
109    pub fn get(chunk_key: &TemplateChunkKey) -> Option<TemplateChunkRecord> {
110        TEMPLATE_CHUNKS.with_borrow(|map| map.get(chunk_key))
111    }
112
113    // Export the full chunk snapshot for ops-owned accounting.
114    #[must_use]
115    pub fn export() -> Vec<(TemplateChunkKey, TemplateChunkRecord)> {
116        TEMPLATE_CHUNKS.with_borrow(|map| {
117            map.iter()
118                .map(|entry| (entry.key().clone(), entry.value()))
119                .collect()
120        })
121    }
122
123    // Count staged chunks by release without cloning chunk payload bytes.
124    #[must_use]
125    pub fn count_by_release() -> StdBTreeMap<TemplateReleaseKey, u32> {
126        TEMPLATE_CHUNKS.with_borrow(|map| {
127            let mut counts: StdBTreeMap<TemplateReleaseKey, u32> = StdBTreeMap::new();
128
129            for entry in map.iter() {
130                let release = entry.key().release.clone();
131                let count = counts.entry(release).or_insert(0);
132                *count = u32::saturating_add(*count, 1);
133            }
134
135            counts
136        })
137    }
138
139    // Clear the chunk store.
140    pub fn clear() {
141        TEMPLATE_CHUNKS.with_borrow_mut(BTreeMap::clear);
142    }
143
144    // Clear the chunk store for isolated unit tests.
145    pub fn clear_for_test() {
146        Self::clear();
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use canic_template_types::ids::{TemplateId, TemplateVersion};
154
155    fn release() -> TemplateReleaseKey {
156        TemplateReleaseKey::new(
157            TemplateId::new("embedded:app"),
158            TemplateVersion::new("0.18.0"),
159        )
160    }
161
162    #[test]
163    fn chunk_set_store_round_trip() {
164        TemplateChunkSetStateStore::clear_for_test();
165        let record = TemplateChunkSetRecord {
166            payload_hash: vec![1; 32],
167            payload_size_bytes: 7,
168            chunk_count: 2,
169            chunk_hashes: vec![vec![2; 32], vec![3; 32]],
170            created_at: 99,
171        };
172
173        TemplateChunkSetStateStore::upsert(release(), record.clone());
174
175        assert_eq!(TemplateChunkSetStateStore::get(&release()), Some(record));
176    }
177
178    #[test]
179    fn chunk_store_round_trip() {
180        TemplateChunkStore::clear_for_test();
181        let chunk_key = TemplateChunkKey::new(release(), 0);
182        let record = TemplateChunkRecord {
183            bytes: vec![1, 2, 3],
184        };
185
186        TemplateChunkStore::upsert(chunk_key.clone(), record.clone());
187
188        assert_eq!(TemplateChunkStore::get(&chunk_key), Some(record));
189    }
190}
canic_template_runtime/storage/template/chunked.rs

canic_template_runtime/storage/template/
chunked.rs
