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

@@ -23,8 +23,5 @@ parking_lot = { workspace = true }
 sha3 = { workspace = true }
 claims = { workspace = true }
 
-[features]
-default = []
-testing = [
-    "move-vm-runtime/testing",
-]
+[dev-dependencies]
+rand = { workspace = true }

crates/storage/src/allocator.rs

@@ -0,0 +1,82 @@
+use std::{
+    alloc::{GlobalAlloc, Layout, System},
+    cell::Cell,
+};
+
+use move_binary_format::errors::VMResult;
+
+thread_local! {
+    static METERING: Cell<bool> = const { Cell::new(false) };
+    static SIZE: Cell<usize> = const { Cell::new(0) };
+}
+
+struct SizeCounterAllocator;
+
+unsafe impl GlobalAlloc for SizeCounterAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        if METERING.with(|metering| metering.get()) {
+            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;
+
+#[inline]
+fn start_metering() {
+    SIZE.with(|size| size.set(0));
+    METERING.with(|metering| metering.set(true));
+}
+
+#[inline]
+fn finish_metering() -> usize {
+    METERING.with(|metering| metering.set(false));
+    SIZE.with(|size| size.get())
+}
+
+#[inline]
+pub(crate) fn get_size<T, O: FnOnce() -> VMResult<T>>(f: O) -> VMResult<(T, usize)> {
+    start_metering();
+    let ret = f()?;
+    let size = finish_metering();
+
+    Ok((ret, size + size_of::<T>()))
+}
+
+#[cfg(test)]
+mod allocator_test {
+    use rand::Rng;
+    use std::thread;
+
+    use super::*;
+
+    #[test]
+    fn test_get_size() {
+        let num_thread = 100;
+        for _ in 0..num_thread {
+            let handle = thread::spawn(|| {
+                let num_bytes = rand::thread_rng().gen_range(0..5120); // < 5KB
+                let (_, size) = get_size(|| {
+                    for _ in 0..num_bytes {
+                        // allocate 1 byte
+                        let _ = vec![0u8; 1];
+                    }
+
+                    Ok(())
+                })
+                .unwrap();
+
+                assert_eq!(size, num_bytes);
+            });
+
+            handle.join().unwrap();
+        }
+    }
+}

crates/storage/src/code_scale.rs

@@ -11,24 +11,45 @@ use crate::module_cache::BytesWithHash;
 use crate::module_cache::NoVersion;
 use crate::state_view::Checksum;
 
-pub struct CodeScale;
+pub struct ScriptScale;
 
-impl WeightScale<Checksum, Code<CompiledScript, Script>> for CodeScale {
-    fn weight(&self, _key: &Checksum, _value: &Code<CompiledScript, Script>) -> usize {
-        1
+impl WeightScale<Checksum, ScriptWrapper> for ScriptScale {
+    fn weight(&self, _key: &Checksum, value: &ScriptWrapper) -> 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
+pub struct ModuleScale;
+
+impl WeightScale<Checksum, ModuleWrapper> for ModuleScale {
+    fn weight(&self, _key: &Checksum, value: &ModuleWrapper) -> usize {
+        value.size
+    }
+}
+
+#[derive(Clone)]
+pub struct ScriptWrapper {
+    pub code: Code<CompiledScript, Script>,
+    pub size: usize,
+}
+
+impl ScriptWrapper {
+    pub fn new(code: Code<CompiledScript, Script>, size: usize) -> Self {
+        Self { code, size }
+    }
+}
+
+#[derive(Clone)]
+pub struct ModuleWrapper {
+    pub module_code: Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
+    pub size: usize,
+}
+
+impl ModuleWrapper {
+    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,
     module_cache::InitiaModuleCache,
     module_storage::{AsInitiaModuleStorage, InitiaModuleStorage},
     script_cache::InitiaScriptCache,
@@ -98,29 +96,43 @@ 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 => {
-                let deserialized_script = self
-                    .runtime_environment()
-                    .deserialize_into_script(serialized_script)?;
-                self.script_cache
-                    .insert_deserialized_script(hash, deserialized_script)?
+                // Deserialize the script and compute its size.
+                let (deserialized_script, allocated_size) = get_size(move || {
+                    self.runtime_environment()
+                        .deserialize_into_script(serialized_script)
+                })?;
+
+                // Cache the deserialized script.
+                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 => {
+                    get_size(move || {
+                        self.runtime_environment()
+                            .deserialize_into_script(serialized_script)
+                            .map(Arc::new)
+                    })?
+                }
+            };
 
         // Locally verify the script.
         let locally_verified_script = self
@@ -137,12 +149,19 @@ impl<M: ModuleStorage> CodeStorage for InitiaCodeStorage<M> {
                     .ok_or_else(|| module_linker_error!(addr, name))
             })
             .collect::<VMResult<Vec<_>>>()?;
-        let verified_script = self
-            .runtime_environment()
-            .build_verified_script(locally_verified_script, &immediate_dependencies)?;
 
-        self.script_cache
-            .insert_verified_script(hash, verified_script)
+        // Verify the script and compute its size.
+        let (verified_script, allocated_size) = get_size(move || {
+            self.runtime_environment()
+                .build_verified_script(locally_verified_script, &immediate_dependencies)
+        })?;
+
+        // Cache the verified script.
+        self.script_cache.insert_verified_script(
+            hash,
+            verified_script,
+            allocated_size + compiled_script_allocated_size,
+        )
     }
 }
 
@@ -162,11 +181,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;