[comments] Addressing Rati's comments + adding size in bytes to modul…

…e cache

aptos-move/aptos-global-cache-manager/Cargo.toml

@@ -20,6 +20,7 @@ aptos-vm-types = { workspace = true }
 move-binary-format = { workspace = true }
 move-core-types = { workspace = true }
 move-vm-runtime = { workspace = true }
+move-vm-types = { workspace = true }
 parking_lot = { workspace = true }
 
 [dev-dependencies]

aptos-move/aptos-global-cache-manager/src/config.rs

@@ -5,23 +5,23 @@
 pub struct GlobalCacheConfig {
     /// If true, when global caches are empty, Aptos framework is prefetched into module cache.
     pub prefetch_framework_code: bool,
-    /// The maximum number of entries stored in module cache. If module cache exceeds this value,
-    /// all its entries should be flushed.
-    pub module_cache_capacity: usize,
-    /// The maximum size of struct name re-indexing map stored in runtime environment.
-    pub struct_name_index_map_capacity: usize,
+    /// The maximum size serialized modules can take in module cache.
+    pub max_module_cache_size_in_bytes: usize,
+    /// The maximum size (in terms of entries) of struct name re-indexing map stored in runtime
+    /// environment.
+    pub max_struct_name_index_map_size: usize,
 }
 
 impl Default for GlobalCacheConfig {
     fn default() -> Self {
         // TODO(loader_v2):
-        //   Right now these are just some numbers, we can set them based on the upper bounds of
-        //   module or identifier sizes, or keep track in read-only module cache how many bytes we
-        //   are using.
+        //   Right now these are hardcoded here, we probably want to add them to gas schedule or
+        //   some on-chain config.
         Self {
             prefetch_framework_code: true,
-            module_cache_capacity: 10_000,
-            struct_name_index_map_capacity: 10_000,
+            // Use 50 Mb for now, should be large enough to cache many modules.
+            max_module_cache_size_in_bytes: 50 * 1024 * 1024,
+            max_struct_name_index_map_size: 100_000,
         }
     }
 }

aptos-move/aptos-global-cache-manager/src/manager.rs

@@ -17,6 +17,7 @@ use move_core_types::{
     vm_status::{StatusCode, VMStatus},
 };
 use move_vm_runtime::{Module, ModuleStorage, WithRuntimeEnvironment};
