massa_module_cache/
controller.rs

1use massa_hash::Hash;
2use massa_models::prehash::BuildHashMapper;
3use massa_sc_runtime::{Compiler, CondomLimits, RuntimeModule};
4use schnellru::{ByLength, LruMap};
5use tracing::debug;
6
7use crate::{
8    config::ModuleCacheConfig, error::CacheError, hd_cache::HDCache, lru_cache::LRUCache,
9    types::ModuleInfo,
10};
11
12/// `LruMap` specialization for `PreHashed` keys
13pub type PreHashLruMap<K, V> = LruMap<K, V, ByLength, BuildHashMapper<K>>;
14
15/// Cache controller of compiled runtime modules
16pub struct ModuleCache {
17    /// Cache config.
18    /// See `CacheConfig` documentation for more information.
19    cfg: ModuleCacheConfig,
20    /// RAM stored LRU cache.
21    /// See `LRUCache` documentation for more information.
22    lru_cache: LRUCache,
23    /// Disk stored cache.
24    /// See the `HDCache` documentation for more information.
25    hd_cache: HDCache,
26}
27
28impl ModuleCache {
29    /// Creates a new `ModuleCache`
30    pub fn new(cfg: ModuleCacheConfig) -> Self {
31        Self {
32            lru_cache: LRUCache::new(cfg.lru_cache_size),
33            hd_cache: HDCache::new(
34                cfg.hd_cache_path.clone(),
35                cfg.hd_cache_size,
36                cfg.snip_amount,
37            ),
38            cfg,
39        }
40    }
41
42    pub fn reset(&mut self) {
43        self.lru_cache.reset();
44        self.hd_cache.reset();
45    }
46
47    /// Internal function to compile and build `ModuleInfo`
48    fn compile_cached(
49        &mut self,
50        bytecode: &[u8],
51        hash: Hash,
52        condom_limits: CondomLimits,
53    ) -> ModuleInfo {
54        match RuntimeModule::new(
55            bytecode,
56            self.cfg.gas_costs.clone(),
57            Compiler::CL,
58            condom_limits,
59        ) {
60            Ok(module) => {
61                debug!("compilation of module {} succeeded", hash);
62                ModuleInfo::Module(module)
63            }
64            Err(e) => {
65                let err_msg = format!("compilation of module {} failed: {}", hash, e);
66                debug!(err_msg);
67                ModuleInfo::Invalid(err_msg)
68            }
69        }
70    }
71
72    /// Save a new or an already existing module in the cache
73    pub fn save_module(&mut self, bytecode: &[u8], condom_limits: CondomLimits) {
74        let hash = Hash::compute_from(bytecode);
75        if let Some(lru_module_info) = self.lru_cache.get(hash) {
76            debug!("save_module: {} present in lru", hash);
77            self.hd_cache.insert(hash, lru_module_info);
78        } else if let Some(hd_module_info) =
79            self.hd_cache
80                .get(hash, self.cfg.gas_costs.clone(), condom_limits.clone())
81        {
82            debug!("save_module: {} missing in lru but present in hd", hash);
83            self.lru_cache.insert(hash, hd_module_info);
84        } else {
85            debug!("save_module: {} missing", hash);
86            let module_info = self.compile_cached(bytecode, hash, condom_limits);
87            self.hd_cache.insert(hash, module_info.clone());
88            self.lru_cache.insert(hash, module_info);
89        }
90    }
91
92    /// Set the initialization cost of a cached module
93    pub fn set_init_cost(&mut self, bytecode: &[u8], init_cost: u64) {
94        let hash = Hash::compute_from(bytecode);
95        self.lru_cache.set_init_cost(hash, init_cost);
96        self.hd_cache.set_init_cost(hash, init_cost);
97    }
98
99    /// Set a cached module as invalid
100    pub fn set_invalid(&mut self, bytecode: &[u8], err_msg: String) {
101        let hash = Hash::compute_from(bytecode);
102        self.lru_cache.set_invalid(hash, err_msg.clone());
103        self.hd_cache.set_invalid(hash, err_msg);
104    }
105
106    /// Load a cached module for execution
107    ///
108    /// Returns the module information, it can be:
109    /// * `ModuleInfo::Invalid` if the module is invalid
110    /// * `ModuleInfo::Module` if the module is valid and has no delta
111    /// * `ModuleInfo::ModuleAndDelta` if the module is valid and has a delta
112    fn load_module_info(&mut self, bytecode: &[u8], condom_limits: CondomLimits) -> ModuleInfo {
113        if bytecode.is_empty() {
114            let error_msg = "load_module: bytecode is absent".to_string();
115            debug!(error_msg);
116            return ModuleInfo::Invalid(error_msg);
117        }
118        if bytecode.len() > self.cfg.max_module_length as usize {
119            let error_msg = format!(
120                "load_module: bytecode length {} exceeds max module length {}",
121                bytecode.len(),
122                self.cfg.max_module_length
123            );
124            debug!(error_msg);
125            return ModuleInfo::Invalid(error_msg);
126        }
127        let hash = Hash::compute_from(bytecode);
128        if let Some(lru_module_info) = self.lru_cache.get(hash) {
129            debug!("load_module: {} present in lru", hash);
130            lru_module_info
131        } else if let Some(hd_module_info) =
132            self.hd_cache
133                .get(hash, self.cfg.gas_costs.clone(), condom_limits.clone())
134        {
135            debug!("load_module: {} missing in lru but present in hd", hash);
136            self.lru_cache.insert(hash, hd_module_info.clone());
137            hd_module_info
138        } else {
139            debug!("load_module: {} missing", hash);
140            let module_info = self.compile_cached(bytecode, hash, condom_limits);
141            self.hd_cache.insert(hash, module_info.clone());
142            self.lru_cache.insert(hash, module_info.clone());
143            module_info
144        }
145    }
146
147    /// Load a cached module for execution and check its validity for execution.
148    /// Also checks that the provided execution gas is enough to pay for the instance creation cost.
149    ///
150    /// Returns the module after loading.
151    pub fn load_module(
152        &mut self,
153        bytecode: &[u8],
154        execution_gas: u64,
155        condom_limits: CondomLimits,
156    ) -> Result<RuntimeModule, CacheError> {
157        // Do not actually debit the instance creation cost from the provided gas
158        // This is only supposed to be a check
159        execution_gas
160            .checked_sub(self.cfg.gas_costs.max_instance_cost)
161            .ok_or(CacheError::LoadError(format!(
162                "Provided gas {} is lower than the base instance creation gas cost {}",
163                execution_gas, self.cfg.gas_costs.max_instance_cost
164            )))?;
165        // TODO: interesting but unimportant optim
166        // remove max_instance_cost hard check if module is cached and has a delta
167        let module_info = self.load_module_info(bytecode, condom_limits);
168        let module = match module_info {
169            ModuleInfo::Invalid(err) => {
170                let err_msg = format!("invalid module: {}", err);
171                return Err(CacheError::LoadError(err_msg));
172            }
173            ModuleInfo::Module(module) => module,
174            ModuleInfo::ModuleAndDelta((module, delta)) => {
175                if delta > execution_gas {
176                    return Err(CacheError::LoadError(format!(
177                        "Provided gas {} is below the gas cost of instance creation ({})",
178                        execution_gas, delta
179                    )));
180                } else {
181                    module
182                }
183            }
184        };
185        Ok(module)
186    }
187
188    /// Load a temporary module from arbitrary bytecode.
189    /// Also checks that the provided execution gas is enough to pay for the instance creation cost.
190    ///
191    /// Returns the module after compilation.
192    pub fn load_tmp_module(
193        &self,
194        bytecode: &[u8],
195        limit: u64,
196        condom_limits: CondomLimits,
197    ) -> Result<RuntimeModule, CacheError> {
198        debug!("load_tmp_module");
199        if bytecode.is_empty() {
200            let error_msg = "load_tmp_module: bytecode is absent".to_string();
201            debug!(error_msg);
202            return Err(CacheError::LoadError(error_msg));
203        }
204        if bytecode.len() > self.cfg.max_module_length as usize {
205            let error_msg = format!(
206                "load_tmp_module: bytecode length {} exceeds max module length {}",
207                bytecode.len(),
208                self.cfg.max_module_length
209            );
210            debug!(error_msg);
211            return Err(CacheError::LoadError(error_msg));
212        }
213        // Do not actually debit the instance creation cost from the provided gas
214        // This is only supposed to be a check
215        limit
216            .checked_sub(self.cfg.gas_costs.max_instance_cost)
217            .ok_or(CacheError::LoadError(format!(
218                "Provided gas {} is lower than the base instance creation gas cost {}",
219                limit, self.cfg.gas_costs.max_instance_cost
220            )))?;
221        let module = RuntimeModule::new(
222            bytecode,
223            self.cfg.gas_costs.clone(),
224            Compiler::SP,
225            condom_limits,
226        )?;
227        Ok(module)
228    }
229
230    /// Returns the memory usage of the LRU cache
231    pub fn get_module_lru_cache_memory_usage(&self) -> usize {
232        self.lru_cache.cache.memory_usage()
233    }
234
235    /// Returns the number of modules currently held in the in-RAM LRU cache
236    pub fn lru_cache_len(&self) -> usize {
237        self.lru_cache.cache.len()
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use massa_sc_runtime::{CondomLimits, GasCosts};
245    use serial_test::serial;
246    use std::sync::atomic::Ordering;
247    use tempfile::TempDir;
248
249    /// Minimal valid wasm module (an `add_one` function).
250    const TEST_BYTECODE: &[u8] = &[
251        0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x06, 0x01, 0x60, 0x01, 0x7f, 0x01,
252        0x7f, 0x03, 0x02, 0x01, 0x00, 0x07, 0x0b, 0x01, 0x07, 0x61, 0x64, 0x64, 0x5f, 0x6f, 0x6e,
253        0x65, 0x00, 0x00, 0x0a, 0x09, 0x01, 0x07, 0x00, 0x20, 0x00, 0x41, 0x01, 0x6a, 0x0b, 0x00,
254        0x1a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x01, 0x0a, 0x01, 0x00, 0x07, 0x61, 0x64, 0x64, 0x5f,
255        0x6f, 0x6e, 0x65, 0x02, 0x07, 0x01, 0x00, 0x01, 0x00, 0x02, 0x70, 0x30,
256    ];
257
258    fn setup(max_module_length: u64) -> (ModuleCache, TempDir) {
259        let cache_dir = TempDir::new().unwrap();
260        let cache = ModuleCache::new(ModuleCacheConfig {
261            hd_cache_path: cache_dir.path().to_path_buf(),
262            gas_costs: GasCosts::default(),
263            lru_cache_size: 10,
264            hd_cache_size: 10,
265            snip_amount: 1,
266            max_module_length,
267            condom_limits: CondomLimits::default(),
268        });
269        (cache, cache_dir)
270    }
271
272    /// When a module is already resident in the LRU cache, `save_module` must not
273    /// perform a redundant RocksDB read/deserialize via the HD cache.
274    #[test]
275    #[serial]
276    fn test_save_module_skips_hd_read_when_in_lru() {
277        let (mut cache, _cache_dir) = setup(1_000_000);
278        let condom_limits = CondomLimits::default();
279
280        // First save: nothing is cached yet, so the HD cache is probed once (a miss)
281        // before compiling and inserting into both caches.
282        cache.save_module(TEST_BYTECODE, condom_limits.clone());
283        assert_eq!(
284            cache.hd_cache.read_count.load(Ordering::Relaxed),
285            1,
286            "first save_module should probe the HD cache exactly once"
287        );
288
289        // Second save: the module is already present in the LRU cache, so we must
290        // short-circuit and never touch the HD cache for a read.
291        cache.save_module(TEST_BYTECODE, condom_limits);
292        assert_eq!(
293            cache.hd_cache.read_count.load(Ordering::Relaxed),
294            1,
295            "save_module must not read the HD cache when the module is in the LRU cache"
296        );
297    }
298
299    #[test]
300    fn test_load_tmp_module_rejects_empty_bytecode() {
301        let (cache, _cache_dir) = setup(4);
302
303        let result = cache.load_tmp_module(&[], u64::MAX, CondomLimits::default());
304
305        assert!(matches!(
306            result,
307            Err(CacheError::LoadError(error)) if error == "load_tmp_module: bytecode is absent"
308        ));
309    }
310
311    #[test]
312    fn test_load_tmp_module_rejects_oversized_bytecode() {
313        let (cache, _cache_dir) = setup(4);
314
315        // Invalid Wasm confirms the size check runs before RuntimeModule compilation.
316        let result = cache.load_tmp_module(&[0; 5], u64::MAX, CondomLimits::default());
317
318        assert!(matches!(
319            result,
320            Err(CacheError::LoadError(error))
321                if error == "load_tmp_module: bytecode length 5 exceeds max module length 4"
322        ));
323    }
324}
massa_module_cache/controller.rs

massa_module_cache/
controller.rs
