junobuild_storage/
runtime.rs

1use crate::certification::cert::update_certified_data;
2use crate::certification::types::certified::CertifiedAssetHashes;
3use crate::memory::STATE;
4use crate::types::config::StorageConfig;
5use crate::types::runtime_state::{
6    BatchId, Batches, ChunkId, Chunks, RuntimeState, StorageRuntimeState,
7};
8use crate::types::store::{Asset, Batch, BatchExpiry, Chunk};
9use ic_cdk::api::time;
10use junobuild_collections::types::core::CollectionKey;
11use junobuild_shared::rate::types::RateConfig;
12use junobuild_shared::rate::utils::increment_and_assert_rate_store;
13use std::collections::HashMap;
14
15// ---------------------------------------------------------
16// Certified assets
17// ---------------------------------------------------------
18
19pub fn init_certified_assets(asset_hashes: &CertifiedAssetHashes) {
20    STATE.with(|state| {
21        init_certified_assets_impl(asset_hashes, &mut state.borrow_mut().runtime.storage)
22    });
23}
24
25pub fn update_certified_asset(asset: &Asset, config: &StorageConfig) {
26    STATE.with(|state| update_certified_asset_impl(asset, config, &mut state.borrow_mut().runtime));
27}
28
29pub fn delete_certified_asset(asset: &Asset) {
30    STATE.with(|state| delete_certified_asset_impl(asset, &mut state.borrow_mut().runtime));
31}
32
33fn init_certified_assets_impl(
34    asset_hashes: &CertifiedAssetHashes,
35    storage: &mut StorageRuntimeState,
36) {
37    // 1. Init all asset in tree
38    storage.asset_hashes = asset_hashes.clone();
39
40    // 2. Update the root hash and the canister certified data
41    update_certified_data(&storage.asset_hashes);
42}
43
44fn update_certified_asset_impl(asset: &Asset, config: &StorageConfig, runtime: &mut RuntimeState) {
45    // 1. Replace or insert the new asset in tree
46    runtime.storage.asset_hashes.insert(asset, config);
47
48    // 2. Update the root hash and the canister certified data
49    update_certified_data(&runtime.storage.asset_hashes);
50}
51
52fn delete_certified_asset_impl(asset: &Asset, runtime: &mut RuntimeState) {
53    // 1. Remove the asset in tree
54    runtime.storage.asset_hashes.delete(asset);
55
56    // 2. Update the root hash and the canister certified data
57    update_certified_data(&runtime.storage.asset_hashes);
58}
59
60// ---------------------------------------------------------
61// Batch
62// ---------------------------------------------------------
63
64pub fn get_batch(batch_id: &BatchId) -> Option<Batch> {
65    STATE.with(|state| {
66        state
67            .borrow()
68            .runtime
69            .storage
70            .batches
71            .get(batch_id)
72            .cloned()
73    })
74}
75
76pub fn insert_batch(batch_id: &BatchId, batch: Batch) {
77    STATE.with(|state| {
78        insert_batch_impl(
79            batch_id,
80            batch,
81            &mut state.borrow_mut().runtime.storage.batches,
82        )
83    })
84}
85
86pub fn clear_expired_batches() {
87    STATE.with(|state| clear_expired_batches_impl(&mut state.borrow_mut().runtime.storage.batches));
88}
89
90pub fn clear_batch(batch_id: &BatchId, chunk_ids: &[ChunkId]) {
91    STATE.with(|state| {
92        clear_batch_impl(batch_id, chunk_ids, &mut state.borrow_mut().runtime.storage)
93    });
94}
95
96fn insert_batch_impl(batch_id: &BatchId, batch: Batch, batches: &mut Batches) {
97    batches.insert(*batch_id, batch);
98}
99
100fn clear_expired_batches_impl(batches: &mut Batches) {
101    clear_expired_items(batches);
102}
103
104fn clear_expired_items<K, V>(batches: &mut HashMap<K, V>)
105where
106    K: Clone + Eq + std::hash::Hash,
107    V: BatchExpiry,
108{
109    let now = time();
110
111    let expired_ids: Vec<K> = batches
112        .iter()
113        .filter_map(|(id, item)| {
114            if now > item.expires_at() {
115                Some(id.clone())
116            } else {
117                None
118            }
119        })
120        .collect();
121
122    for id in expired_ids {
123        batches.remove(&id);
124    }
125}
126
127fn clear_batch_impl(batch_id: &BatchId, chunk_ids: &[ChunkId], state: &mut StorageRuntimeState) {
128    for chunk_id in chunk_ids.iter() {
129        state.chunks.remove(chunk_id);
130    }
131
132    state.batches.remove(batch_id);
133}
134
135// ---------------------------------------------------------
136// Chunks
137// ---------------------------------------------------------
138
139pub fn get_chunk(chunk_id: &ChunkId) -> Option<Chunk> {
140    STATE.with(|state| {
141        let chunks = state.borrow().runtime.storage.chunks.clone();
142        let chunk = chunks.get(chunk_id);
143        chunk.cloned()
144    })
145}
146
147pub fn clear_expired_chunks() {
148    STATE.with(|state| clear_expired_chunks_impl(&mut state.borrow_mut().runtime.storage));
149}
150
151pub fn insert_chunk(chunk_id: &ChunkId, chunk: Chunk) {
152    STATE.with(|state| {
153        insert_chunk_impl(
154            chunk_id,
155            chunk,
156            &mut state.borrow_mut().runtime.storage.chunks,
157        )
158    })
159}
160
161fn clear_expired_chunks_impl(state: &mut StorageRuntimeState) {
162    let cloned_chunks = state.chunks.clone();
163
164    for (chunk_id, chunk) in cloned_chunks.iter() {
165        if !state.batches.contains_key(&chunk.batch_id) {
166            state.chunks.remove(chunk_id);
167        }
168    }
169}
170
171fn insert_chunk_impl(chunk_id: &ChunkId, chunk: Chunk, chunks: &mut Chunks) {
172    chunks.insert(*chunk_id, chunk);
173}
174
175// ---------------------------------------------------------
176// Rates
177// ---------------------------------------------------------
178
179pub fn increment_and_assert_rate(
180    collection: &CollectionKey,
181    config: &Option<RateConfig>,
182) -> Result<(), String> {
183    STATE.with(|state| {
184        increment_and_assert_rate_store(
185            collection,
186            config,
187            &mut state.borrow_mut().runtime.storage.rate_tokens,
188        )
189    })
190}
junobuild_storage/runtime.rs

junobuild_storage/
runtime.rs
