Feat/cache weight allocated size

crates/e2e-move-tests/Cargo.toml

@@ -46,6 +46,5 @@ default = []
 testing = [
     "initia-move-gas/testing",
     "initia-move-natives/testing",
-    "initia-move-storage/testing",
     "move-vm-runtime/testing",
 ]

crates/storage/Cargo.toml

@@ -22,9 +22,3 @@ clru = { workspace = true }
 parking_lot = { workspace = true }
 sha3 = { workspace = true }
 claims = { workspace = true }
-
-[features]
-default = []
-testing = [
-    "move-vm-runtime/testing",
-]

crates/storage/src/allocator.rs

@@ -0,0 +1,32 @@
+use std::{
+    alloc::{GlobalAlloc, Layout, System},
+    cell::Cell,
+};
+
+thread_local! {
+    static SIZE: Cell<usize> = const { Cell::new(0) };
+}
+
+struct SizeCounterAllocator;
+
+unsafe impl GlobalAlloc for SizeCounterAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        SIZE.with(|size| size.set(size.get() + layout.size()));
+        System.alloc(layout)
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        System.dealloc(ptr, layout)
+    }
+}
+
+#[global_allocator]
+static GLOBAL: SizeCounterAllocator = SizeCounterAllocator;
+
+pub(crate) fn initialize_size() {
+    SIZE.with(|size| size.set(0));
+}
+
+pub(crate) fn get_size() -> usize {
+    SIZE.with(|size| size.get())
+}

crates/storage/src/code_scale.rs

@@ -13,22 +13,43 @@ use crate::state_view::Checksum;
 
 pub struct CodeScale;
 
