Skip to content

Commit

Permalink
refactory PR161 (#163)
Browse files Browse the repository at this point in the history
  • Loading branch information
@beer-1
beer-1 authored Nov 1, 2024
1 parent 56f0fdb commit f8dd0ef
Show file tree
Hide file tree
Showing 7 changed files with 138 additions and 95 deletions.
3 changes: 3 additions & 0 deletions crates/storage/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,6 @@ clru = { workspace = true }
parking_lot = { workspace = true }
sha3 = { workspace = true }
claims = { workspace = true }

[dev-dependencies]
rand = { workspace = true }
56 changes: 53 additions & 3 deletions crates/storage/src/allocator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,21 @@ use std::{
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 {
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)
}

Expand All @@ -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<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();
}
}
}
20 changes: 10 additions & 10 deletions crates/storage/src/code_scale.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Checksum, CodeWrapper> for CodeScale {
fn weight(&self, _key: &Checksum, value: &CodeWrapper) -> usize {
impl WeightScale<Checksum, ScriptWrapper> for ScriptScale {
fn weight(&self, _key: &Checksum, value: &ScriptWrapper) -> usize {
value.size
}
}

pub struct ModuleCodeScale;
pub struct ModuleScale;

impl WeightScale<Checksum, ModuleCodeWrapper> for ModuleCodeScale {
fn weight(&self, _key: &Checksum, value: &ModuleCodeWrapper) -> usize {
impl WeightScale<Checksum, ModuleWrapper> for ModuleScale {
fn weight(&self, _key: &Checksum, value: &ModuleWrapper) -> usize {
value.size
}
}

#[derive(Clone)]
pub struct CodeWrapper {
pub struct ScriptWrapper {
pub code: Code<CompiledScript, Script>,
pub size: usize,
}

impl CodeWrapper {
impl ScriptWrapper {
pub fn new(code: Code<CompiledScript, Script>, size: usize) -> Self {
Self { code, size }
}
}

#[derive(Clone)]
pub struct ModuleCodeWrapper {
pub struct ModuleWrapper {
pub module_code: Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
pub size: usize,
}

impl ModuleCodeWrapper {
impl ModuleWrapper {
pub fn new(
module_code: Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>,
size: usize,
Expand Down
47 changes: 26 additions & 21 deletions crates/storage/src/code_storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -98,11 +98,13 @@ impl<M: ModuleStorage> CodeStorage for InitiaCodeStorage<M> {
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,
Expand All @@ -120,16 +122,15 @@ impl<M: ModuleStorage> CodeStorage for InitiaCodeStorage<M> {
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)
})?
}
};

Expand All @@ -149,14 +150,18 @@ impl<M: ModuleStorage> CodeStorage for InitiaCodeStorage<M> {
})
.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, 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,
)
}
}

Expand Down
44 changes: 13 additions & 31 deletions crates/storage/src/module_cache.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -58,16 +58,15 @@ pub struct NoVersion;

pub struct InitiaModuleCache {
#[allow(clippy::type_complexity)]
pub(crate) module_cache:
Mutex<CLruCache<Checksum, ModuleCodeWrapper, RandomState, ModuleCodeScale>>,
pub(crate) module_cache: Mutex<CLruCache<Checksum, ModuleWrapper, RandomState, ModuleScale>>,
}

impl InitiaModuleCache {
pub fn new(cache_capacity: usize) -> Arc<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),
)),
})
}
Expand Down Expand Up @@ -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(())
}
Expand All @@ -113,29 +112,15 @@ impl InitiaModuleCache {
version: NoVersion,
) -> VMResult<Arc<ModuleCode<CompiledModule, Module, BytesWithHash, NoVersion>>> {
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)
}
Expand All @@ -154,14 +139,12 @@ impl InitiaModuleCache {
Extension = BytesWithHash,
Version = NoVersion,
>,
) -> VMResult<Option<ModuleCodeWrapper>> {
) -> VMResult<Option<ModuleWrapper>> {
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 {
Expand All @@ -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(|_| {
Expand Down
Loading

0 comments on commit f8dd0ef

Please sign in to comment.