calculate cache module & script allocated size

Cargo.toml

@@ -114,6 +114,8 @@ ripemd = "0.1.1"
 tiny-keccak = { version = "2.0.2", features = ["keccak", "sha3"] }
 get-size = { version = "0.1.4", features = ["derive"]}
 pretty_assertions = "1.2.1"
+dashmap = { version = "5.5.3"}
+lazy_static = "1.4.0"
 
 # Note: the BEGIN and END comments below are required for external tooling. Do not remove.
 # BEGIN MOVE DEPENDENCIES

crates/storage/Cargo.toml

@@ -28,6 +28,8 @@ smallvec = { workspace = true }
 triomphe = { workspace = true }
 primitive-types = { workspace = true, features = ["impl-serde"] }
 pretty_assertions = { workspace = true }
+dashmap = { workspace = true}
+lazy_static = { workspace = true }
 
 [dev-dependencies]
 move-vm-test-utils = { workspace = true }

crates/storage/src/allocator.rs

@@ -1,23 +1,16 @@
 use std::{
-    alloc::{GlobalAlloc, Layout, System}, sync::Arc, thread::ThreadId
+    alloc::{GlobalAlloc, Layout, System}, cell::Cell
 };
 
-static mut SIZE_COUNTER: usize = 0;
-static mut REQUEST_THREAD_ID: Option<ThreadId> = None;
+thread_local! {
+    static SIZE: Cell<usize> = Cell::new(0);
+}
 
 struct SizeCounterAllocator;
 
 unsafe impl GlobalAlloc for SizeCounterAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-
-        unsafe {
-            if let Some(tid) = REQUEST_THREAD_ID {
-                if tid == std::thread::current().id() {
-                    SIZE_COUNTER += layout.size();
-                }
-            }
-        }
-
+        SIZE.with(|size| size.set(size.get() + layout.size()));
         System.alloc(layout)
     }
 
@@ -29,20 +22,10 @@ unsafe impl GlobalAlloc for SizeCounterAllocator {
 #[global_allocator]
 static GLOBAL: SizeCounterAllocator = SizeCounterAllocator;
 
-static SIZE_COUNTER_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
-pub(crate) fn get_size_of<T: Clone>(t: Arc<T>) -> usize {
-    let lock = SIZE_COUNTER_LOCK.lock().unwrap();
-
-    let msize: usize;
-    unsafe {
-        REQUEST_THREAD_ID = Some(std::thread::current().id());
-        let _ = (*t).clone();
-        msize = SIZE_COUNTER + size_of::<T>();
-        SIZE_COUNTER = 0;
-        REQUEST_THREAD_ID = None;
-    }
-
-    drop(lock);
-
-    msize
+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

@@ -7,52 +7,54 @@ use move_vm_runtime::Module;
 use move_vm_runtime::Script;
 use move_vm_types::code::{Code, ModuleCode};
 
-use crate::allocator::get_size_of;
 use crate::module_cache::BytesWithHash;
 use crate::module_cache::NoVersion;
 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 {
-        match value {
-            Code::Deserialized(compiled_script) => get_size_of(compiled_script.clone()),
-            Code::Verified(script) => get_size_of(script.clone()),
-        }
+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
+impl WeightScale<Checksum, ModuleCodeWrapper> for ModuleCodeScale
 {
     fn weight(
         &self,
         _key: &Checksum,
-        value: &Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
+        value: &ModuleCodeWrapper,
     ) -> usize {
-        match value.code() {
-            Code::Deserialized(compiled_module) => {
-                get_size_of(compiled_module.clone())
-            },
-            Code::Verified(module) => {
-                get_size_of(module.clone())
-            },
-        }
+        value.size
     }
 }
 
-#[cfg(test)]
-mod test {
-    use std::sync::Arc;
-    use move_binary_format::file_format::basic_test_module;
+#[derive(Clone)]
+pub struct CodeWrapper {
+    pub code: Code<CompiledScript, Script>,
+    pub size: usize,
+}
 
-    #[test]
-    fn test_get_size_of_compiled_module() {
-        let module = basic_test_module();
-        let size = crate::code_scale::get_size_of(Arc::new(module));
-        assert!(size > 0);
+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

@@ -11,17 +11,14 @@ use move_vm_runtime::{
     Script, WithRuntimeEnvironment,
 };
 use move_vm_types::{
-    code::{Code, ModuleBytesStorage},
+    code::ModuleBytesStorage,
     module_linker_error,
     sha3_256
 };
 use std::sync::Arc;
 
 use crate::{
-    module_cache::InitiaModuleCache,
-    module_storage::{AsInitiaModuleStorage, InitiaModuleStorage},
-    script_cache::InitiaScriptCache,
-    state_view::ChecksumStorage,
+    allocator::{initialize_size, get_size}, module_cache::InitiaModuleCache, module_storage::{AsInitiaModuleStorage, InitiaModuleStorage}, script_cache::InitiaScriptCache, state_view::ChecksumStorage
 };
 
 /// Code storage that stores both modules and scripts (not thread-safe).
@@ -99,28 +96,37 @@ 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)?;
+                let allocated_size = get_size();
                 self.script_cache
-                    .insert_deserialized_script(hash, deserialized_script)?
+                    .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
+        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)?,
+                .map(Arc::new)?;
+                let allocated_size = get_size();
+                (compiled_script, allocated_size)
+            }
         };
 
         // Locally verify the script.
@@ -138,12 +144,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)
     }
 }
 
@@ -163,11 +172,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())
         }
     }
 }