MystenLabs · lxfind · Nov 10, 2023 · Nov 10, 2023 · mystenmark · Nov 10, 2023
diff --git a/crates/sui-core/src/authority.rs b/crates/sui-core/src/authority.rs
@@ -4601,7 +4601,7 @@ impl NodeStateDump {
             .runtime_packages_loaded_from_db
             .values()
         {
-            runtime_reads.push(obj.clone());
+            runtime_reads.push(obj.object().clone());
         }
 
         // All other input objects should already be in `inner_temporary_store.objects`

diff --git a/crates/sui-single-node-benchmark/src/mock_storage.rs b/crates/sui-single-node-benchmark/src/mock_storage.rs
@@ -8,28 +8,31 @@ use once_cell::unsync::OnceCell;
 use prometheus::core::{Atomic, AtomicU64};
 use std::collections::HashMap;
 use std::sync::Arc;
+use sui_storage::package_object_cache::PackageObjectCache;
 use sui_types::base_types::{
     EpochId, ObjectID, ObjectRef, SequenceNumber, TransactionDigest, VersionNumber,
 };
 use sui_types::error::{SuiError, SuiResult};
 use sui_types::object::{Object, Owner};
 use sui_types::storage::{
-    get_module_by_id, load_package_object_from_object_store, BackingPackageStore,
-    ChildObjectResolver, GetSharedLocks, ObjectStore, PackageObjectArc, ParentSync,
+    get_module_by_id, BackingPackageStore, ChildObjectResolver, GetSharedLocks, ObjectStore,
+    PackageObjectArc, ParentSync,
 };
 use sui_types::transaction::{InputObjectKind, InputObjects, ObjectReadResult};
 
 // TODO: We won't need a special purpose InMemoryObjectStore once the InMemoryCache is ready.
 #[derive(Clone)]
 pub(crate) struct InMemoryObjectStore {
     objects: Arc<HashMap<ObjectID, Object>>,
+    package_cache: Arc<PackageObjectCache>,
     num_object_reads: Arc<AtomicU64>,
 }
 
 impl InMemoryObjectStore {
     pub(crate) fn new(objects: HashMap<ObjectID, Object>) -> Self {
         Self {
             objects: Arc::new(objects),
+            package_cache: PackageObjectCache::new(),
             num_object_reads: Arc::new(AtomicU64::new(0)),
         }
     }
@@ -134,7 +137,7 @@ impl ObjectStore for InMemoryObjectStore {
 
 impl BackingPackageStore for InMemoryObjectStore {
     fn get_package_object(&self, package_id: &ObjectID) -> SuiResult<Option<PackageObjectArc>> {
-        load_package_object_from_object_store(self, package_id)
+        self.package_cache.get_package_object(package_id, self)
     }
 }
 

diff --git a/crates/sui-types/src/inner_temporary_store.rs b/crates/sui-types/src/inner_temporary_store.rs
@@ -5,6 +5,7 @@ use crate::base_types::{SequenceNumber, VersionDigest};
 use crate::effects::{TransactionEffects, TransactionEffectsAPI, TransactionEvents};
 use crate::execution::DynamicallyLoadedObjectMetadata;
 use crate::storage::InputKey;
+use crate::storage::PackageObjectArc;
 use crate::{
     base_types::ObjectID,
     object::{Object, Owner},
@@ -20,7 +21,7 @@ pub type WrittenObjects = BTreeMap<ObjectID, Object>;
 pub type ObjectMap = BTreeMap<ObjectID, Arc<Object>>;
 pub type TxCoins = (ObjectMap, WrittenObjects);
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone)]
 pub struct InnerTemporaryStore {
     pub input_objects: ObjectMap,
     pub mutable_inputs: BTreeMap<ObjectID, (VersionDigest, Owner)>,
@@ -30,7 +31,7 @@ pub struct InnerTemporaryStore {
     pub events: TransactionEvents,
     pub max_binary_format_version: u32,
     pub no_extraneous_module_bytes: bool,
-    pub runtime_packages_loaded_from_db: BTreeMap<ObjectID, Object>,
+    pub runtime_packages_loaded_from_db: BTreeMap<ObjectID, PackageObjectArc>,
     pub lamport_version: SequenceNumber,
 }
 

diff --git a/crates/sui-types/src/storage/mod.rs b/crates/sui-types/src/storage/mod.rs
@@ -192,7 +192,7 @@ pub trait Storage {
 
 pub type PackageFetchResults<Package> = Result<Vec<Package>, Vec<ObjectID>>;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct PackageObjectArc {
     package_object: Arc<Object>,
 }

diff --git a/sui-execution/latest/sui-adapter/src/temporary_store.rs b/sui-execution/latest/sui-adapter/src/temporary_store.rs
@@ -54,7 +54,7 @@ pub struct TemporaryStore<'backing> {
 
     /// Every package that was loaded from DB store during execution.
     /// These packages were not previously loaded into the temporary store.
-    runtime_packages_loaded_from_db: RwLock<BTreeMap<ObjectID, Object>>,
+    runtime_packages_loaded_from_db: RwLock<BTreeMap<ObjectID, PackageObjectArc>>,
 
     /// The set of objects that we may receive during execution. Not guaranteed to receive all, or
     /// any of the objects referenced in this set.
@@ -1107,16 +1107,25 @@ impl<'backing> BackingPackageStore for TemporaryStore<'backing> {
         // look into the types of each written object in the output, and some of them need the
         // newly written packages for type checking.
         // In production path though, this should never happen.
-        if let Some(obj) = self.read_object(package_id) {
+        if let Some(obj) = self.execution_results.written_objects.get(package_id) {
             Ok(Some(PackageObjectArc::new(obj.clone())))
         } else {
             self.store.get_package_object(package_id).map(|obj| {
                 // Track object but leave unchanged
                 if let Some(v) = &obj {
-                    // TODO: Can this lock ever block execution?
-                    self.runtime_packages_loaded_from_db
-                        .write()
-                        .insert(*package_id, v.object().clone());
+                    if !self
+                        .runtime_packages_loaded_from_db
+                        .read()
+                        .contains_key(package_id)
+                    {
+                        // TODO: Can this lock ever block execution?
+                        // TODO: Another way to avoid the cost of maintaining this map is to not
+                        // enable it in normal runs, and if a fork is detected, rerun it with a flag
+                        // turned on and start populating this field.
+                        self.runtime_packages_loaded_from_db
+                            .write()
+                            .insert(*package_id, v.clone());
+                    }
                 }
                 obj
             })

diff --git a/sui-execution/v0/sui-adapter/src/temporary_store.rs b/sui-execution/v0/sui-adapter/src/temporary_store.rs
@@ -63,7 +63,7 @@ pub struct TemporaryStore<'backing> {
 
     /// Every package that was loaded from DB store during execution.
     /// These packages were not previously loaded into the temporary store.
-    runtime_packages_loaded_from_db: RwLock<BTreeMap<ObjectID, Object>>,
+    runtime_packages_loaded_from_db: RwLock<BTreeMap<ObjectID, PackageObjectArc>>,
 }
 
 impl<'backing> TemporaryStore<'backing> {
@@ -1003,10 +1003,17 @@ impl<'backing> BackingPackageStore for TemporaryStore<'backing> {
             self.store.get_package_object(package_id).map(|obj| {
                 // Track object but leave unchanged
                 if let Some(v) = &obj {
-                    // TODO: Can this lock ever block execution?
-                    self.runtime_packages_loaded_from_db
-                        .write()
-                        .insert(*package_id, v.object().clone());
+                    if !self
+                        .runtime_packages_loaded_from_db
+                        .read()
+                        .contains_key(package_id)
+                    {
+                        // TODO: Can this lock ever block execution?
+                        // TODO: Why do we need a RwLock anyway???
+                        self.runtime_packages_loaded_from_db
+                            .write()
+                            .insert(*package_id, v.clone());
+                    }
                 }
                 obj
             })

diff --git a/sui-execution/v1/sui-adapter/src/temporary_store.rs b/sui-execution/v1/sui-adapter/src/temporary_store.rs
@@ -54,7 +54,7 @@ pub struct TemporaryStore<'backing> {
 
     /// Every package that was loaded from DB store during execution.
     /// These packages were not previously loaded into the temporary store.
-    runtime_packages_loaded_from_db: RwLock<BTreeMap<ObjectID, Object>>,
+    runtime_packages_loaded_from_db: RwLock<BTreeMap<ObjectID, PackageObjectArc>>,
 
     /// The set of objects that we may receive during execution. Not guaranteed to receive all, or
     /// any of the objects referenced in this set.
@@ -1113,10 +1113,17 @@ impl<'backing> BackingPackageStore for TemporaryStore<'backing> {
             self.store.get_package_object(package_id).map(|obj| {
                 // Track object but leave unchanged
                 if let Some(v) = &obj {
-                    // TODO: Can this lock ever block execution?
-                    self.runtime_packages_loaded_from_db
-                        .write()
-                        .insert(*package_id, v.object().clone());
+                    if !self
+                        .runtime_packages_loaded_from_db
+                        .read()
+                        .contains_key(package_id)
+                    {
+                        // TODO: Can this lock ever block execution?
+                        // TODO: Why do we need a RwLock anyway???
+                        self.runtime_packages_loaded_from_db
+                            .write()
+                            .insert(*package_id, v.clone());
+                    }
                 }
                 obj
             })