Add PinnedMemoryCache::size

Author: webmaster128

packages/vm/src/cache.rs

@@ -123,7 +123,7 @@ where
             stats: cache.stats,
             elements_pinned_memory_cache: cache.pinned_memory_cache.len(),
             elements_memory_cache: cache.memory_cache.len(),
-            size_pinned_memory_cache: 0,
+            size_pinned_memory_cache: cache.pinned_memory_cache.size(),
             size_memory_cache: cache.memory_cache.size(),
         }
     }
@@ -185,22 +185,28 @@ where
         // Try to get module from the memory cache
         if let Some(module) = cache.memory_cache.load(checksum)? {
             cache.stats.hits_memory_cache += 1;
-            return cache.pinned_memory_cache.store(checksum, module);
+            return cache
+                .pinned_memory_cache
+                .store(checksum, module.module, module.size);
         }
 
         // Try to get module from file system cache
         let store = make_runtime_store(Some(cache.instance_memory_limit));
-        if let Some((module, _)) = cache.fs_cache.load(checksum, &store)? {
+        if let Some((module, module_size)) = cache.fs_cache.load(checksum, &store)? {
             cache.stats.hits_fs_cache += 1;
-            return cache.pinned_memory_cache.store(checksum, module);
+            return cache
+                .pinned_memory_cache
+                .store(checksum, module, module_size);
         }
 
         // Re-compile from original Wasm bytecode
         let code = self.load_wasm_with_path(&cache.wasm_path, checksum)?;
         let module = compile(&code, Some(cache.instance_memory_limit))?;
         // Store into the fs cache too
-        cache.fs_cache.store(checksum, &module)?;
-        cache.pinned_memory_cache.store(checksum, module)
+        let module_size = cache.fs_cache.store(checksum, &module)?;
+        cache
+            .pinned_memory_cache
+            .store(checksum, module, module_size)
     }
 
     /// Unpins a Module, i.e. removes it from the pinned memory cache.
@@ -235,8 +241,12 @@ where
         // Get module from memory cache
         if let Some(module) = cache.memory_cache.load(checksum)? {
             cache.stats.hits_memory_cache += 1;
-            let instance =
-                Instance::from_module(&module, backend, options.gas_limit, options.print_debug)?;
+            let instance = Instance::from_module(
+                &module.module,
+                backend,
+                options.gas_limit,
+                options.print_debug,
+            )?;
             return Ok(instance);
         }
 

packages/vm/src/modules/in_memory_cache.rs

@@ -1,9 +1,10 @@
 use clru::{CLruCache, CLruCacheConfig, WeightScale};
 use std::collections::hash_map::RandomState;
+use std::num::NonZeroUsize;
 use wasmer::Module;
 
+use super::sized_module::SizedModule;
 use crate::{Checksum, Size, VmError, VmResult};
-use std::num::NonZeroUsize;
 
 // Minimum module size.
 // Based on `examples/module_size.sh`, and the cosmwasm-plus contracts.
@@ -14,12 +15,6 @@ use std::num::NonZeroUsize;
 // Which is a very small percentage (~0.03%) of our typical cache memory budget (2 GB).
 const MINIMUM_MODULE_SIZE: Size = Size::kibi(250);
 
-#[derive(Debug)]
-struct SizedModule {
-    pub module: Module,
-    pub size: usize,
-}
-
 #[derive(Debug)]
 struct SizeScale;
 
@@ -64,10 +59,10 @@ impl InMemoryCache {
     }
 
     /// Looks up a module in the cache and creates a new module