-impl WeightScale<Checksum, Code<CompiledScript, Script>> for CodeScale {
-    fn weight(&self, _key: &Checksum, _value: &Code<CompiledScript, Script>) -> usize {
-        1
+impl WeightScale<Checksum, CodeWrapper> for CodeScale {
+    fn weight(&self, _key: &Checksum, value: &CodeWrapper) -> usize {
+        value.size
     }
 }
 
 pub struct ModuleCodeScale;
 
-impl WeightScale<Checksum, Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>>
-    for ModuleCodeScale
-{
-    fn weight(
-        &self,
-        _key: &Checksum,
-        _value: &Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
-    ) -> usize {
-        1
+impl WeightScale<Checksum, ModuleCodeWrapper> for ModuleCodeScale {
+    fn weight(&self, _key: &Checksum, value: &ModuleCodeWrapper) -> usize {
+        value.size
+    }
+}
+
+#[derive(Clone)]
+pub struct CodeWrapper {
+    pub code: Code<CompiledScript, Script>,
+    pub size: usize,
+}
+
+impl CodeWrapper {
+    pub fn new(code: Code<CompiledScript, Script>, size: usize) -> Self {
+        Self { code, size }
+    }
+}
+
+#[derive(Clone)]
+pub struct ModuleCodeWrapper {
+    pub module_code: Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
+    pub size: usize,
+}
+
+impl ModuleCodeWrapper {
+    pub fn new(
+        module_code: Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
+        size: usize,
+    ) -> Self {
+        Self { module_code, size }
     }
 }

crates/storage/src/code_storage.rs

@@ -10,13 +10,11 @@ use move_vm_runtime::{
     logging::expect_no_verification_errors, CodeStorage, Module, ModuleStorage, RuntimeEnvironment,
     Script, WithRuntimeEnvironment,
 };
-use move_vm_types::{
-    code::{Code, ModuleBytesStorage},
-    module_linker_error, sha3_256,
-};
+use move_vm_types::{code::ModuleBytesStorage, module_linker_error, sha3_256};
 use std::sync::Arc;
 
 use crate::{
+    allocator::{get_size, initialize_size},
     module_cache::InitiaModuleCache,
     module_storage::{AsInitiaModuleStorage, InitiaModuleStorage},
     script_cache::InitiaScriptCache,
@@ -98,29 +96,42 @@ impl<M: ModuleStorage> CodeStorage for InitiaCodeStorage<M> {
     ) -> VMResult<Arc<CompiledScript>> {
         let hash = sha3_256(serialized_script);
         Ok(match self.script_cache.get_script(&hash) {
-            Some(script) => script.deserialized().clone(),
+            Some(script) => script.code.deserialized().clone(),
             None => {
+                initialize_size();
                 let deserialized_script = self
                     .runtime_environment()
                     .deserialize_into_script(serialized_script)?;
-                self.script_cache
-                    .insert_deserialized_script(hash, deserialized_script)?
+                let allocated_size = get_size();
+                self.script_cache.insert_deserialized_script(
+                    hash,
+                    deserialized_script,
+                    allocated_size,
+                )?
             }
         })
     }
 
     fn verify_and_cache_script(&self, serialized_script: &[u8]) -> VMResult<Arc<Script>> {
-        use Code::*;
-
         let hash = sha3_256(serialized_script);
-        let deserialized_script = match self.script_cache.get_script(&hash) {
-            Some(Verified(script)) => return Ok(script),
-            Some(Deserialized(deserialized_script)) => deserialized_script,
-            None => self
-                .runtime_environment()
-                .deserialize_into_script(serialized_script)
-                .map(Arc::new)?,
-        };
+        let (deserialized_script, compiled_script_allocated_size) =
+            match self.script_cache.get_script(&hash) {
+                Some(code_wrapper) => {
+                    if code_wrapper.code.is_verified() {
+                        return Ok(code_wrapper.code.verified().clone());
+                    }
+                    (code_wrapper.code.deserialized().clone(), code_wrapper.size)
+                }
+                None => {
+                    initialize_size();
+                    let compiled_script = self
+                        .runtime_environment()
+                        .deserialize_into_script(serialized_script)
+                        .map(Arc::new)?;
+                    let allocated_size = get_size();
+                    (compiled_script, allocated_size)
+                }
+            };
 
         // Locally verify the script.
         let locally_verified_script = self
@@ -137,12 +148,15 @@ impl<M: ModuleStorage> CodeStorage for InitiaCodeStorage<M> {
                     .ok_or_else(|| module_linker_error!(addr, name))
             })
             .collect::<VMResult<Vec<_>>>()?;
+
+        initialize_size();
         let verified_script = self
             .runtime_environment()
             .build_verified_script(locally_verified_script, &immediate_dependencies)?;
+        let allocated_size = get_size() + compiled_script_allocated_size;
 
         self.script_cache
-            .insert_verified_script(hash, verified_script)
+            .insert_verified_script(hash, verified_script, allocated_size)
     }
 }
 
@@ -162,11 +176,11 @@ impl<M: ModuleStorage> InitiaCodeStorage<M> {
         );
         for hash in deserialized {
             let script = claims::assert_some!(self.script_cache.get_script(hash));
-            assert!(!script.is_verified())
+            assert!(!script.code.is_verified())
         }
         for hash in verified {
             let script = claims::assert_some!(self.script_cache.get_script(hash));
-            assert!(script.is_verified())
+            assert!(script.code.is_verified())
         }
     }
 }

crates/storage/src/lib.rs

@@ -9,4 +9,5 @@ pub mod module_cache;
 pub mod module_storage;
 pub mod script_cache;
 
+mod allocator;
 pub mod code_scale;

crates/storage/src/module_cache.rs

@@ -11,7 +11,11 @@ use move_vm_runtime::Module;
 use move_vm_types::code::{ModuleCode, ModuleCodeBuilder, WithBytes, WithHash};
 use parking_lot::Mutex;
 
-use crate::{code_scale::ModuleCodeScale, state_view::Checksum};
+use crate::{
+    allocator::{get_size, initialize_size},
+    code_scale::{ModuleCodeScale, ModuleCodeWrapper},
+    state_view::Checksum,
+};
 
 fn handle_cache_error(module_id: ModuleId) -> VMError {
     PartialVMError::new(StatusCode::MEMORY_LIMIT_EXCEEDED)
@@ -54,19 +58,13 @@ pub struct NoVersion;
 
 pub struct InitiaModuleCache {
     #[allow(clippy::type_complexity)]
-    pub(crate) module_cache: Mutex<
-        CLruCache<
-            Checksum,
-            Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
-            RandomState,
-            ModuleCodeScale,
-        >,
-    >,
+    pub(crate) module_cache:
+        Mutex<CLruCache<Checksum, ModuleCodeWrapper, RandomState, ModuleCodeScale>>,
 }
 
 impl InitiaModuleCache {
     pub fn new(cache_capacity: usize) -> Arc<InitiaModuleCache> {
-        let capacity = NonZeroUsize::new(cache_capacity).unwrap();
+        let capacity = NonZeroUsize::new(cache_capacity * 1024 * 1024).unwrap();
         Arc::new(InitiaModuleCache {
             module_cache: Mutex::new(CLruCache::with_config(
                 CLruCacheConfig::new(capacity).with_scale(ModuleCodeScale),
@@ -82,6 +80,7 @@ impl InitiaModuleCache {
         &self,
         key: Checksum,
         deserialized_code: CompiledModule,
+        allocated_size: usize,
         extension: Arc<BytesWithHash>,
         version: NoVersion,
     ) -> VMResult<()> {
@@ -98,7 +97,7 @@ impl InitiaModuleCache {
                     version,
                 ));
                 module_cache
-                    .put_with_weight(key, module)
+                    .put_with_weight(key, ModuleCodeWrapper::new(module, allocated_size))
                     .map_err(|_| handle_cache_error(module_id))?;
                 Ok(())
             }
@@ -109,21 +108,25 @@ impl InitiaModuleCache {
         &self,
         key: Checksum,
         verified_code: Module,
+        allocated_size: usize,
         extension: Arc<BytesWithHash>,
         version: NoVersion,
     ) -> VMResult<Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>> {
         let mut module_cache = self.module_cache.lock();
 
         match module_cache.get(&key) {
-            Some(code) => {
-                if code.code().is_verified() {
-                    Ok(code.clone())
+            Some(code_wrapper) => {
+                if code_wrapper.module_code.code().is_verified() {
+                    Ok(code_wrapper.module_code.clone())
                 } else {
                     let module_id = verified_code.self_id();
                     let module =
                         Arc::new(ModuleCode::from_verified(verified_code, extension, version));
                     module_cache
-                        .put_with_weight(key, module.clone())
+                        .put_with_weight(
+                            key,
+                            ModuleCodeWrapper::new(module.clone(), allocated_size),
+                        )
                         .map_err(|_| handle_cache_error(module_id))?;
                     Ok(module)
                 }
@@ -132,7 +135,7 @@ impl InitiaModuleCache {
                 let module_id = verified_code.self_id();
                 let module = Arc::new(ModuleCode::from_verified(verified_code, extension, version));
                 module_cache
-                    .put_with_weight(key, module.clone())
+                    .put_with_weight(key, ModuleCodeWrapper::new(module.clone(), allocated_size))
                     .map_err(|_| handle_cache_error(module_id))?;
                 Ok(module)
             }
@@ -151,32 +154,38 @@ impl InitiaModuleCache {
             Extension = BytesWithHash,
             Version = NoVersion,
         >,
-    ) -> VMResult<Option<Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>>> {
+    ) -> VMResult<Option<ModuleCodeWrapper>> {
         let mut module_cache = self.module_cache.lock();
         Ok(match module_cache.get(checksum) {
-            Some(code) => Some(code.clone()),
-            None => match builder.build(id)? {
-                Some(code) => {
-                    if code.extension().hash() != checksum {
-                        return Err(PartialVMError::new(
-                            StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR,
-                        )
-                        .with_message("Checksum mismatch".to_string())
-                        .finish(Location::Module(id.clone())));
+            Some(code) => Some(ModuleCodeWrapper::new(code.module_code.clone(), code.size)),
+            None => {
+                initialize_size();
+                let build_result = builder.build(id)?;
+                let allocated_size = get_size();
+                match build_result {
+                    Some(code) => {
+                        if code.extension().hash() != checksum {
+                            return Err(PartialVMError::new(
+                                StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR,
+                            )
+                            .with_message("Checksum mismatch".to_string())
+                            .finish(Location::Module(id.clone())));
+                        }
+
+                        let code = Arc::new(code);
+                        let code_wrapper = ModuleCodeWrapper::new(code.clone(), allocated_size);
+                        module_cache
+                            .put_with_weight(*checksum, code_wrapper.clone())
+                            .map_err(|_| {
+                                PartialVMError::new(StatusCode::MEMORY_LIMIT_EXCEEDED)
+                                    .with_message("Module storage cache eviction error".to_string())
+                                    .finish(Location::Module(id.clone()))
+                            })?;
+                        Some(code_wrapper)
                     }
-
-                    let code = Arc::new(code);
-                    module_cache
-                        .put_with_weight(*checksum, code.clone())
-                        .map_err(|_| {
-                            PartialVMError::new(StatusCode::MEMORY_LIMIT_EXCEEDED)
-                                .with_message("Module storage cache eviction error".to_string())
-                                .finish(Location::Module(id.clone()))
-                        })?;
-                    Some(code)
+                    None => None,
                 }
-                None => None,
-            },
+            }
         })
     }