namada_vm/wasm/compilation_cache/
common.rs

1//! WASM compilation cache.
2//!
3//! The cache is backed by in-memory LRU cache with configurable size
4//! limit and a file system cache of serialized modules.
5
6use std::collections::hash_map::RandomState;
7use std::fs;
8use std::marker::PhantomData;
9use std::num::NonZeroUsize;
10use std::path::{Path, PathBuf};
11use std::sync::{Arc, RwLock};
12use std::thread::sleep;
13use std::time::Duration;
14
15use clru::{CLruCache, CLruCacheConfig, WeightScale};
16use namada_core::collections::HashMap;
17use namada_core::control_flow::time::{ExponentialBackoff, SleepStrategy};
18use namada_core::hash::Hash;
19use namada_gas::GasMeterKind;
20use wasmer::{Module, Store};
21use wasmer_cache::{FileSystemCache, Hash as CacheHash};
22
23use crate::wasm::run::untrusted_wasm_store;
24use crate::wasm::{self, memory};
25use crate::{WasmCacheAccess, WasmCacheRoAccess};
26
27/// Cache handle. Thread-safe.
28#[derive(Debug, Clone)]
29pub struct Cache<N, A> {
30    /// Cached files directory
31    dir: PathBuf,
32    /// Compilation progress
33    progress: Arc<RwLock<HashMap<CacheKey, Compilation>>>,
34    /// In-memory LRU cache of compiled modules
35    in_memory: Arc<RwLock<MemoryCache>>,
36    /// The cache's name
37    name: PhantomData<N>,
38    /// Cache access level
39    access: PhantomData<A>,
40    /// Wasmer store - The store that's used to compile the modules.
41    // N.B.: This has to be kept alive to avoid segfaults when running them
42    // from cache (otherwise `wasmer_compiler::FRAME_INFO` gets cleared out
43    // when the `Store` is dropped and the next run of this `Module`
44    // doesn't re-instantiate it and crashes when it tries to access it).
45    // Related issues:
46    // - https://github.com/wasmerio/wasmer/issues/4377
47    // - https://github.com/wasmerio/wasmer/issues/4949
48    store: Arc<Store>,
49}
50
51/// This trait is used to give names to different caches
52pub trait CacheName: Clone + std::fmt::Debug {
53    /// Get the name of the cache
54    fn name() -> &'static str;
55}
56
57/// WASM cache key consists of the code hash and gas meter kind.
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
59pub struct CacheKey {
60    code_hash: Hash,
61    gas_meter_kind: GasMeterKind,
62}
63
64/// In-memory LRU cache of compiled modules
65type MemoryCache = CLruCache<CacheKey, Module, RandomState, ModuleCacheScale>;
66
67/// Compilation progress
68#[derive(Debug)]
69enum Compilation {
70    Compiling,
71    Done,
72}
73
74/// Configures the cache scale of modules that limits the maximum capacity
75/// of the cache (CLruCache::len + CLruCache::weight <= CLruCache::capacity).
76#[derive(Debug)]
77struct ModuleCacheScale;
78
79impl WeightScale<CacheKey, Module> for ModuleCacheScale {
80    fn weight(&self, _key: &CacheKey, _value: &Module) -> usize {
81        1
82    }
83}
84
85impl<N: CacheName, A: WasmCacheAccess> Cache<N, A> {
86    /// Create a wasm in-memory cache with a given size limit and a file
87    /// system cache.
88    ///
89    /// # Panics
90    /// The `max_bytes` must be non-zero.
91    pub fn new(dir: impl Into<PathBuf>, max_bytes: usize) -> Self {
92        let cache = CLruCache::with_config(
93            CLruCacheConfig::new(NonZeroUsize::new(max_bytes).unwrap())
94                .with_scale(ModuleCacheScale),
95        );
96        let in_memory = Arc::new(RwLock::new(cache));
97
98        let target_hash = {
99            use std::hash::{Hash, Hasher};
100            let mut hasher = std::hash::DefaultHasher::new();
101            wasmer::Target::default().hash(&mut hasher);
102            hasher.finish()
103        };
104        let version_prefix =
105            option_env!("GIT_DESCRIBED").unwrap_or(env!("CARGO_PKG_VERSION"));
106        let version = format!(
107            "{version_prefix}{}{:x}",
108            concat!("_", env!("RUSTUP_TOOLCHAIN"), "_"),
109            target_hash,
110        );
111        let dir = dir.into().join(version);
112
113        fs::create_dir_all(&dir)
114            .expect("Couldn't create the wasm cache directory");
115
116        Self {
117            dir,
118            progress: Default::default(),
119            in_memory,
120            name: Default::default(),
121            access: Default::default(),
122            store: Arc::new(store()),
123        }
124    }
125
126    /// Get a WASM module from LRU cache, from a file or compile it and cache
127    /// it. If the cache access is set to [`crate::WasmCacheRwAccess`], it
128    /// updates the position in the LRU cache. Otherwise, the compiled
129    /// module will not be cached, if it's not already.
130    pub fn fetch(
131        &mut self,
132        code_hash: &Hash,
133        gas_meter_kind: GasMeterKind,
134    ) -> Result<Option<(Module, Store)>, wasm::run::Error> {
135        if A::is_read_write() {
136            let module = self.get(code_hash, gas_meter_kind)?;
137            Ok(module.map(|module| (module, store())))
138        } else {
139            let store = store();
140            let module = self.peek(code_hash, gas_meter_kind, &store)?;
141            Ok(module.map(|module| (module, store)))
142        }
143    }
144
145    /// Get the current number of items in the cache
146    pub fn get_size(&self) -> usize {
147        self.in_memory.read().unwrap().len()
148    }
149
150    /// Get the current weight of the cache
151    pub fn get_cache_size(&self) -> usize {
152        self.in_memory.read().unwrap().weight()
153    }
154
155    /// Get a WASM module from LRU cache, from a file or compile it and cache
156    /// it. Updates the position in the LRU cache.
157    fn get(
158        &mut self,
159        hash: &Hash,
160        gas_meter_kind: GasMeterKind,
161    ) -> Result<Option<Module>, wasm::run::Error> {
162        let key = CacheKey {
163            code_hash: *hash,
164            gas_meter_kind,
165        };
166        let mut in_memory = self.in_memory.write().unwrap();
167        if let Some(module) = in_memory.get(&key) {
168            tracing::trace!(
169                "{} found {} in cache.",
170                N::name(),
171                hash.to_string()
172            );
173            return Ok(Some(module.clone()));
174        }
175        drop(in_memory);
176
177        let mut iter = 0;
178        let exponential_backoff = ExponentialBackoff {
179            base: 2,
180            as_duration: |backoff: u64| {
181                Duration::from_millis(backoff.saturating_mul(10))
182            },
183        };
184        loop {
185            let progress = self.progress.read().unwrap();
186            match progress.get(&key) {
187                Some(Compilation::Done) => {
188                    drop(progress);
189                    let mut in_memory = self.in_memory.write().unwrap();
190                    if let Some(module) = in_memory.get(&key) {
191                        tracing::info!(
192                            "{} found {} in memory cache.",
193                            N::name(),
194                            hash.to_string()
195                        );
196                        return Ok(Some(module.clone()));
197                    }
198
199                    if let Ok(module) = file_load_module(
200                        &self.dir,
201                        hash,
202                        gas_meter_kind,
203                        &self.store,
204                    ) {
205                        tracing::info!(
206                            "{} found {} in file cache.",
207                            N::name(),
208                            hash.to_string()
209                        );
210                        // Put into cache, ignore result if it's full
211                        let _ = in_memory.put_with_weight(key, module.clone());
212
213                        return Ok(Some(module));
214                    } else {
215                        return Ok(None);
216                    }
217                }
218                Some(Compilation::Compiling) => {
219                    drop(progress);
220                    tracing::info!(
221                        "Waiting for {} {} ...",
222                        N::name(),
223                        hash.to_string()
224                    );
225                    sleep(exponential_backoff.backoff(&iter));
226                    // Cannot overflow
227                    #[allow(clippy::arithmetic_side_effects)]
228                    {
229                        iter += 1;
230                    }
231                    continue;
232                }
233                None => {
234                    drop(progress);
235                    let module = if module_file_exists(
236                        &self.dir,
237                        hash,
238                        gas_meter_kind,
239                    ) {
240                        tracing::info!(
241                            "Trying to load {} {} from file.",
242                            N::name(),
243                            hash.to_string()
244                        );
245                        if let Ok(module) = file_load_module(
246                            &self.dir,
247                            hash,
248                            gas_meter_kind,
249                            &self.store,
250                        ) {
251                            module
252                        } else {
253                            return Ok(None);
254                        }
255                    } else {
256                        return Ok(None);
257                    };
258
259                    // Update progress
260                    let mut progress = self.progress.write().unwrap();
261                    progress.insert(key, Compilation::Done);
262
263                    // Put into cache, ignore the result (fails if the module
264                    // cannot fit into the cache)
265                    let mut in_memory = self.in_memory.write().unwrap();
266                    let _ = in_memory.put_with_weight(key, module.clone());
267
268                    return Ok(Some(module));
269                }
270            }
271        }
272    }
273
274    /// Peak-only is used for dry-ran txs (and VPs that the tx triggers).
275    /// It doesn't update the in-memory cache.
276    fn peek(
277        &self,
278        hash: &Hash,
279        gas_meter_kind: GasMeterKind,
280        store: &Store,
281    ) -> Result<Option<Module>, wasm::run::Error> {
282        let key = CacheKey {
283            code_hash: *hash,
284            gas_meter_kind,
285        };
286        let in_memory = self.in_memory.read().unwrap();
287        if let Some(module) = in_memory.peek(&key) {
288            tracing::info!(
289                "{} found {} in cache.",
290                N::name(),
291                hash.to_string()
292            );
293            return Ok(Some(module.clone()));
294        }
295        drop(in_memory);
296
297        let mut iter = 0;
298        let exponential_backoff = ExponentialBackoff {
299            base: 2,
300            as_duration: |backoff: u64| {
301                Duration::from_millis(backoff.saturating_mul(10))
302            },
303        };
304        loop {
305            let progress = self.progress.read().unwrap();
306            match progress.get(&key) {
307                Some(Compilation::Done) => {
308                    drop(progress);
309                    let in_memory = self.in_memory.read().unwrap();
310                    if let Some(module) = in_memory.peek(&key) {
311                        tracing::info!(
312                            "{} found {} in memory cache.",
313                            N::name(),
314                            hash.to_string()
315                        );
316                        return Ok(Some(module.clone()));
317                    }
318
319                    if let Ok(module) =
320                        file_load_module(&self.dir, hash, gas_meter_kind, store)
321                    {
322                        tracing::info!(
323                            "{} found {} in file cache.",
324                            N::name(),
325                            hash.to_string()
326                        );
327                        return Ok(Some(module));
328                    } else {
329                        return Ok(None);
330                    }
331                }
332                Some(Compilation::Compiling) => {
333                    drop(progress);
334                    tracing::info!(
335                        "Waiting for {} {} ...",
336                        N::name(),
337                        hash.to_string()
338                    );
339                    sleep(exponential_backoff.backoff(&iter));
340                    // Cannot overflow
341                    #[allow(clippy::arithmetic_side_effects)]
342                    {
343                        iter += 1;
344                    }
345                    continue;
346                }
347                None => {
348                    drop(progress);
349
350                    return if module_file_exists(
351                        &self.dir,
352                        hash,
353                        gas_meter_kind,
354                    ) {
355                        tracing::info!(
356                            "Trying to load {} {} from file.",
357                            N::name(),
358                            hash.to_string()
359                        );
360                        if let Ok(module) = file_load_module(
361                            &self.dir,
362                            hash,
363                            gas_meter_kind,
364                            store,
365                        ) {
366                            return Ok(Some(module));
367                        } else {
368                            return Ok(None);
369                        }
370                    } else {
371                        Ok(None)
372                    };
373                }
374            }
375        }
376    }
377
378    /// Compile a WASM module and persist the compiled modules to files.
379    pub fn compile_or_fetch(
380        &mut self,
381        code: impl AsRef<[u8]>,
382        gas_meter_kind: GasMeterKind,
383    ) -> Result<Option<(Module, Store)>, wasm::run::Error> {
384        let hash = hash_of_code(&code);
385        let key = CacheKey {
386            code_hash: hash,
387            gas_meter_kind,
388        };
389
390        if !A::is_read_write() {
391            // It doesn't update the cache and files
392            let progress = self.progress.read().unwrap();
393            match progress.get(&key) {
394                Some(_) => {
395                    let store = store();
396                    let module = self.peek(&hash, gas_meter_kind, &store)?;
397                    return Ok(module.map(|module| (module, store)));
398                }
399                None => {
400                    let code =
401                        wasm::run::prepare_wasm_code(code, gas_meter_kind)?;
402                    let store = store();
403                    let module = compile(code, &store)?;
404                    return Ok(Some((module, store)));
405                }
406            }
407        }
408
409        let mut progress = self.progress.write().unwrap();
410        if progress.get(&key).is_some() {
411            drop(progress);
412            return self.fetch(&hash, gas_meter_kind);
413        }
414        progress.insert(key, Compilation::Compiling);
415        drop(progress);
416
417        tracing::info!("Compiling {} {}.", N::name(), hash.to_string());
418
419        match wasm::run::prepare_wasm_code(code, gas_meter_kind) {
420            Ok(code) => match compile(code, &self.store) {
421                Ok(module) => {
422                    // Write the file
423                    file_write_module(
424                        &self.dir,
425                        &module,
426                        &hash,
427                        gas_meter_kind,
428                    );
429
430                    // Update progress
431                    let mut progress = self.progress.write().unwrap();
432                    progress.insert(key, Compilation::Done);
433
434                    // Put into cache, ignore result if it's full
435                    let mut in_memory = self.in_memory.write().unwrap();
436                    let _ = in_memory.put_with_weight(
437                        CacheKey {
438                            code_hash: hash,
439                            gas_meter_kind,
440                        },
441                        module.clone(),
442                    );
443
444                    Ok(Some((module, store())))
445                }
446                Err(err) => {
447                    tracing::info!(
448                        "Failed to compile WASM {} with {}",
449                        hash.to_string(),
450                        err
451                    );
452                    let mut progress = self.progress.write().unwrap();
453                    progress.swap_remove(&key);
454                    Err(err)
455                }
456            },
457            Err(err) => {
458                tracing::info!(
459                    "Failed to prepare WASM {} with {}",
460                    hash.to_string(),
461                    err
462                );
463                let mut progress = self.progress.write().unwrap();
464                progress.swap_remove(&key);
465                Err(err)
466            }
467        }
468    }
469
470    /// Pre-compile a WASM module to a file. The compilation runs in a new OS
471    /// thread and the function returns immediately.
472    pub fn pre_compile(
473        &mut self,
474        code: impl AsRef<[u8]>,
475        gas_meter_kind: GasMeterKind,
476    ) {
477        if A::is_read_write() {
478            let hash = hash_of_code(&code);
479            let key = CacheKey {
480                code_hash: hash,
481                gas_meter_kind,
482            };
483            let mut progress = self.progress.write().unwrap();
484            match progress.get(&key) {
485                Some(_) => {
486                    // Already known, do nothing
487                }
488                None => {
489                    if module_file_exists(&self.dir, &hash, gas_meter_kind) {
490                        progress.insert(key, Compilation::Done);
491                        return;
492                    }
493                    progress.insert(key, Compilation::Compiling);
494                    drop(progress);
495                    let progress = self.progress.clone();
496                    let code = code.as_ref().to_vec();
497                    let dir = self.dir.clone();
498                    let store = self.store.clone();
499                    std::thread::spawn(move || {
500                        tracing::info!("Compiling WASM {}.", hash.to_string());
501
502                        let _module = match wasm::run::prepare_wasm_code(
503                            code,
504                            gas_meter_kind,
505                        ) {
506                            Ok(code) => {
507                                match compile(code, &store) {
508                                    Ok(module) => {
509                                        // Write the file
510                                        file_write_module(
511                                            &dir,
512                                            &module,
513                                            &hash,
514                                            gas_meter_kind,
515                                        );
516
517                                        // Update progress
518                                        let mut progress =
519                                            progress.write().unwrap();
520                                        progress.insert(key, Compilation::Done);
521                                        tracing::info!(
522                                            "Finished compiling WASM {hash}."
523                                        );
524                                        if progress.values().all(
525                                            |compilation| {
526                                                matches!(
527                                                    compilation,
528                                                    Compilation::Done
529                                                )
530                                            },
531                                        ) {
532                                            tracing::info!(
533                                                "Finished compiling all {}.",
534                                                N::name()
535                                            )
536                                        }
537                                        module
538                                    }
539                                    Err(err) => {
540                                        let mut progress =
541                                            progress.write().unwrap();
542                                        tracing::info!(
543                                            "Failed to compile WASM {} with {}",
544                                            hash.to_string(),
545                                            err
546                                        );
547                                        progress.swap_remove(&key);
548                                        return Err(err);
549                                    }
550                                }
551                            }
552                            Err(err) => {
553                                let mut progress = progress.write().unwrap();
554                                tracing::info!(
555                                    "Failed to prepare WASM {} with {}",
556                                    hash.to_string(),
557                                    err
558                                );
559                                progress.swap_remove(&key);
560                                return Err(err);
561                            }
562                        };
563
564                        let res: Result<(), wasm::run::Error> = Ok(());
565                        res
566                    });
567                }
568            }
569        }
570    }
571
572    /// Get a read-only cache handle.
573    pub fn read_only(&self) -> Cache<N, WasmCacheRoAccess> {
574        Cache {
575            dir: self.dir.clone(),
576            progress: self.progress.clone(),
577            in_memory: self.in_memory.clone(),
578            name: Default::default(),
579            access: Default::default(),
580            store: self.store.clone(),
581        }
582    }
583}
584
585fn hash_of_code(code: impl AsRef<[u8]>) -> Hash {
586    Hash::sha256(code.as_ref())
587}
588
589fn compile(
590    code: impl AsRef<[u8]>,
591    store: &Store,
592) -> Result<Module, wasm::run::Error> {
593    universal::compile(code, store).map_err(wasm::run::Error::CompileError)
594}
595
596fn file_ext() -> &'static str {
597    // This has to be using the file_ext matching the compilation method in the
598    // `fn compile`
599    universal::FILE_EXT
600}
601
602pub(crate) fn store() -> Store {
603    // This has to be using the store matching the compilation method in the
604    // `fn compile`
605    universal::store()
606}
607
608fn file_write_module(
609    dir: impl AsRef<Path>,
610    module: &Module,
611    hash: &Hash,
612    gas_meter_kind: GasMeterKind,
613) {
614    use wasmer_cache::Cache;
615    let mut fs_cache = fs_cache(dir, hash, gas_meter_kind);
616    fs_cache.store(CacheHash::new(hash.0), module).unwrap();
617}
618
619fn file_load_module(
620    dir: impl AsRef<Path>,
621    hash: &Hash,
622    gas_meter_kind: GasMeterKind,
623    store: &Store,
624) -> Result<Module, wasmer::DeserializeError> {
625    use wasmer_cache::Cache;
626    let fs_cache = fs_cache(dir, hash, gas_meter_kind);
627    let hash = CacheHash::new(hash.0);
628    let module = unsafe { fs_cache.load(store, hash) };
629    if let Err(err) = module.as_ref() {
630        tracing::error!(
631            "Error loading cached wasm {}: {err}.",
632            hash.to_string()
633        );
634    }
635    module
636}
637
638fn fs_cache(
639    dir: impl AsRef<Path>,
640    hash: &Hash,
641    gas_meter_kind: GasMeterKind,
642) -> FileSystemCache {
643    let kind = gas_meter_kind_dir(gas_meter_kind);
644    let path = dir
645        .as_ref()
646        .join(kind)
647        .join(hash.to_string().to_lowercase());
648    let mut fs_cache = FileSystemCache::new(path).unwrap();
649    fs_cache.set_cache_extension(Some(file_ext()));
650    fs_cache
651}
652
653fn gas_meter_kind_dir(gas_meter_kind: GasMeterKind) -> &'static str {
654    match gas_meter_kind {
655        GasMeterKind::HostFn => "host_fn",
656        GasMeterKind::MutGlobal => "mut_global",
657    }
658}
659
660fn module_file_exists(
661    dir: impl AsRef<Path>,
662    hash: &Hash,
663    gas_meter_kind: GasMeterKind,
664) -> bool {
665    let file = dir
666        .as_ref()
667        .join(gas_meter_kind_dir(gas_meter_kind))
668        .join(hash.to_string().to_lowercase())
669        .join(format!(
670            "{}.{}",
671            hash.to_string().to_lowercase(),
672            file_ext()
673        ));
674    file.exists()
675}
676
677/// A universal engine compilation. The module can be serialized to/from bytes.
678mod universal {
679    use super::*;
680
681    #[allow(dead_code)]
682    pub const FILE_EXT: &str = "bin";
683
684    /// Compile wasm with a universal engine.
685    #[allow(dead_code)]
686    pub fn compile(
687        code: impl AsRef<[u8]>,
688        store: &Store,
689    ) -> Result<Module, wasmer::CompileError> {
690        Module::new(store, code.as_ref())
691    }
692
693    /// Universal WASM store
694    #[allow(dead_code)]
695    pub fn store() -> Store {
696        untrusted_wasm_store(memory::vp_limit())
697    }
698}
699
700/// Testing helpers
701#[cfg(any(test, feature = "testing"))]
702pub mod testing {
703    use tempfile::{TempDir, tempdir};
704
705    use super::*;
706    use crate::WasmCacheRwAccess;
707    use crate::wasm::{TxCache, VpCache};
708
709    /// Instantiate the default wasmer store.
710    pub fn store() -> Store {
711        super::store()
712    }
713
714    /// VP Cache with a temp dir for testing
715    pub fn vp_cache() -> (VpCache<WasmCacheRwAccess>, TempDir) {
716        cache::<super::super::vp::Name>()
717    }
718
719    /// Tx Cache with a temp dir for testing
720    pub fn tx_cache() -> (TxCache<WasmCacheRwAccess>, TempDir) {
721        cache::<super::super::tx::Name>()
722    }
723
724    /// Generic Cache with a temp dir for testing
725    pub fn cache<N: CacheName>() -> (Cache<N, WasmCacheRwAccess>, TempDir) {
726        let dir = tempdir().unwrap();
727        let cache = Cache::new(
728            dir.path(),
729            50 * 1024 * 1024, // 50 MiB
730        );
731        (cache, dir)
732    }
733}
734
735#[allow(clippy::arithmetic_side_effects)]
736#[cfg(test)]
737mod test {
738    use std::cmp::max;
739
740    use assert_matches::assert_matches;
741    use byte_unit::{Byte, UnitType};
742    use namada_test_utils::TestWasms;
743    use tempfile::{TempDir, tempdir};
744    use test_log::test;
745
746    use super::*;
747    use crate::WasmCacheRwAccess;
748
749    #[test]
750    fn test_fetch_or_compile_valid_wasm() {
751        // Load some WASMs and find their hashes and in-memory size
752        let tx_read_storage_key = load_wasm(TestWasms::TxReadStorageKey.path());
753        let tx_no_op = load_wasm(TestWasms::TxNoOp.path());
754        let gas_meter_kind = GasMeterKind::MutGlobal;
755        // Create a new cache with the limit set to
756        // `max(tx_read_storage_key.size, tx_no_op.size) + 1`
757        {
758            let max_bytes = max(tx_read_storage_key.size, tx_no_op.size) + 1;
759            println!(
760                "Using cache with max_bytes {} ({})",
761                Byte::from_u128(max_bytes as u128)
762                    .unwrap()
763                    .get_appropriate_unit(UnitType::Binary),
764                max_bytes
765            );
766            let (mut cache, _tmp_dir) = cache(max_bytes);
767
768            // Fetch `tx_read_storage_key`
769            {
770                let fetched = cache
771                    .fetch(&tx_read_storage_key.hash, gas_meter_kind)
772                    .unwrap();
773                assert_matches!(
774                    fetched,
775                    None,
776                    "The module should not be in cache"
777                );
778
779                let fetched = cache
780                    .compile_or_fetch(
781                        &tx_read_storage_key.code,
782                        GasMeterKind::MutGlobal,
783                    )
784                    .unwrap();
785                assert_matches!(
786                    fetched,
787                    Some(_),
788                    "The code should be compiled"
789                );
790
791                let in_memory = cache.in_memory.read().unwrap();
792                assert_matches!(
793                    in_memory.peek(&CacheKey {
794                        code_hash: tx_read_storage_key.hash,
795                        gas_meter_kind
796                    }),
797                    Some(_),
798                    "The module must be in memory"
799                );
800
801                let progress = cache.progress.read().unwrap();
802                assert_matches!(
803                    progress.get(&CacheKey {
804                        code_hash: tx_read_storage_key.hash,
805                        gas_meter_kind
806                    }),
807                    Some(Compilation::Done),
808                    "The progress must be updated"
809                );
810
811                assert!(
812                    module_file_exists(
813                        &cache.dir,
814                        &tx_read_storage_key.hash,
815                        gas_meter_kind
816                    ),
817                    "The file must be written"
818                );
819            }
820
821            // Fetch `tx_no_op`. Fetching another module should get us over the
822            // limit, so the previous one should be popped from the cache
823            {
824                let fetched =
825                    cache.fetch(&tx_no_op.hash, gas_meter_kind).unwrap();
826                assert_matches!(
827                    fetched,
828                    None,
829                    "The module must not be in cache"
830                );
831
832                let fetched = cache
833                    .compile_or_fetch(&tx_no_op.code, GasMeterKind::MutGlobal)
834                    .unwrap();
835                assert_matches!(
836                    fetched,
837                    Some(_),
838                    "The code should be compiled"
839                );
840
841                let in_memory = cache.in_memory.read().unwrap();
842                assert_matches!(
843                    in_memory.peek(&CacheKey {
844                        code_hash: tx_no_op.hash,
845                        gas_meter_kind
846                    }),
847                    Some(_),
848                    "The module must be in memory"
849                );
850
851                let progress = cache.progress.read().unwrap();
852                assert_matches!(
853                    progress.get(&CacheKey {
854                        code_hash: tx_no_op.hash,
855                        gas_meter_kind
856                    }),
857                    Some(Compilation::Done),
858                    "The progress must be updated"
859                );
860
861                assert!(
862                    module_file_exists(
863                        &cache.dir,
864                        &tx_no_op.hash,
865                        gas_meter_kind
866                    ),
867                    "The file must be written"
868                );
869
870                // The previous module's file should still exist
871                assert!(
872                    module_file_exists(
873                        &cache.dir,
874                        &tx_read_storage_key.hash,
875                        gas_meter_kind
876                    ),
877                    "The file must be written"
878                );
879                // But it should not be in-memory
880                assert_matches!(
881                    in_memory.peek(&CacheKey {
882                        code_hash: tx_read_storage_key.hash,
883                        gas_meter_kind
884                    }),
885                    None,
886                    "The module should have been popped from memory"
887                );
888            }
889
890            // Reset the in-memory cache and progress and fetch
891            // `tx_read_storage_key` again, this time it should get loaded
892            // from file
893            let in_memory_cache = CLruCache::with_config(
894                CLruCacheConfig::new(NonZeroUsize::new(max_bytes).unwrap())
895                    .with_scale(ModuleCacheScale),
896            );
897            let in_memory = Arc::new(RwLock::new(in_memory_cache));
898            cache.in_memory = in_memory;
899            cache.progress = Default::default();
900            {
901                let fetched = cache
902                    .fetch(&tx_read_storage_key.hash, gas_meter_kind)
903                    .unwrap();
904                assert_matches!(
905                    fetched,
906                    Some(_),
907                    "The module must be in file cache"
908                );
909
910                let in_memory = cache.in_memory.read().unwrap();
911                assert_matches!(
912                    in_memory.peek(&CacheKey {
913                        code_hash: tx_read_storage_key.hash,
914                        gas_meter_kind
915                    }),
916                    Some(_),
917                    "The module must be in memory"
918                );
919
920                let progress = cache.progress.read().unwrap();
921                assert_matches!(
922                    progress.get(&CacheKey {
923                        code_hash: tx_read_storage_key.hash,
924                        gas_meter_kind
925                    }),
926                    Some(Compilation::Done),
927                    "The progress must be updated"
928                );
929
930                assert!(
931                    module_file_exists(
932                        &cache.dir,
933                        &tx_read_storage_key.hash,
934                        gas_meter_kind
935                    ),
936                    "The file must be written"
937                );
938
939                // The previous module's file should still exist
940                assert!(
941                    module_file_exists(
942                        &cache.dir,
943                        &tx_no_op.hash,
944                        gas_meter_kind
945                    ),
946                    "The file must be written"
947                );
948                // But it should not be in-memory
949                assert_matches!(
950                    in_memory.peek(&CacheKey {
951                        code_hash: tx_no_op.hash,
952                        gas_meter_kind
953                    }),
954                    None,
955                    "The module should have been popped from memory"
956                );
957            }
958
959            // Fetch `tx_read_storage_key` again, now it should be in-memory
960            {
961                let fetched = cache
962                    .fetch(&tx_read_storage_key.hash, gas_meter_kind)
963                    .unwrap();
964                assert_matches!(
965                    fetched,
966                    Some(_),
967                    "The module must be in memory"
968                );
969
970                let in_memory = cache.in_memory.read().unwrap();
971                assert_matches!(
972                    in_memory.peek(&CacheKey {
973                        code_hash: tx_read_storage_key.hash,
974                        gas_meter_kind
975                    }),
976                    Some(_),
977                    "The module must be in memory"
978                );
979
980                let progress = cache.progress.read().unwrap();
981                assert_matches!(
982                    progress.get(&CacheKey {
983                        code_hash: tx_read_storage_key.hash,
984                        gas_meter_kind
985                    }),
986                    Some(Compilation::Done),
987                    "The progress must be updated"
988                );
989
990                assert!(
991                    module_file_exists(
992                        &cache.dir,
993                        &tx_read_storage_key.hash,
994                        gas_meter_kind
995                    ),
996                    "The file must be written"
997                );
998
999                // The previous module's file should still exist
1000                assert!(
1001                    module_file_exists(
1002                        &cache.dir,
1003                        &tx_no_op.hash,
1004                        gas_meter_kind
1005                    ),
1006                    "The file must be written"
1007                );
1008                // But it should not be in-memory
1009                assert_matches!(
1010                    in_memory.peek(&CacheKey {
1011                        code_hash: tx_no_op.hash,
1012                        gas_meter_kind
1013                    }),
1014                    None,
1015                    "The module should have been popped from memory"
1016                );
1017            }
1018
1019            // Fetch `tx_no_op` with read/only access
1020            {
1021                let mut cache = cache.read_only();
1022
1023                let fetched =
1024                    cache.fetch(&tx_no_op.hash, gas_meter_kind).unwrap();
1025                assert_matches!(
1026                    fetched,
1027                    Some(_),
1028                    "The module must be in cache"
1029                );
1030
1031                // Fetching with read-only should not modify the in-memory cache
1032                let fetched = cache
1033                    .compile_or_fetch(&tx_no_op.code, gas_meter_kind)
1034                    .unwrap();
1035                assert_matches!(
1036                    fetched,
1037                    Some(_),
1038                    "The module should be compiled"
1039                );
1040
1041                let in_memory = cache.in_memory.read().unwrap();
1042                assert_matches!(
1043                    in_memory.peek(&CacheKey {
1044                        code_hash: tx_no_op.hash,
1045                        gas_meter_kind
1046                    }),
1047                    None,
1048                    "The module should not be added back to in-memory cache"
1049                );
1050
1051                let in_memory = cache.in_memory.read().unwrap();
1052                assert_matches!(
1053                    in_memory.peek(&CacheKey {
1054                        code_hash: tx_read_storage_key.hash,
1055                        gas_meter_kind
1056                    }),
1057                    Some(_),
1058                    "The previous module must still be in memory"
1059                );
1060            }
1061        }
1062    }
1063
1064    #[test]
1065    fn test_fetch_or_compile_invalid_wasm() {
1066        // Some random bytes
1067        let invalid_wasm = vec![1_u8, 0, 8, 10, 6, 1];
1068        let hash = hash_of_code(&invalid_wasm);
1069        let gas_meter_kind = GasMeterKind::HostFn;
1070        let (mut cache, _) = testing::cache::<TestCache>();
1071
1072        // Try to compile it
1073        let error = cache
1074            .compile_or_fetch(&invalid_wasm, GasMeterKind::MutGlobal)
1075            .expect_err("Compilation should fail");
1076        println!("Error: {}", error);
1077
1078        let in_memory = cache.in_memory.read().unwrap();
1079        assert_matches!(
1080            in_memory.peek(&CacheKey {
1081                code_hash: hash,
1082                gas_meter_kind
1083            }),
1084            None,
1085            "There should be no entry for this hash in memory"
1086        );
1087
1088        let progress = cache.progress.read().unwrap();
1089        assert_matches!(
1090            progress.get(&CacheKey {
1091                code_hash: hash,
1092                gas_meter_kind
1093            }),
1094            None,
1095            "Any progress is removed"
1096        );
1097
1098        assert!(
1099            !module_file_exists(&cache.dir, &hash, gas_meter_kind),
1100            "The file must not be written"
1101        );
1102    }
1103
1104    #[test]
1105    fn test_pre_compile_valid_wasm() {
1106        // Load some WASMs and find their hashes and in-memory size
1107        let vp_always_true = load_wasm(TestWasms::VpAlwaysTrue.path());
1108        let vp_eval = load_wasm(TestWasms::VpEval.path());
1109
1110        // Create a new cache with the limit set to
1111        // `max(vp_always_true.size, vp_eval.size) + 1 + extra_bytes`
1112        {
1113            let max_bytes = max(vp_always_true.size, vp_eval.size) + 1;
1114            println!(
1115                "Using cache with max_bytes {} ({})",
1116                Byte::from_u128(max_bytes as u128)
1117                    .unwrap()
1118                    .get_appropriate_unit(UnitType::Binary),
1119                max_bytes
1120            );
1121            let (mut cache, _tmp_dir) = cache(max_bytes);
1122            let gas_meter_kind = GasMeterKind::MutGlobal;
1123
1124            // Pre-compile `vp_always_true`
1125            {
1126                cache
1127                    .pre_compile(&vp_always_true.code, GasMeterKind::MutGlobal);
1128
1129                let progress = cache.progress.read().unwrap();
1130                assert_matches!(
1131                    progress.get(&CacheKey {
1132                        code_hash: vp_always_true.hash,
1133                        gas_meter_kind
1134                    }),
1135                    Some(Compilation::Done | Compilation::Compiling),
1136                    "The progress must be updated"
1137                );
1138            }
1139
1140            // Now fetch it to wait for it finish compilation
1141            {
1142                let fetched =
1143                    cache.fetch(&vp_always_true.hash, gas_meter_kind).unwrap();
1144                assert_matches!(
1145                    fetched,
1146                    Some(_),
1147                    "The module must be in cache"
1148                );
1149
1150                let in_memory = cache.in_memory.read().unwrap();
1151                assert_matches!(
1152                    in_memory.peek(&CacheKey {
1153                        code_hash: vp_always_true.hash,
1154                        gas_meter_kind
1155                    }),
1156                    Some(_),
1157                    "The module must be in memory"
1158                );
1159
1160                let progress = cache.progress.read().unwrap();
1161                assert_matches!(
1162                    progress.get(&CacheKey {
1163                        code_hash: vp_always_true.hash,
1164                        gas_meter_kind
1165                    }),
1166                    Some(Compilation::Done),
1167                    "The progress must be updated"
1168                );
1169
1170                assert!(
1171                    module_file_exists(
1172                        &cache.dir,
1173                        &vp_always_true.hash,
1174                        gas_meter_kind
1175                    ),
1176                    "The file must be written"
1177                );
1178            }
1179
1180            // Pre-compile `vp_eval`. Pre-compiling another module should get us
1181            // over the limit, so the previous one should be popped
1182            // from the cache
1183            {
1184                cache.pre_compile(&vp_eval.code, GasMeterKind::MutGlobal);
1185
1186                let progress = cache.progress.read().unwrap();
1187                assert_matches!(
1188                    progress.get(&CacheKey {
1189                        code_hash: vp_eval.hash,
1190                        gas_meter_kind
1191                    }),
1192                    Some(Compilation::Done | Compilation::Compiling),
1193                    "The progress must be updated"
1194                );
1195            }
1196
1197            // Now fetch it to wait for it finish compilation
1198            {
1199                let fetched =
1200                    cache.fetch(&vp_eval.hash, gas_meter_kind).unwrap();
1201                assert_matches!(
1202                    fetched,
1203                    Some(_),
1204                    "The module must be in cache"
1205                );
1206
1207                let in_memory = cache.in_memory.read().unwrap();
1208                assert_matches!(
1209                    in_memory.peek(&CacheKey {
1210                        code_hash: vp_eval.hash,
1211                        gas_meter_kind
1212                    }),
1213                    Some(_),
1214                    "The module must be in memory"
1215                );
1216
1217                assert!(
1218                    module_file_exists(
1219                        &cache.dir,
1220                        &vp_eval.hash,
1221                        gas_meter_kind
1222                    ),
1223                    "The file must be written"
1224                );
1225
1226                // The previous module's file should still exist
1227                assert!(
1228                    module_file_exists(
1229                        &cache.dir,
1230                        &vp_always_true.hash,
1231                        gas_meter_kind
1232                    ),
1233                    "The file must be written"
1234                );
1235                // But it should not be in-memory
1236                assert_matches!(
1237                    in_memory.peek(&CacheKey {
1238                        code_hash: vp_always_true.hash,
1239                        gas_meter_kind
1240                    }),
1241                    None,
1242                    "The module should have been popped from memory"
1243                );
1244            }
1245        }
1246    }
1247
1248    #[test]
1249    fn test_pre_compile_invalid_wasm() {
1250        // Some random bytes
1251        let invalid_wasm = vec![1_u8];
1252        let hash = hash_of_code(&invalid_wasm);
1253        let (mut cache, _) = testing::cache::<TestCache>();
1254        let gas_meter_kind = GasMeterKind::HostFn;
1255
1256        // Try to pre-compile it
1257        {
1258            cache.pre_compile(&invalid_wasm, GasMeterKind::MutGlobal);
1259            let progress = cache.progress.read().unwrap();
1260            assert_matches!(
1261                progress.get(&CacheKey {
1262                    code_hash: hash,
1263                    gas_meter_kind
1264                }),
1265                Some(Compilation::Done | Compilation::Compiling) | None,
1266                "The progress must be updated"
1267            );
1268        }
1269
1270        // Now fetch it to wait for it finish compilation
1271        {
1272            let fetched = cache.fetch(&hash, gas_meter_kind).unwrap();
1273            assert_matches!(
1274                fetched,
1275                None,
1276                "There should be no entry for this hash in cache"
1277            );
1278
1279            let in_memory = cache.in_memory.read().unwrap();
1280            assert_matches!(
1281                in_memory.peek(&CacheKey {
1282                    code_hash: hash,
1283                    gas_meter_kind
1284                }),
1285                None,
1286                "There should be no entry for this hash in memory"
1287            );
1288
1289            let progress = cache.progress.read().unwrap();
1290            assert_matches!(
1291                progress.get(&CacheKey {
1292                    code_hash: hash,
1293                    gas_meter_kind
1294                }),
1295                None,
1296                "Any progress is removed"
1297            );
1298
1299            assert!(
1300                !module_file_exists(&cache.dir, &hash, gas_meter_kind),
1301                "The file must not be written"
1302            );
1303        }
1304    }
1305
1306    /// Get the WASM code bytes, its hash and find the compiled module's size
1307    fn load_wasm(file: impl AsRef<Path>) -> WasmWithMeta {
1308        let file = file.as_ref();
1309        let code = fs::read(file).unwrap();
1310        let hash = hash_of_code(&code);
1311        // Find the size of the compiled module
1312        let size = {
1313            let (mut cache, _tmp_dir) = cache(
1314                // No in-memory cache needed, but must be non-zero
1315                1,
1316            );
1317            let (_module, _store) = cache
1318                .compile_or_fetch(&code, GasMeterKind::MutGlobal)
1319                .unwrap()
1320                .unwrap();
1321            1
1322        };
1323        println!(
1324            "Compiled module {} size including the hash: {} ({})",
1325            file.to_string_lossy(),
1326            Byte::from_u128(size as u128)
1327                .unwrap()
1328                .get_appropriate_unit(UnitType::Binary),
1329            size,
1330        );
1331        WasmWithMeta { code, hash, size }
1332    }
1333
1334    /// A test helper for loading WASM and finding its hash and size
1335    #[derive(Clone, Debug)]
1336    struct WasmWithMeta {
1337        code: Vec<u8>,
1338        hash: Hash,
1339        /// Compiled module's in-memory size
1340        size: usize,
1341    }
1342
1343    /// A `CacheName` implementation for unit tests
1344    #[derive(Clone, Debug)]
1345    struct TestCache;
1346    impl CacheName for TestCache {
1347        fn name() -> &'static str {
1348            "test"
1349        }
1350    }
1351
1352    /// A cache with a temp dir for unit tests
1353    fn cache(
1354        max_bytes: usize,
1355    ) -> (Cache<TestCache, WasmCacheRwAccess>, TempDir) {
1356        let dir = tempdir().unwrap();
1357        let cache = Cache::new(dir.path(), max_bytes);
1358        (cache, dir)
1359    }
1360}