-    pub fn load(&mut self, checksum: &Checksum) -> VmResult<Option<Module>> {
+    pub fn load(&mut self, checksum: &Checksum) -> VmResult<Option<SizedModule>> {
         if let Some(modules) = &mut self.modules {
             match modules.get(checksum) {
-                Some(sized_module) => Ok(Some(sized_module.module.clone())),
+                Some(module) => Ok(Some(module.clone())),
                 None => Ok(None),
             }
         } else {
@@ -167,7 +162,7 @@ mod tests {
 
         // Ensure cached module can be executed
         {
-            let instance = WasmerInstance::new(&cached, &imports! {}).unwrap();
+            let instance = WasmerInstance::new(&cached.module, &imports! {}).unwrap();
             set_remaining_points(&instance, TESTING_GAS_LIMIT);
             let add_one = instance.exports.get_function("add_one").unwrap();
             let result = add_one.call(&[42.into()]).unwrap();

packages/vm/src/modules/mod.rs

@@ -1,6 +1,7 @@
 mod file_system_cache;
 mod in_memory_cache;
 mod pinned_memory_cache;
+mod sized_module;
 
 pub use file_system_cache::FileSystemCache;
 pub use in_memory_cache::InMemoryCache;

packages/vm/src/modules/pinned_memory_cache.rs

@@ -1,11 +1,12 @@
 use std::collections::HashMap;
 use wasmer::Module;
 
+use super::sized_module::SizedModule;
 use crate::{Checksum, VmResult};
 
 /// An pinned in memory module cache
 pub struct PinnedMemoryCache {
-    modules: HashMap<Checksum, Module>,
+    modules: HashMap<Checksum, SizedModule>,
 }
 
 impl PinnedMemoryCache {
@@ -16,8 +17,8 @@ impl PinnedMemoryCache {
         }
     }
 
-    pub fn store(&mut self, checksum: &Checksum, module: Module) -> VmResult<()> {
-        self.modules.insert(*checksum, module);
+    pub fn store(&mut self, checksum: &Checksum, module: Module, size: usize) -> VmResult<()> {
+        self.modules.insert(*checksum, SizedModule { module, size });
         Ok(())
     }
 
@@ -31,7 +32,7 @@ impl PinnedMemoryCache {
     /// Looks up a module in the cache and creates a new module
     pub fn load(&mut self, checksum: &Checksum) -> VmResult<Option<Module>> {
         match self.modules.get(checksum) {
-            Some(module) => Ok(Some(module.clone())),
+            Some(module) => Ok(Some(module.module.clone())),
             None => Ok(None),
         }
     }
@@ -45,6 +46,14 @@ impl PinnedMemoryCache {
     pub fn len(&self) -> usize {
         self.modules.len()
     }
+
+    /// Returns cumulative size of all elements in the cache.
+    ///
+    /// This is based on the values provided with `store`. No actual
+    /// memory size is measured here.
+    pub fn size(&self) -> usize {
+        self.modules.iter().map(|(_, module)| module.size).sum()
+    }
 }
 
 #[cfg(test)]
@@ -90,7 +99,7 @@ mod tests {
         }
 
         // Store module
-        cache.store(&checksum, original).unwrap();
+        cache.store(&checksum, original, 0).unwrap();
 
         // Load module
         let cached = cache.load(&checksum).unwrap().unwrap();
@@ -126,7 +135,7 @@ mod tests {
 
         // Add
         let original = compile(&wasm, None).unwrap();
-        cache.store(&checksum, original).unwrap();
+        cache.store(&checksum, original, 0).unwrap();
 
         assert_eq!(cache.has(&checksum), true);
 
@@ -157,7 +166,7 @@ mod tests {
 
         // Add
         let original = compile(&wasm, None).unwrap();
-        cache.store(&checksum, original).unwrap();
+        cache.store(&checksum, original, 0).unwrap();
 
         assert_eq!(cache.len(), 1);
 
@@ -166,4 +175,53 @@ mod tests {
 
         assert_eq!(cache.len(), 0);
     }
+
+    #[test]
+    fn size_works() {
+        let mut cache = PinnedMemoryCache::new();
+
+        // Create module
+        let wasm1 = wat::parse_str(
+            r#"(module
+            (type $t0 (func (param i32) (result i32)))
+            (func $add_one (export "add_one") (type $t0) (param $p0 i32) (result i32)
+                get_local $p0
+                i32.const 1
+                i32.add)
+            )"#,
+        )
+        .unwrap();
+        let checksum1 = Checksum::generate(&wasm1);
+        let wasm2 = wat::parse_str(
+            r#"(module
+            (type $t0 (func (param i32) (result i32)))
+            (func $add_one (export "add_two") (type $t0) (param $p0 i32) (result i32)
+                get_local $p0
+                i32.const 2
+                i32.add)
+            )"#,
+        )
+        .unwrap();
+        let checksum2 = Checksum::generate(&wasm2);
+
+        assert_eq!(cache.size(), 0);
+
+        // Add 1
+        let original = compile(&wasm1, None).unwrap();
+        cache.store(&checksum1, original, 500).unwrap();
+        assert_eq!(cache.size(), 500);
+
+        // Add 2
+        let original = compile(&wasm2, None).unwrap();
+        cache.store(&checksum2, original, 300).unwrap();
+        assert_eq!(cache.size(), 800);
+
+        // Remove 1
+        cache.remove(&checksum1).unwrap();
+        assert_eq!(cache.size(), 300);
+
+        // Remove 2
+        cache.remove(&checksum2).unwrap();
+        assert_eq!(cache.size(), 0);
+    }
 }

packages/vm/src/modules/sized_module.rs

@@ -0,0 +1,7 @@
+use wasmer::Module;
+
+#[derive(Debug, Clone)]
+pub struct SizedModule {
+    pub module: Module,
+    pub size: usize,
+}