diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index d36b90df..f251015e 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -22,3 +22,6 @@ clru = { workspace = true } parking_lot = { workspace = true } sha3 = { workspace = true } claims = { workspace = true } + +[dev-dependencies] +rand = { workspace = true } diff --git a/crates/storage/src/allocator.rs b/crates/storage/src/allocator.rs index 7729f087..230dd39b 100644 --- a/crates/storage/src/allocator.rs +++ b/crates/storage/src/allocator.rs @@ -3,7 +3,10 @@ use std::{ cell::Cell, }; +use move_binary_format::errors::VMResult; + thread_local! { + static METERING: Cell = const { Cell::new(false) }; static SIZE: Cell = const { Cell::new(0) }; } @@ -11,7 +14,10 @@ struct SizeCounterAllocator; unsafe impl GlobalAlloc for SizeCounterAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - SIZE.with(|size| size.set(size.get() + layout.size())); + if METERING.with(|metering| metering.get()) { + SIZE.with(|size| size.set(size.get() + layout.size())); + } + System.alloc(layout) } @@ -23,10 +29,54 @@ unsafe impl GlobalAlloc for SizeCounterAllocator { #[global_allocator] static GLOBAL: SizeCounterAllocator = SizeCounterAllocator; -pub(crate) fn initialize_size() { +#[inline] +fn start_metering() { SIZE.with(|size| size.set(0)); + METERING.with(|metering| metering.set(true)); } -pub(crate) fn get_size() -> usize { +#[inline] +fn finish_metering() -> usize { + METERING.with(|metering| metering.set(false)); SIZE.with(|size| size.get()) } + +#[inline] +pub(crate) fn get_size VMResult>(f: O) -> VMResult<(T, usize)> { + start_metering(); + let ret = f()?; + let size = finish_metering(); + + Ok((ret, size + size_of::())) +} + +#[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(); + } + } +} diff --git a/crates/storage/src/code_scale.rs b/crates/storage/src/code_scale.rs index 7058d477..eda0b95b 100644 --- a/crates/storage/src/code_scale.rs +++ b/crates/storage/src/code_scale.rs @@ -11,41 +11,41 @@ use crate::module_cache::BytesWithHash; use crate::module_cache::NoVersion; use crate::state_view::Checksum; -pub struct CodeScale; +pub struct ScriptScale; -impl WeightScale for CodeScale { - fn weight(&self, _key: &Checksum, value: &CodeWrapper) -> usize { +impl WeightScale for ScriptScale { + fn weight(&self, _key: &Checksum, value: &ScriptWrapper) -> usize { value.size } } -pub struct ModuleCodeScale; +pub struct ModuleScale; -impl WeightScale for ModuleCodeScale { - fn weight(&self, _key: &Checksum, value: &ModuleCodeWrapper) -> usize { +impl WeightScale for ModuleScale { + fn weight(&self, _key: &Checksum, value: &ModuleWrapper) -> usize { value.size } } #[derive(Clone)] -pub struct CodeWrapper { +pub struct ScriptWrapper { pub code: Code, pub size: usize, } -impl CodeWrapper { +impl ScriptWrapper { pub fn new(code: Code, size: usize) -> Self { Self { code, size } } } #[derive(Clone)] -pub struct ModuleCodeWrapper { +pub struct ModuleWrapper { pub module_code: Arc>, pub size: usize, } -impl ModuleCodeWrapper { +impl ModuleWrapper { pub fn new( module_code: Arc>, size: usize, diff --git a/crates/storage/src/code_storage.rs b/crates/storage/src/code_storage.rs index 148e2b74..36c8588f 100644 --- a/crates/storage/src/code_storage.rs +++ b/crates/storage/src/code_storage.rs @@ -14,7 +14,7 @@ use move_vm_types::{code::ModuleBytesStorage, module_linker_error, sha3_256}; use std::sync::Arc; use crate::{ - allocator::{get_size, initialize_size}, + allocator::get_size, module_cache::InitiaModuleCache, module_storage::{AsInitiaModuleStorage, InitiaModuleStorage}, script_cache::InitiaScriptCache, @@ -98,11 +98,13 @@ impl CodeStorage for InitiaCodeStorage { Ok(match self.script_cache.get_script(&hash) { 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(); + // 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, @@ -120,16 +122,15 @@ impl CodeStorage for InitiaCodeStorage { 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) + get_size(move || { + self.runtime_environment() + .deserialize_into_script(serialized_script) + .map(Arc::new) + })? } }; @@ -149,14 +150,18 @@ impl CodeStorage for InitiaCodeStorage { }) .collect::>>()?; - 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, allocated_size) + // 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, + ) } } diff --git a/crates/storage/src/module_cache.rs b/crates/storage/src/module_cache.rs index 8cae2343..3efdf5f1 100644 --- a/crates/storage/src/module_cache.rs +++ b/crates/storage/src/module_cache.rs @@ -12,8 +12,8 @@ use move_vm_types::code::{ModuleCode, ModuleCodeBuilder, WithBytes, WithHash}; use parking_lot::Mutex; use crate::{ - allocator::{get_size, initialize_size}, - code_scale::{ModuleCodeScale, ModuleCodeWrapper}, + allocator::get_size, + code_scale::{ModuleScale, ModuleWrapper}, state_view::Checksum, }; @@ -58,8 +58,7 @@ pub struct NoVersion; pub struct InitiaModuleCache { #[allow(clippy::type_complexity)] - pub(crate) module_cache: - Mutex>, + pub(crate) module_cache: Mutex>, } impl InitiaModuleCache { @@ -67,7 +66,7 @@ impl InitiaModuleCache { 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), + CLruCacheConfig::new(capacity).with_scale(ModuleScale), )), }) } @@ -97,7 +96,7 @@ impl InitiaModuleCache { version, )); module_cache - .put_with_weight(key, ModuleCodeWrapper::new(module, allocated_size)) + .put_with_weight(key, ModuleWrapper::new(module, allocated_size)) .map_err(|_| handle_cache_error(module_id))?; Ok(()) } @@ -113,29 +112,15 @@ impl InitiaModuleCache { version: NoVersion, ) -> VMResult>> { let mut module_cache = self.module_cache.lock(); - match module_cache.get(&key) { - 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, - ModuleCodeWrapper::new(module.clone(), allocated_size), - ) - .map_err(|_| handle_cache_error(module_id))?; - Ok(module) - } + Some(module_wrapper) if module_wrapper.module_code.code().is_verified() => { + Ok(module_wrapper.module_code.clone()) } - None => { + _ => { let module_id = verified_code.self_id(); let module = Arc::new(ModuleCode::from_verified(verified_code, extension, version)); module_cache - .put_with_weight(key, ModuleCodeWrapper::new(module.clone(), allocated_size)) + .put_with_weight(key, ModuleWrapper::new(module.clone(), allocated_size)) .map_err(|_| handle_cache_error(module_id))?; Ok(module) } @@ -154,14 +139,12 @@ impl InitiaModuleCache { Extension = BytesWithHash, Version = NoVersion, >, - ) -> VMResult> { + ) -> VMResult> { let mut module_cache = self.module_cache.lock(); Ok(match module_cache.get(checksum) { - Some(code) => Some(ModuleCodeWrapper::new(code.module_code.clone(), code.size)), + Some(module_wrapper) => Some(module_wrapper.clone()), None => { - initialize_size(); - let build_result = builder.build(id)?; - let allocated_size = get_size(); + let (build_result, allocated_size) = get_size(move || builder.build(id))?; match build_result { Some(code) => { if code.extension().hash() != checksum { @@ -172,8 +155,7 @@ impl InitiaModuleCache { .finish(Location::Module(id.clone()))); } - let code = Arc::new(code); - let code_wrapper = ModuleCodeWrapper::new(code.clone(), allocated_size); + let code_wrapper = ModuleWrapper::new(Arc::new(code), allocated_size); module_cache .put_with_weight(*checksum, code_wrapper.clone()) .map_err(|_| { diff --git a/crates/storage/src/module_storage.rs b/crates/storage/src/module_storage.rs index 4d68fc03..10758f54 100644 --- a/crates/storage/src/module_storage.rs +++ b/crates/storage/src/module_storage.rs @@ -2,16 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 use crate::{ - allocator::{get_size, initialize_size}, + allocator::get_size, + code_scale::ModuleWrapper, module_cache::{BytesWithHash, InitiaModuleCache, NoVersion}, state_view::{Checksum, ChecksumStorage}, }; use bytes::Bytes; -use move_binary_format::{errors::VMResult, CompiledModule}; +use move_binary_format::{ + errors::{Location, PartialVMError, VMResult}, + CompiledModule, +}; use move_core_types::{ account_address::AccountAddress, identifier::IdentStr, language_storage::ModuleId, - metadata::Metadata, + metadata::Metadata, vm_status::StatusCode, }; use move_vm_runtime::{Module, ModuleStorage, RuntimeEnvironment, WithRuntimeEnvironment}; use move_vm_types::{ @@ -169,7 +173,11 @@ impl<'s, S: ModuleBytesStorage + ChecksumStorage> ModuleCodeBuilder for InitiaMo .deserialize_into_compiled_module(&bytes)?; if checksum != hash { - return Err(module_linker_error!(key.address(), key.name())); + return Err( + PartialVMError::new(StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR) + .with_message("Checksum mismatch".to_string()) + .finish(Location::Module(compiled_module.self_id())), + ); } let extension = Arc::new(BytesWithHash::new(bytes, hash)); @@ -291,8 +299,7 @@ impl<'a, S: ModuleBytesStorage + ChecksumStorage> ModuleStorage for InitiaModule Ok(Some(visit_dependencies_and_verify( id, checksum, - module_code_wrapper.module_code, - module_code_wrapper.size, + module_code_wrapper, &mut visited, self, )?)) @@ -315,12 +322,12 @@ impl<'a, S: ModuleBytesStorage + ChecksumStorage> ModuleStorage for InitiaModule fn visit_dependencies_and_verify( module_id: ModuleId, checksum: Checksum, - module: Arc>, - allocated_size: usize, + unverified_module_wrapper: ModuleWrapper, visited: &mut HashSet, module_cache_with_context: &InitiaModuleStorage<'_, S>, ) -> VMResult> { let runtime_environment = module_cache_with_context.runtime_environment; + let module = unverified_module_wrapper.module_code; // Step 1: Local verification. runtime_environment.paranoid_check_module_address_and_name( @@ -364,8 +371,7 @@ fn visit_dependencies_and_verify( let verified_dependency = visit_dependencies_and_verify( dependency_id.clone(), dependency_checksum, - dependency.module_code, - dependency.size, + dependency, visited, module_cache_with_context, )?; @@ -387,17 +393,18 @@ fn visit_dependencies_and_verify( }; } - initialize_size(); - let verified_code = - runtime_environment.build_verified_module(locally_verified_code, &verified_dependencies)?; - let module_allocated_size = get_size() + allocated_size; + // Build verified module and the compute size of the verified module + let (verified_code, allocated_size_for_verified) = get_size(move || { + runtime_environment.build_verified_module(locally_verified_code, &verified_dependencies) + })?; + // Cache the verified module let module = module_cache_with_context .module_cache .insert_verified_module( checksum, verified_code, - module_allocated_size, + allocated_size_for_verified + unverified_module_wrapper.size, module.extension().clone(), module.version(), )?; diff --git a/crates/storage/src/script_cache.rs b/crates/storage/src/script_cache.rs index 3003199d..b7ee58c7 100644 --- a/crates/storage/src/script_cache.rs +++ b/crates/storage/src/script_cache.rs @@ -11,12 +11,12 @@ use move_vm_types::code::Code; use parking_lot::Mutex; use crate::{ - code_scale::{CodeScale, CodeWrapper}, + code_scale::{ScriptScale, ScriptWrapper}, state_view::Checksum, }; pub struct InitiaScriptCache { - pub(crate) script_cache: Mutex>, + pub(crate) script_cache: Mutex>, } impl InitiaScriptCache { @@ -24,7 +24,7 @@ impl InitiaScriptCache { Arc::new(InitiaScriptCache { script_cache: Mutex::new(CLruCache::with_config( CLruCacheConfig::new(NonZeroUsize::new(cache_capacity * 1024 * 1024).unwrap()) - .with_scale(CodeScale), + .with_scale(ScriptScale), )), }) } @@ -48,7 +48,7 @@ impl InitiaScriptCache { // error occurs when the new script has a weight greater than the cache capacity script_cache - .put_with_weight(key, CodeWrapper::new(new_script, allocated_size)) + .put_with_weight(key, ScriptWrapper::new(new_script, allocated_size)) .map_err(|_| { PartialVMError::new(StatusCode::MEMORY_LIMIT_EXCEEDED) .with_message("Script storage cache eviction error".to_string()) @@ -69,13 +69,13 @@ impl InitiaScriptCache { let mut script_cache = self.script_cache.lock(); let (new_script, verified_script) = match script_cache.get(&key) { - Some(code) => { - if !code.code.is_verified() { + Some(script_wrapper) => { + if !script_wrapper.code.is_verified() { let new_script = Code::from_verified(verified_script); let verified_script = new_script.verified().clone(); (Some(new_script), verified_script) } else { - (None, code.code.verified().clone()) + (None, script_wrapper.code.verified().clone()) } } None => { @@ -87,7 +87,7 @@ impl InitiaScriptCache { if new_script.is_some() { script_cache - .put_with_weight(key, CodeWrapper::new(new_script.unwrap(), allocated_size)) + .put_with_weight(key, ScriptWrapper::new(new_script.unwrap(), allocated_size)) .map_err(|_| { PartialVMError::new(StatusCode::MEMORY_LIMIT_EXCEEDED) .with_message("Script storage cache eviction error".to_string()) @@ -97,16 +97,12 @@ impl InitiaScriptCache { Ok(verified_script) } - pub(crate) fn get_script(&self, key: &Checksum) -> Option { - let mut script_cache = self.script_cache.lock(); - script_cache - .get(key) - .map(|k| CodeWrapper::new(k.code.clone(), k.size)) + pub(crate) fn get_script(&self, key: &Checksum) -> Option { + self.script_cache.lock().get(key).cloned() } #[allow(unused)] pub(crate) fn num_scripts(&self) -> usize { - let script_cache = self.script_cache.lock(); - script_cache.len() + self.script_cache.lock().len() } }