[loader-v2] More tests for global manager (aptos-labs#15563)

Author: georgemitenkov

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

@@ -312,7 +312,15 @@ mod test {
         state_store::{state_key::StateKey, state_value::StateValue, MockStateView},
     };
     use claims::assert_ok;
-    use std::collections::HashMap;
+    use move_core_types::identifier::Identifier;
+    use move_vm_types::{
+        code::{mock_verified_code, MockExtension},
+        loaded_data::runtime_types::StructIdentifier,
+    };
+    use std::{
+        collections::HashMap,
+        sync::atomic::{AtomicU64, Ordering},
+    };
 
     #[test]
     fn test_prefetch_existing_aptos_framework() {
@@ -346,19 +354,48 @@ mod test {
         assert_eq!(module_cache.num_modules(), 0);
     }
 
-    #[allow(dead_code)]
-    fn state_view_with_changed_feature_flag(
-        feature_flag: Option<FeatureFlag>,
-    ) -> MockStateView<StateKey> {
+    fn add_struct_identifier<K, D, V, E>(manager: &mut ModuleCacheManager<K, D, V, E>, name: &str)
+    where
+        K: Hash + Eq + Clone,
+        V: Deref<Target = Arc<D>>,
+        E: WithSize,
+    {
+        assert_ok!(manager
+            .environment
+            .as_mut()
+            .unwrap()
+            .runtime_environment()
+            .struct_name_to_idx_for_test(StructIdentifier {
+                module: ModuleId::new(AccountAddress::ZERO, Identifier::new("m").unwrap()),
+                name: Identifier::new(name).unwrap()
+            }));
+    }
+
+    fn assert_struct_name_index_map_size_eq<K, D, V, E>(
+        manager: &ModuleCacheManager<K, D, V, E>,
+        expected: usize,
+    ) where
+        K: Hash + Eq + Clone,
+        V: Deref<Target = Arc<D>>,
+        E: WithSize,
+    {
+        let actual = assert_ok!(manager
+            .environment
+            .as_ref()
+            .unwrap()
+            .runtime_environment()
+            .struct_name_index_map_size());
+        assert_eq!(actual, expected);
+    }
+
+    fn state_view_with_changed_feature_flag(feature_flag: FeatureFlag) -> MockStateView<StateKey> {
         // Tweak feature flags to force a different config.
         let mut features = Features::default();
 
-        if let Some(feature_flag) = feature_flag {
-            if features.is_enabled(feature_flag) {
-                features.disable(feature_flag);
-            } else {
-                features.enable(feature_flag);
-            }
+        if features.is_enabled(feature_flag) {
+            features.disable(feature_flag);
+        } else {
+            features.enable(feature_flag);
         }
 
         MockStateView::new(HashMap::from([(
@@ -368,28 +405,197 @@ mod test {
     }
 
     #[test]
-    fn test_check_ready_sets_transaction_slice_metadata() {
+    fn test_check_ready() {
+        let mut manager = ModuleCacheManager::new();
+        assert_eq!(
+            manager.transaction_slice_metadata,
+            TransactionSliceMetadata::Unknown
+        );
+        assert_eq!(manager.module_cache.num_modules(), 0);
+
         let state_view = MockStateView::empty();
         let config = BlockExecutorModuleCacheLocalConfig {
             prefetch_framework_code: false,
-            max_module_cache_size_in_bytes: 8,
+            max_module_cache_size_in_bytes: 32,
             max_struct_name_index_map_num_entries: 2,
         };
 
+        // Populate the cache for testing.
+        manager
+            .module_cache
+            .insert(0, mock_verified_code(0, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(1, mock_verified_code(1, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(2, mock_verified_code(2, MockExtension::new(8)));
+
+        // Case 1: Initial set-up, modules should not be cached. Metadata and environment are set.
+        let metadata_1 = TransactionSliceMetadata::block_from_u64(0, 1);
+        assert_ok!(manager.check_ready(AptosEnvironment::new(&state_view), &config, metadata_1));
+        assert_eq!(manager.transaction_slice_metadata, metadata_1);
+        assert!(manager.environment.is_some());
+        assert_eq!(manager.module_cache.num_modules(), 0);
+
+        add_struct_identifier(&mut manager, "foo");
+        assert_struct_name_index_map_size_eq(&manager, 1);
+        manager
+            .module_cache
+            .insert(0, mock_verified_code(0, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(1, mock_verified_code(1, MockExtension::new(8)));
+        assert_eq!(manager.module_cache.num_modules(), 2);
+
+        // Case 2: Different metadata => cache is flushed. Here we pass a deep copy of environment.
+        let metadata_2 = TransactionSliceMetadata::block_from_u64(2, 3);
+        assert_ok!(manager.check_ready(AptosEnvironment::new(&state_view), &config, metadata_2));
+        assert_eq!(manager.transaction_slice_metadata, metadata_2);
+        assert!(manager.environment.is_some());
+        assert_eq!(manager.module_cache.num_modules(), 0);
+        assert_struct_name_index_map_size_eq(&manager, 0);
+
+        add_struct_identifier(&mut manager, "foo");
+        add_struct_identifier(&mut manager, "bar");
+        assert_struct_name_index_map_size_eq(&manager, 2);
+        manager
+            .module_cache
+            .insert(0, mock_verified_code(0, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(1, mock_verified_code(1, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(2, mock_verified_code(2, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(3, mock_verified_code(3, MockExtension::new(8)));
+        assert_eq!(manager.module_cache.num_modules(), 4);
+
+        // Case 3: Metadata follows immediately after and environment is the same. Cache is not
+        //         flushed.
+        let metadata_3 = TransactionSliceMetadata::block_from_u64(3, 4);
+        assert!(metadata_3.is_immediately_after(&metadata_2));
+
+        assert_ok!(manager.check_ready(AptosEnvironment::new(&state_view), &config, metadata_3));
+        assert_eq!(manager.transaction_slice_metadata, metadata_3);
+        assert!(manager.environment.is_some());
+        assert_eq!(manager.module_cache.num_modules(), 4);
+        assert_eq!(manager.module_cache.size_in_bytes(), 32);
+        assert_struct_name_index_map_size_eq(&manager, 2);
+
+        manager
+            .module_cache
+            .insert(4, mock_verified_code(4, MockExtension::new(8)));
+        assert_eq!(manager.module_cache.num_modules(), 5);
+        assert_eq!(manager.module_cache.size_in_bytes(), 40);
+
+        // Case 4: Too many modules cached.
+        let metadata_4 = TransactionSliceMetadata::block_from_u64(4, 5);
+        assert!(metadata_4.is_immediately_after(&metadata_3));
+
+        assert_ok!(manager.check_ready(AptosEnvironment::new(&state_view), &config, metadata_4));
+        assert_eq!(manager.transaction_slice_metadata, metadata_4);
+        assert!(manager.environment.is_some());
+        assert_eq!(manager.module_cache.num_modules(), 0);
+        assert_struct_name_index_map_size_eq(&manager, 2);
+
+        manager
+            .module_cache
+            .insert(0, mock_verified_code(0, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(1, mock_verified_code(1, MockExtension::new(8)));
+        assert_eq!(manager.module_cache.num_modules(), 2);
+
+        // Case 5: Environment changes.
+        let metadata_5 = TransactionSliceMetadata::block_from_u64(5, 6);
+        assert!(metadata_5.is_immediately_after(&metadata_4));
+
+        let state_view = state_view_with_changed_feature_flag(FeatureFlag::EMIT_FEE_STATEMENT);
+
+        assert_ok!(manager.check_ready(AptosEnvironment::new(&state_view), &config, metadata_5));
+        assert_eq!(manager.transaction_slice_metadata, metadata_5);
+        assert!(manager.environment.is_some());
+        assert_eq!(manager.module_cache.num_modules(), 0);
+        assert_struct_name_index_map_size_eq(&manager, 0);
+
+        add_struct_identifier(&mut manager, "foo");
+        add_struct_identifier(&mut manager, "bar");
+        add_struct_identifier(&mut manager, "baz");
+        assert_struct_name_index_map_size_eq(&manager, 3);
+        manager
+            .module_cache
+            .insert(0, mock_verified_code(0, MockExtension::new(8)));
+        manager
+            .module_cache
+            .insert(1, mock_verified_code(1, MockExtension::new(8)));
+        assert_eq!(manager.module_cache.num_modules(), 2);
+        assert_eq!(manager.module_cache.size_in_bytes(), 16);
+
+        // Case 6: Type cache is too large.
+        let metadata_6 = TransactionSliceMetadata::block_from_u64(6, 5);
+        assert!(metadata_6.is_immediately_after(&metadata_5));
+
+        assert_ok!(manager.check_ready(AptosEnvironment::new(&state_view), &config, metadata_6));
+        assert_eq!(manager.transaction_slice_metadata, metadata_6);
+        assert!(manager.environment.is_some());
+        assert_eq!(manager.module_cache.num_modules(), 0);
+        assert_struct_name_index_map_size_eq(&manager, 0);
+    }
+
+    #[test]
+    fn test_try_lock_inner_single_thread() {
         let manager = AptosModuleCacheManager::new();
-        assert_eq!(
-            manager.inner.lock().transaction_slice_metadata,
-            TransactionSliceMetadata::Unknown
-        );
 
-        let metadata_1 = TransactionSliceMetadata::block_from_u64(0, 1);
-        assert_ok!(manager.try_lock(&state_view, &config, metadata_1));
-        assert_eq!(manager.inner.lock().transaction_slice_metadata, metadata_1);
+        let state_view = MockStateView::empty();
+        let config = BlockExecutorModuleCacheLocalConfig::default();
+        let metadata = TransactionSliceMetadata::block_from_u64(0, 1);
 
-        let metadata_2 = TransactionSliceMetadata::block_from_u64(1, 2);
-        assert_ok!(manager.try_lock(&state_view, &config, metadata_2));
-        assert_eq!(manager.inner.lock().transaction_slice_metadata, metadata_2);
+        let guard = assert_ok!(manager.try_lock(&state_view, &config, metadata));
+        assert!(matches!(guard, AptosModuleCacheManagerGuard::Guard { .. }));
     }
 
-    // TODO(loader_v2): Add more unit tests like with previous commits.
+    #[test]
+    fn test_try_lock_inner_multiple_threads() {
+        let manager = Arc::new(AptosModuleCacheManager::new());
+
+        let state_view = Arc::new(MockStateView::empty());
+        let config = Arc::new(BlockExecutorModuleCacheLocalConfig::default());
+        let metadata = TransactionSliceMetadata::block_from_u64(0, 1);
+
+        let counter = Arc::new(AtomicU64::new(0));
+        let num_threads = 8;
+        let mut handles = Vec::with_capacity(num_threads);
+
+        for _ in 0..num_threads {
+            let handle = std::thread::spawn({
+                let manager = manager.clone();
+                let state_view = state_view.clone();
+                let config = config.clone();
+                let counter = counter.clone();
+
+                move || {
+                    let guard = assert_ok!(manager.try_lock_inner(&state_view, &config, metadata));
+
+                    // Wait for all threads to complete.
+                    counter.fetch_add(1, Ordering::SeqCst);
+                    loop {
+                        if counter.load(Ordering::SeqCst) == num_threads as u64 {
+                            break;
+                        }
+                    }
+                    if matches!(guard, AptosModuleCacheManagerGuard::Guard { .. }) {
+                        1
+                    } else {
+                        0
+                    }
+                }
+            });
+            handles.push(handle);
+        }
+        let sum = handles.into_iter().map(|h| h.join().unwrap()).sum::<i32>();
+        assert_eq!(sum, 1);
+    }
 }