[Authority + Adapter] Add package caching, and remove redundant diges…

…t checks + clones (MystenLabs#963)

* Added cache and keep digest checks

Co-authored-by: George Danezis <[email protected]>

sui/src/rest_server.rs

@@ -1193,7 +1193,7 @@ async fn handle_move_call(
 
     // Pass in the objects for a deeper check
     resolve_and_type_check(
-        package_object.clone(),
+        &package_object,
         &module,
         &function,
         &type_args,

sui/src/wallet_commands.rs

@@ -262,7 +262,7 @@ impl WalletCommands {
                 // Pass in the objects for a deeper check
                 // We can technically move this to impl MovePackage
                 resolve_and_type_check(
-                    package_obj.clone(),
+                    package_obj,
                     module,
                     function,
                     type_args,

sui_core/src/authority.rs

@@ -19,7 +19,7 @@ use std::{
     pin::Pin,
     sync::Arc,
 };
-use sui_adapter::adapter;
+use sui_adapter::adapter::{self, SuiMoveVM};
 use sui_types::{
     base_types::*,
     batch::UpdateItem,
@@ -77,7 +77,7 @@ pub struct AuthorityState {
 
     /// Move native functions that are available to invoke
     _native_functions: NativeFunctionTable,
-    move_vm: Arc<adapter::MoveVM>,
+    move_vm: Arc<adapter::SuiMoveVM>,
 
     /// The database
     pub(crate) _database: Arc<AuthorityStore>, // TODO: remove pub
@@ -284,7 +284,7 @@ impl AuthorityState {
             transaction_digest,
             objects_by_kind,
             &self.move_vm,
-            self._native_functions.clone(),
+            &self._native_functions,
         )?;
         let signed_effects = effects.to_sign_effects(&self.name, &*self.secret);
 
@@ -607,6 +607,8 @@ impl AuthorityState {
         let package_id = ObjectID::from(*modules[0].self_id().address());
         let natives = self._native_functions.clone();
         let vm = adapter::verify_and_link(&temporary_store, &modules, package_id, natives)?;
+        let vm = SuiMoveVM::new(vm);
+
         if let ExecutionStatus::Failure { error, .. } = adapter::store_package_and_init_modules(
             &mut temporary_store,
             &vm,

sui_core/src/execution_engine.rs

@@ -23,8 +23,8 @@ pub fn execute_transaction_to_effects<S: BackingPackageStore>(
     transaction: Transaction,
     transaction_digest: TransactionDigest,
     objects_by_kind: Vec<(InputObjectKind, Object)>,
-    move_vm: &Arc<adapter::MoveVM>,
-    native_functions: NativeFunctionTable,
+    move_vm: &Arc<adapter::SuiMoveVM>,
+    native_functions: &NativeFunctionTable,
 ) -> SuiResult<TransactionEffects> {
     let mut transaction_dependencies: BTreeSet<_> = objects_by_kind
         .iter()
@@ -64,8 +64,8 @@ fn execute_transaction<S: BackingPackageStore>(
     transaction: Transaction,
     mut objects_by_kind: Vec<(InputObjectKind, Object)>,
     tx_ctx: &mut TxContext,
-    move_vm: &Arc<adapter::MoveVM>,
-    native_functions: NativeFunctionTable,
+    move_vm: &Arc<adapter::SuiMoveVM>,
+    native_functions: &NativeFunctionTable,
 ) -> SuiResult<ExecutionStatus> {
     // unwraps here are safe because we built `inputs`
     let mut gas_object = objects_by_kind.pop().unwrap().1;
@@ -90,8 +90,8 @@ fn execute_transaction<S: BackingPackageStore>(
                 adapter::execute(
                     move_vm,
                     temporary_store,
-                    native_functions.clone(),
-                    package,
+                    native_functions,
+                    &package,
                     &c.module,
                     &c.function,
                     c.type_arguments.clone(),

sui_programmability/adapter/Cargo.toml

@@ -12,6 +12,7 @@ anyhow = "1.0.55"
 bcs = "0.1.3"
 once_cell = "1.9.0"
 structopt = "0.3.26"
+parking_lot = "0.12.0"
 
 move-binary-format = { git = "https://github.com/diem/move", rev = "2ef516919d5bbc728a57a2c0073b85c46d9fcf5a" }
 move-bytecode-utils = { git = "https://github.com/diem/move", rev = "2ef516919d5bbc728a57a2c0073b85c46d9fcf5a" }

sui_programmability/adapter/src/adapter.rs

@@ -50,10 +50,46 @@ macro_rules! exec_failure {
 #[path = "unit_tests/adapter_tests.rs"]
 mod adapter_tests;
 
-pub fn new_move_vm(natives: NativeFunctionTable) -> Result<Arc<MoveVM>, SuiError> {
-    Ok(Arc::new(
-        MoveVM::new(natives).map_err(|_| SuiError::ExecutionInvariantViolation)?,
-    ))
+// This structure holds a VM and a cache of packages that contain
+// in turn caches of their Function entry points,
+// that involve re-computation / deserialization to otherwise get.
+pub struct SuiMoveVM {
+    movevm: MoveVM,
+
+    // TODO: change this to an LRU cache, to avoid running out of
+    //       memory.
+    package_cache: parking_lot::RwLock<HashMap<ObjectID, Arc<Object>>>,
+}
+
+impl SuiMoveVM {
+    /// Make a new struct from a VM
+    pub fn new(vm: MoveVM) -> SuiMoveVM {
+        SuiMoveVM {
+            movevm: vm,
+            package_cache: parking_lot::RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Returns an object from the cache if one is available
+    /// and otherwise caches a copy of this object.
+    pub fn get_package(&self, object: &Object) -> Arc<Object> {
+        let id = object.id();
+
+        if let Some(cached_object) = self.package_cache.read().get(&id) {
+            return cached_object.clone();
+        }
+
+        let arc_object = Arc::new(object.clone());
+        self.package_cache.write().insert(id, arc_object.clone());
+        arc_object
+    }
+}
+
+pub fn new_move_vm(natives: NativeFunctionTable) -> Result<Arc<SuiMoveVM>, SuiError> {
+    Ok(Arc::new(SuiMoveVM {
+        movevm: MoveVM::new(natives).map_err(|_| SuiError::ExecutionInvariantViolation)?,
+        package_cache: parking_lot::RwLock::new(HashMap::new()),
+    }))
 }
 
 /// Execute `module::function<type_args>(object_args ++ pure_args)` as a call from `sender` with the given `gas_budget`.
@@ -65,10 +101,10 @@ pub fn new_move_vm(natives: NativeFunctionTable) -> Result<Arc<MoveVM>, SuiError
 /// into ExecutionStatus::Failure.
 #[allow(clippy::too_many_arguments)]
 pub fn execute<E: Debug, S: ResourceResolver<Error = E> + ModuleResolver<Error = E> + Storage>(
-    vm: &MoveVM,
+    vm: &SuiMoveVM,
     state_view: &mut S,
-    _natives: NativeFunctionTable,
-    package_object: Object,
+    _natives: &NativeFunctionTable,
+    package_object: &Object,
     module: &Identifier,
     function: &Identifier,
     type_args: Vec<TypeTag>,
@@ -87,14 +123,16 @@ pub fn execute<E: Debug, S: ResourceResolver<Error = E> + ModuleResolver<Error =
             object_owner_map.insert(obj.id().into(), owner);
         }
     }
+    let cached_package = vm.get_package(package_object);
+
     let TypeCheckSuccess {
         module_id,
         args,
         mutable_ref_objects,
         by_value_objects,
         return_types,
     } = match resolve_and_type_check(
-        package_object,
+        &cached_package,
         module,
         function,
         &type_args,
@@ -133,7 +171,7 @@ fn execute_internal<
     E: Debug,
     S: ResourceResolver<Error = E> + ModuleResolver<Error = E> + Storage,
 >(
-    vm: &MoveVM,
+    vm: &SuiMoveVM,
     state_view: &mut S,
     module_id: &ModuleId,
     function: &Identifier,
@@ -155,7 +193,7 @@ fn execute_internal<
             Ok(ok) => ok,
             Err(err) => return ExecutionStatus::new_failure(gas::MIN_MOVE, err),
         };
-    let session = vm.new_session(state_view);
+    let session = vm.movevm.new_session(state_view);
     match session.execute_function_for_effects(
         module_id,
         function,
@@ -364,6 +402,7 @@ pub fn publish<E: Debug, S: ResourceResolver<Error = E> + ModuleResolver<Error =
         Err(err) => exec_failure!(gas::MIN_MOVE, err),
     };
 
+    let vm = SuiMoveVM::new(vm);
     let gas_used_for_init = match store_package_and_init_modules(
         state_view,
         &vm,
@@ -393,7 +432,7 @@ pub fn store_package_and_init_modules<
     S: ResourceResolver<Error = E> + ModuleResolver<Error = E> + Storage,
 >(
     state_view: &mut S,
-    vm: &MoveVM,
+    vm: &SuiMoveVM,
     modules: Vec<CompiledModule>,
     ctx: &mut TxContext,
     gas_budget: u64,
@@ -417,7 +456,7 @@ pub fn store_package_and_init_modules<
 /// Modules in module_ids_to_init must have the init method defined
 fn init_modules<E: Debug, S: ResourceResolver<Error = E> + ModuleResolver<Error = E> + Storage>(
     state_view: &mut S,
-    vm: &MoveVM,
+    vm: &SuiMoveVM,
     module_ids_to_init: Vec<ModuleId>,
     ctx: &mut TxContext,
     gas_budget: u64,

sui_programmability/adapter/src/unit_tests/adapter_tests.rs

@@ -198,8 +198,8 @@ fn call(
     adapter::execute(
         &vm,
         storage,
-        native_functions.clone(),
-        package,
+        native_functions,
+        &package,
         &Identifier::new(module_name).unwrap(),
         &Identifier::new(fun_name).unwrap(),
         type_args,
@@ -705,8 +705,8 @@ fn test_move_call_incorrect_function() {
     let status = adapter::execute(
         &vm,
         &mut storage,
-        native_functions.clone(),
-        gas_object,
+        &native_functions,
+        &gas_object,
         &Identifier::new("ObjectBasics").unwrap(),
         &Identifier::new("create").unwrap(),
         vec![],

sui_types/Cargo.toml

@@ -25,6 +25,7 @@ serde_json = "1.0.79"
 serde_with = "1.12.0"
 signature = "1.5.0"
 static_assertions = "1.1.0"
+parking_lot = "0.12.0"
 
 name_variant = { git = "https://github.com/MystenLabs/mysten-infra", rev ="97a056f85555fa2afe497d6abb7cf6bf8faa63cf"}
 typed-store = { git = "https://github.com/MystenLabs/mysten-infra", rev ="e44bca4513a6ff6c97399cd79e82e4bc00571ac3"}

sui_types/src/move_package.rs

@@ -19,7 +19,11 @@ use move_core_types::{
 };
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
-use std::{collections::BTreeMap, u32};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+    u32,
+};
 
 // TODO: robust MovePackage tests
 // #[cfg(test)]
@@ -36,11 +40,36 @@ pub struct TypeCheckSuccess {
     pub return_types: Vec<Type>, // to validate return types after the call
 }
 
+/// An inner structure used to cache expensive calls that
+/// otherwise involve repeated deserialization.
+#[derive(Debug, Clone, Default)]
+pub struct PackageCache {
+    pub function_signature_cache:
+        Arc<parking_lot::RwLock<HashMap<(Identifier, Identifier), Function>>>,
+}
+
+impl PartialEq for PackageCache {
+    fn eq(&self, _other: &Self) -> bool {
+        true
+    }
+}
+impl Eq for PackageCache {}
+
+use std::hash::Hash;
+use std::hash::Hasher;
+
+impl Hash for PackageCache {
+    fn hash<H: Hasher>(&self, _state: &mut H) {}
+}
+
 // serde_bytes::ByteBuf is an analog of Vec<u8> with built-in fast serialization.
 #[derive(Eq, PartialEq, Debug, Clone, Deserialize, Serialize, Hash)]
 pub struct MovePackage {
     id: ObjectID,
     module_map: BTreeMap<String, ByteBuf>,
+
+    #[serde(skip)]
+    cache: PackageCache,
 }
 
 impl MovePackage {
@@ -52,6 +81,7 @@ impl MovePackage {
         Self {
             id,
             module_map: module_map.clone(),
+            cache: PackageCache::default(),
         }
     }
 
@@ -60,13 +90,7 @@ impl MovePackage {
     }
 
     pub fn module_id(&self, module: &Identifier) -> Result<ModuleId, SuiError> {
-        let ser = self
-            .serialized_module_map()
-            .get(module.as_str())
-            .ok_or_else(|| SuiError::ModuleNotFound {
-                module_name: module.to_string(),
-            })?;
-        Ok(CompiledModule::deserialize(ser)?.self_id())
+        Ok(ModuleId::new(*self.id, module.clone()))
     }
 
     /// Get the function signature for the specified function
@@ -75,6 +99,15 @@ impl MovePackage {
         module: &Identifier,
         function: &Identifier,
     ) -> Result<Function, SuiError> {
+        if let Some(func) = self
+            .cache
+            .function_signature_cache
+            .read()
+            .get(&(module.clone(), function.clone()))
+        {
+            return Ok(func.clone());
+        }
+
         let bytes = self
             .serialized_module_map()
             .get(module.as_str())
@@ -84,14 +117,22 @@ impl MovePackage {
         let m = CompiledModule::deserialize(bytes)
             .expect("Unwrap safe because Sui serializes/verifies modules before publishing them");
 
-        Function::new_from_name(&m, function).ok_or_else(|| SuiError::FunctionNotFound {
-            error: format!(
-                "Could not resolve function '{}' in module {}::{}",
-                function,
-                self.id(),
-                module
-            ),
-        })
+        let func =
+            Function::new_from_name(&m, function).ok_or_else(|| SuiError::FunctionNotFound {
+                error: format!(
+                    "Could not resolve function '{}' in module {}::{}",
+                    function,
+                    self.id(),
+                    module
+                ),
+            })?;
+
+        self.cache
+            .function_signature_cache
+            .write()
+            .insert((module.clone(), function.clone()), func.clone());
+
+        Ok(func)
     }
 
     /// Checks if the specified function is an entry function and returns the function if so
@@ -161,15 +202,15 @@ impl From<&Vec<CompiledModule>> for MovePackage {
 /// - Return the ID of the resolved module, a vector of BCS encoded arguments to pass to the VM, and a partitioning
 /// of the input objects into objects passed by value vs by mutable reference
 pub fn resolve_and_type_check(
-    package_object: Object,
+    package_object: &Object,
     module: &Identifier,
     function: &Identifier,
     type_args: &[TypeTag],
     object_args: Vec<Object>,
     mut pure_args: Vec<Vec<u8>>,
 ) -> Result<TypeCheckSuccess, SuiError> {
     // Resolve the function we are calling
-    let (function_signature, module_id) = match package_object.data {
+    let (function_signature, module_id) = match &package_object.data {
         Data::Package(package) => (
             package.check_and_get_entry_function(module, function)?,
             package.module_id(module)?,