+use move_vm_types::code::WithSize;
 use parking_lot::Mutex;
 use std::{
     hash::Hash,
@@ -54,11 +55,11 @@ impl BlockId {
 
 /// Manages global caches, e.g., modules or execution environment. Should not be used concurrently.
 struct GlobalCacheManagerInner<K, DC, VC, E> {
-    /// Different configurations used for handling global caches.
     config: GlobalCacheConfig,
+
     /// Cache for modules. It is read-only for any concurrent execution, and can only be mutated
     /// when it is known that there are no concurrent accesses, e.g., at block boundaries.
-    /// [GlobalCacheManagerInner] tries to ensure that these invariants always hold.
+    /// [GlobalCacheManagerInner] must ensure that these invariants always hold.
     module_cache: Arc<ReadOnlyModuleCache<K, DC, VC, E>>,
     /// Identifies previously executed block, initially [BlockId::Unset].
     previous_block_id: Mutex<BlockId>,
@@ -75,6 +76,7 @@ impl<K, DC, VC, E> GlobalCacheManagerInner<K, DC, VC, E>
 where
     K: Hash + Eq + Clone,
     VC: Deref<Target = Arc<DC>>,
+    E: WithSize,
 {
     /// Returns a new instance of [GlobalCacheManagerInner] with default [GlobalCacheConfig].
     fn new_with_default_config() -> Self {
@@ -174,10 +176,10 @@ where
             Ok(size) => size,
         };
 
-        if struct_name_index_map_size > self.config.struct_name_index_map_capacity {
+        if struct_name_index_map_size > self.config.max_struct_name_index_map_size {
             flush_all_caches = true;
         }
-        if self.module_cache.num_modules() > self.config.module_cache_capacity {
+        if self.module_cache.size_in_bytes() > self.config.max_module_cache_size_in_bytes {
             // Technically, if we flush modules we do not need to flush type caches, but we unify
             // flushing logic for easier reasoning.
             flush_all_caches = true;
@@ -234,7 +236,7 @@ where
     }
 
     /// Resets all states (under a lock) as if global caches are empty and no blocks have been
-    /// executed so far. Reruns an invariant violation error.
+    /// executed so far. Returns an invariant violation error.
     fn reset_and_return_invariant_violation(&self, msg: &str) -> VMStatus {
         // Lock to reset the state under lock.
         let mut previous_block_id = self.previous_block_id.lock();
@@ -278,7 +280,7 @@ impl GlobalCacheManager {
     ///   1. Previously executed block ID does not match the provided value.
     ///   2. The environment has changed for this state.
     ///   3. The size of the struct name re-indexing map is too large.
-    ///   4. The size of the module cache is too large.
+    ///   4. The size (in bytes) of the module cache is too large.
     ///
     /// Additionally, if cache is empty, prefetches the framework code into it.
     ///
@@ -480,11 +482,11 @@ mod test {
     }
 
     #[test]
-    fn mark_execution_start_when_too_many_modules() {
+    fn mark_execution_start_when_module_cache_is_too_large() {
         let state_view = MockStateView::empty();
 
         let config = GlobalCacheConfig {
-            module_cache_capacity: 1,
+            max_module_cache_size_in_bytes: 8,
             ..Default::default()
         };
         let global_cache_manager = GlobalCacheManagerInner::new_with_config(config);
@@ -494,14 +496,16 @@ mod test {
             .insert(0, mock_verified_code(0, MockExtension::new(8)));
         global_cache_manager
             .module_cache
-            .insert(1, mock_verified_code(1, MockExtension::new(8)));
+            .insert(1, mock_verified_code(1, MockExtension::new(24)));
         assert_eq!(global_cache_manager.module_cache.num_modules(), 2);
+        assert_eq!(global_cache_manager.module_cache.size_in_bytes(), 32);
 
         // Cache is too large, should be flushed for next block.
         assert_ok!(
             global_cache_manager.mark_block_execution_start(&state_view, Some(HashValue::random()))
         );
         assert_eq!(global_cache_manager.module_cache.num_modules(), 0);
+        assert_eq!(global_cache_manager.module_cache.size_in_bytes(), 0);
     }
 
     #[test_case(None)]
@@ -565,7 +569,7 @@ mod test {
             u32,
             MockDeserializedCode,
             MockVerifiedCode,
-            (),
+            MockExtension,
         >::new_with_default_config());
         assert!(global_cache_manager.ready_for_next_block());
 
@@ -587,7 +591,7 @@ mod test {
             u32,
             MockDeserializedCode,
             MockVerifiedCode,
-            (),
+            MockExtension,
         >::new_with_default_config();
         assert!(global_cache_manager.previous_block_id.lock().is_unset());
 
@@ -619,7 +623,7 @@ mod test {
             u32,
             MockDeserializedCode,
             MockVerifiedCode,
-            (),
+            MockExtension,
         >::new_with_default_config());
         global_cache_manager.mark_not_ready_for_next_block();
 

aptos-move/block-executor/src/captured_reads.rs

@@ -27,7 +27,7 @@ use aptos_types::{
 use aptos_vm_types::resolver::ResourceGroupSize;
 use derivative::Derivative;
 use move_core_types::value::MoveTypeLayout;
-use move_vm_types::code::{ModuleCode, SyncModuleCache, WithAddress, WithName};
+use move_vm_types::code::{ModuleCode, SyncModuleCache, WithAddress, WithName, WithSize};
 use std::{
     collections::{
         hash_map::{
@@ -351,6 +351,7 @@ where
     T: Transaction,
     K: Hash + Eq + Ord + Clone,
     VC: Deref<Target = Arc<DC>>,
+    S: WithSize,
 {
     // Return an iterator over the captured reads.
     pub(crate) fn get_read_values_with_delayed_fields(

third_party/move/move-vm/runtime/src/storage/module_storage.rs

@@ -11,7 +11,7 @@ use move_core_types::{
     metadata::Metadata,
 };
 use move_vm_types::{
-    code::{ModuleCache, ModuleCode, ModuleCodeBuilder, WithBytes, WithHash},
+    code::{ModuleCache, ModuleCode, ModuleCodeBuilder, WithBytes, WithHash, WithSize},
     module_cyclic_dependency_error, module_linker_error,
 };
 use std::sync::Arc;
@@ -117,7 +117,7 @@ where
             Verified = Module,
             Extension = E,
         >,
-    E: WithBytes + WithHash,
+    E: WithBytes + WithSize + WithHash,
     V: Clone + Default + Ord,
 {
     fn check_module_exists(
@@ -237,7 +237,7 @@ where
             Verified = Module,
             Extension = E,
         >,
-    E: WithBytes + WithHash,
+    E: WithBytes + WithSize + WithHash,
     V: Clone + Default + Ord,
 {
     let runtime_environment = module_cache_with_context.runtime_environment();

third_party/move/move-vm/types/src/code/cache/module_cache.rs

@@ -138,6 +138,7 @@ pub trait ModuleCache {
     fn num_modules(&self) -> usize;
 }
 
+/// Same as [ModuleCode], additionally storing a version.
 struct VersionedModuleCode<DC, VC, E, V> {
     module_code: Arc<ModuleCode<DC, VC, E>>,
     version: V,
@@ -147,29 +148,35 @@ impl<DC, VC, E, V> VersionedModuleCode<DC, VC, E, V>
 where
     V: Default + Clone + Ord,
 {
+    /// Returns new [ModuleCode] with the specified version.
     fn new(module_code: ModuleCode<DC, VC, E>, version: V) -> Self {
         Self {
             module_code: Arc::new(module_code),
             version,
         }
     }
 
+    /// Returns new [ModuleCode] with the default (storage) version.
     fn new_with_default_version(module_code: ModuleCode<DC, VC, E>) -> Self {
         Self::new(module_code, V::default())
     }
 
+    /// Returns the reference to the stored module.
     fn module_code(&self) -> &Arc<ModuleCode<DC, VC, E>> {
         &self.module_code
     }
 
+    /// Returns the stored module.
     fn into_module_code(self) -> Arc<ModuleCode<DC, VC, E>> {
         self.module_code
     }
 
+    /// Returns the version associated with this module.
     fn version(&self) -> V {
         self.version.clone()
     }
 
+    /// Returns the clone of the module along with its version.
     fn as_module_code_and_version(&self) -> (Arc<ModuleCode<DC, VC, E>>, V) {
         (self.module_code.clone(), self.version.clone())
     }

third_party/move/move-vm/types/src/code/cache/test_types.rs

@@ -1,7 +1,7 @@
 // Copyright (c) The Move Contributors
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::code::ModuleCode;
+use crate::code::{ModuleCode, WithSize};
 use std::{ops::Deref, sync::Arc};
 
 #[derive(Clone, Debug)]
@@ -56,7 +56,6 @@ pub fn mock_verified_code<E>(
 
 #[derive(Clone, Debug)]
 pub struct MockExtension {
-    #[allow(dead_code)]
     size: usize,
 }
 
@@ -66,6 +65,12 @@ impl MockExtension {
     }
 }
 
+impl WithSize for MockExtension {
+    fn size_in_bytes(&self) -> usize {
+        self.size
+    }
+}
+
 pub fn mock_extension(size: usize) -> Arc<MockExtension> {
     Arc::new(MockExtension::new(size))
 }

third_party/move/move-vm/types/src/code/cache/types.rs

@@ -9,7 +9,13 @@ use std::{ops::Deref, sync::Arc};
 
 pub trait WithBytes {
     fn bytes(&self) -> &Bytes;
+}
+
+pub trait WithSize {
+    fn size_in_bytes(&self) -> usize;
+}
 
+impl<T: WithBytes> WithSize for T {
     fn size_in_bytes(&self) -> usize {
         self.bytes().len()
     }

third_party/move/move-vm/types/src/code/mod.rs

@@ -16,6 +16,6 @@ pub use cache::{
         UnsyncModuleCache,
     },
     script_cache::{ambassador_impl_ScriptCache, ScriptCache, SyncScriptCache, UnsyncScriptCache},
-    types::{Code, WithAddress, WithBytes, WithHash, WithName},
+    types::{Code, WithAddress, WithBytes, WithHash, WithName, WithSize},
 };
 pub use storage::ModuleBytesStorage;