massa_module_cache/
hd_cache.rs

1use crate::types::{
2    ModuleInfo, ModuleMetadata, ModuleMetadataDeserializer, ModuleMetadataSerializer,
3};
4use massa_hash::Hash;
5use massa_sc_runtime::{CondomLimits, GasCosts, RuntimeModule};
6use massa_serialization::{DeserializeError, Deserializer, Serializer};
7use rand::RngCore;
8use rocksdb::{Direction, IteratorMode, Options, WriteBatch, DB};
9use std::path::PathBuf;
10use tracing::{debug, warn};
11
12const OPEN_ERROR: &str = "critical: rocksdb open operation failed";
13const CRUD_ERROR: &str = "critical: rocksdb crud operation failed";
14const DATA_SER_ERROR: &str = "critical: metadata serialization failed";
15const DATA_DESER_ERROR: &str = "critical: metadata deserialization failed";
16const MOD_SER_ERROR: &str = "critical: module serialization failed";
17const MOD_DESER_ERROR: &str = "critical: module deserialization failed";
18const MODULE_IDENT: u8 = 0u8;
19const DATA_IDENT: u8 = 1u8;
20
21/// Module key formatting macro
22#[macro_export]
23macro_rules! module_key {
24    ($bc_hash:expr) => {
25        [&$bc_hash.to_bytes()[..], &[MODULE_IDENT]].concat()
26    };
27}
28
29/// Delta key formatting macro
30#[macro_export]
31macro_rules! metadata_key {
32    ($bc_hash:expr) => {
33        [&$bc_hash.to_bytes()[..], &[DATA_IDENT]].concat()
34    };
35}
36
37pub(crate) struct HDCache {
38    /// RocksDB database
39    db: Option<DB>,
40    /// How many entries are in the db. Count is initialized at creation time by iterating
41    /// over all the entries in the db then it is maintained in memory
42    entry_count: usize,
43    /// Maximum number of entries we want to keep in the db.
44    /// When this maximum is reached `snip_amount` entries are removed
45    max_entry_count: usize,
46    /// How many entries are removed when `entry_count` reaches `max_entry_count`
47    snip_amount: usize,
48    /// Module metadata serializer
49    meta_ser: ModuleMetadataSerializer,
50    /// Module metadata deserializer
51    meta_deser: ModuleMetadataDeserializer,
52    /// Test-only counter of `get` calls, used to assert that no redundant
53    /// RocksDB read/deserialize happens when a module is already in the LRU cache.
54    #[cfg(test)]
55    pub(crate) read_count: std::sync::atomic::AtomicUsize,
56}
57
58impl HDCache {
59    /// Create a new HDCache
60    ///
61    /// # Arguments
62    /// * path: where to store the db
63    /// * max_entry_count: maximum number of entries we want to keep in the db
64    /// * amount_to_remove: how many entries are removed when `entry_count` reaches `max_entry_count`
65    pub fn new(path: PathBuf, max_entry_count: usize, snip_amount: usize) -> Self {
66        // Reset the DB if it already exists
67        if path.exists() {
68            if let Err(e) = DB::destroy(&Options::default(), path.clone()) {
69                warn!("Failed to destroy the db: {:?}", e);
70            }
71        }
72        let db = DB::open_default(path).expect(OPEN_ERROR);
73        let entry_count = 0;
74
75        Self {
76            db: Some(db),
77            entry_count,
78            max_entry_count,
79            snip_amount,
80            meta_ser: ModuleMetadataSerializer::new(),
81            meta_deser: ModuleMetadataDeserializer::new(),
82            #[cfg(test)]
83            read_count: std::sync::atomic::AtomicUsize::new(0),
84        }
85    }
86
87    pub fn reset(&mut self) {
88        let path = self.db.as_ref().unwrap().path().to_path_buf();
89
90        // Close the existing database by dropping it
91        let _ = self.db.take();
92
93        // Destroy the database files
94        if path.exists() {
95            if let Err(e) = DB::destroy(&Options::default(), path.clone()) {
96                warn!("Failed to destroy the db: {:?}", e);
97            }
98        }
99        // Reopen the database
100        let db = DB::open_default(&path).expect(OPEN_ERROR);
101        self.db = Some(db);
102        self.entry_count = 0;
103    }
104
105    /// Insert a new module in the cache
106    pub fn insert(&mut self, hash: Hash, module_info: ModuleInfo) {
107        if self.entry_count >= self.max_entry_count {
108            self.snip();
109        }
110
111        let mut ser_metadata = Vec::new();
112        let ser_module = match module_info {
113            ModuleInfo::Invalid(err_msg) => {
114                self.meta_ser
115                    .serialize(&ModuleMetadata::Invalid(err_msg), &mut ser_metadata)
116                    .expect(DATA_SER_ERROR);
117                Vec::new()
118            }
119            ModuleInfo::Module(module) => {
120                self.meta_ser
121                    .serialize(&ModuleMetadata::NotExecuted, &mut ser_metadata)
122                    .expect(DATA_SER_ERROR);
123                module.serialize().expect(MOD_SER_ERROR)
124            }
125            ModuleInfo::ModuleAndDelta((module, delta)) => {
126                self.meta_ser
127                    .serialize(&ModuleMetadata::Delta(delta), &mut ser_metadata)
128                    .expect(DATA_SER_ERROR);
129                module.serialize().expect(MOD_SER_ERROR)
130            }
131        };
132
133        let mut batch = WriteBatch::default();
134        batch.put(module_key!(hash), ser_module);
135        batch.put(metadata_key!(hash), ser_metadata);
136        self.db
137            .as_ref()
138            .expect(CRUD_ERROR)
139            .write(&batch)
140            .expect(CRUD_ERROR);
141
142        self.entry_count = self.entry_count.saturating_add(1);
143
144        debug!("(HD insert) entry_count is: {}", self.entry_count);
145    }
146
147    /// Sets the initialization cost of a given module separately
148    ///
149    /// # Arguments
150    /// * `hash`: hash associated to the module for which we want to set the cost
151    /// * `init_cost`: the new cost associated to the module
152    pub fn set_init_cost(&self, hash: Hash, init_cost: u64) {
153        let mut ser_metadata = Vec::new();
154        self.meta_ser
155            .serialize(&ModuleMetadata::Delta(init_cost), &mut ser_metadata)
156            .expect(DATA_SER_ERROR);
157        self.db
158            .as_ref()
159            .expect(CRUD_ERROR)
160            .put(metadata_key!(hash), ser_metadata)
161            .expect(CRUD_ERROR);
162    }
163
164    /// Sets a given module as invalid
165    pub fn set_invalid(&self, hash: Hash, err_msg: String) {
166        let mut ser_metadata = Vec::new();
167        self.meta_ser
168            .serialize(&ModuleMetadata::Invalid(err_msg), &mut ser_metadata)
169            .expect(DATA_SER_ERROR);
170        self.db
171            .as_ref()
172            .expect(CRUD_ERROR)
173            .put(metadata_key!(hash), ser_metadata)
174            .expect(CRUD_ERROR);
175    }
176
177    /// Retrieve a module
178    pub fn get(
179        &self,
180        hash: Hash,
181        gas_costs: GasCosts,
182        condom_limits: CondomLimits,
183    ) -> Option<ModuleInfo> {
184        #[cfg(test)]
185        self.read_count
186            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
187
188        let mut iterator = self
189            .db
190            .as_ref()
191            .expect(CRUD_ERROR)
192            .iterator(IteratorMode::From(&module_key!(hash), Direction::Forward));
193
194        if let (Some(Ok((key_1, ser_module))), Some(Ok((key_2, ser_metadata)))) =
195            (iterator.next(), iterator.next())
196        {
197            if *key_1 == module_key!(hash) && *key_2 == metadata_key!(hash) {
198                let (_, metadata) = self
199                    .meta_deser
200                    .deserialize::<DeserializeError>(&ser_metadata)
201                    .expect(DATA_DESER_ERROR);
202                if let ModuleMetadata::Invalid(err_msg) = metadata {
203                    return Some(ModuleInfo::Invalid(err_msg));
204                }
205                let module = RuntimeModule::deserialize(
206                    &ser_module,
207                    gas_costs.max_instance_cost,
208                    gas_costs,
209                    condom_limits,
210                )
211                .expect(MOD_DESER_ERROR);
212                let result = match metadata {
213                    ModuleMetadata::Invalid(err_msg) => ModuleInfo::Invalid(err_msg),
214                    ModuleMetadata::NotExecuted => ModuleInfo::Module(module),
215                    ModuleMetadata::Delta(delta) => ModuleInfo::ModuleAndDelta((module, delta)),
216                };
217                Some(result)
218            } else {
219                None
220            }
221        } else {
222            None
223        }
224    }
225
226    /// Try to remove as much as `self.amount_to_snip` entries from the db
227    fn snip(&mut self) {
228        let mut iter = self.db.as_ref().expect(CRUD_ERROR).raw_iterator();
229        let mut batch = WriteBatch::default();
230        let mut snipped_count: usize = 0;
231
232        while snipped_count < self.snip_amount {
233            // generate a random key
234            let mut rbytes = [0u8; 16];
235            rand::thread_rng().fill_bytes(&mut rbytes);
236            let key = *Hash::compute_from(&rbytes).to_bytes();
237
238            // take the upper existing key
239            iter.seek_for_prev(key);
240
241            // check iterator validity
242            if !iter.valid() {
243                continue;
244            }
245
246            // unwrap justified by above conditional statement.
247            // seeking the previous key of a randombly generated one
248            // will always end up on a metadata key.
249            let metadata_key = iter.key().unwrap();
250            batch.delete(metadata_key);
251            iter.prev();
252            let module_key = iter.key().unwrap();
253            batch.delete(module_key);
254
255            // increase snipped_count
256            snipped_count += 1;
257        }
258
259        // safety check
260        if batch.len() / 2 != snipped_count {
261            panic!("snipped_count incoherence");
262        }
263
264        // delete the key and reduce entry_count
265        self.db
266            .as_ref()
267            .expect(CRUD_ERROR)
268            .write(&batch)
269            .expect(CRUD_ERROR);
270        self.entry_count -= snipped_count;
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use massa_hash::Hash;
278    use massa_sc_runtime::{Compiler, GasCosts, RuntimeModule};
279    use rand::thread_rng;
280    use serial_test::serial;
281    use tempfile::TempDir;
282
283    fn make_default_module_info() -> ModuleInfo {
284        let bytecode: Vec<u8> = vec![
285            0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x06, 0x01, 0x60, 0x01, 0x7f,
286            0x01, 0x7f, 0x03, 0x02, 0x01, 0x00, 0x07, 0x0b, 0x01, 0x07, 0x61, 0x64, 0x64, 0x5f,
287            0x6f, 0x6e, 0x65, 0x00, 0x00, 0x0a, 0x09, 0x01, 0x07, 0x00, 0x20, 0x00, 0x41, 0x01,
288            0x6a, 0x0b, 0x00, 0x1a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x01, 0x0a, 0x01, 0x00, 0x07,
289            0x61, 0x64, 0x64, 0x5f, 0x6f, 0x6e, 0x65, 0x02, 0x07, 0x01, 0x00, 0x01, 0x00, 0x02,
290            0x70, 0x30,
291        ];
292        ModuleInfo::Module(
293            RuntimeModule::new(
294                &bytecode,
295                GasCosts::default(),
296                Compiler::CL,
297                CondomLimits::default(),
298            )
299            .unwrap(),
300        )
301    }
302
303    fn setup() -> HDCache {
304        let tmp_path = TempDir::new().unwrap().path().to_path_buf();
305        HDCache::new(tmp_path, 1000, 10)
306    }
307
308    #[test]
309    #[serial]
310    fn test_basic_crud() {
311        let mut cache = setup();
312        let hash = Hash::compute_from(b"test_hash");
313        let module = make_default_module_info();
314
315        let init_cost = 100;
316        let gas_costs = GasCosts::default();
317        let condom_limits = CondomLimits::default();
318
319        cache.insert(hash, module);
320        let cached_module_v1 = cache
321            .get(hash, gas_costs.clone(), condom_limits.clone())
322            .unwrap();
323        assert!(matches!(cached_module_v1, ModuleInfo::Module(_)));
324
325        cache.set_init_cost(hash, init_cost);
326        let cached_module_v2 = cache
327            .get(hash, gas_costs.clone(), condom_limits.clone())
328            .unwrap();
329        assert!(matches!(cached_module_v2, ModuleInfo::ModuleAndDelta(_)));
330
331        let err_msg = "test_error".to_string();
332        cache.set_invalid(hash, err_msg.clone());
333        let cached_module_v3 = cache.get(hash, gas_costs, condom_limits.clone()).unwrap();
334        let ModuleInfo::Invalid(res_err) = cached_module_v3 else {
335            panic!("expected ModuleInfo::Invalid");
336        };
337        assert_eq!(res_err, err_msg);
338    }
339
340    #[test]
341    #[serial]
342    fn test_insert_more_than_max_entry() {
343        let mut cache = setup();
344        let module = make_default_module_info();
345
346        // fill the db: add cache.max_entry_count entries
347        for count in 0..cache.max_entry_count {
348            let key = Hash::compute_from(count.to_string().as_bytes());
349            cache.insert(key, module.clone());
350        }
351        assert_eq!(cache.entry_count, cache.max_entry_count);
352
353        // insert one more entry
354        let key = Hash::compute_from(cache.max_entry_count.to_string().as_bytes());
355        cache.insert(key, module);
356        assert_eq!(
357            cache.entry_count,
358            cache.max_entry_count - cache.snip_amount + 1
359        );
360        dbg!(cache.entry_count);
361    }
362
363    #[test]
364    #[serial]
365    fn test_missing_module() {
366        let mut cache = setup();
367        let module = make_default_module_info();
368
369        let gas_costs = GasCosts::default();
370        let condom_limits = CondomLimits::default();
371
372        for count in 0..cache.max_entry_count {
373            let key = Hash::compute_from(count.to_string().as_bytes());
374            cache.insert(key, module.clone());
375        }
376
377        for _ in 0..cache.max_entry_count {
378            let mut rbytes = [0u8; 16];
379            thread_rng().fill_bytes(&mut rbytes);
380            let get_key = Hash::compute_from(&rbytes);
381            let cached_module = cache.get(get_key, gas_costs.clone(), condom_limits.clone());
382            assert!(cached_module.is_none());
383        }
384    }
385}
massa_module_cache/hd_cache.rs

massa_module_cache/
hd_cache.rs
