[cleanup] Add errors, documentation, reset state on panics

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

@@ -9,7 +9,7 @@ use aptos_types::{
 };
 use aptos_vm_environment::environment::AptosEnvironment;
 use aptos_vm_types::module_and_script_storage::AsAptosCodeStorage;
-use move_binary_format::{errors::Location, CompiledModule};
+use move_binary_format::CompiledModule;
 use move_core_types::{
     account_address::AccountAddress,
     ident_str,
@@ -162,15 +162,19 @@ where
             .as_ref()
             .expect("Environment has to be set")
             .runtime_environment();
-        let struct_name_index_map_size =
-            runtime_environment
-                .struct_name_index_map_size()
-                .map_err(|err| {
-                    // TODO(loader_v2):
-                    //   Is this fine to fail here? We leave the state as not ready forever?
-                    //   Seems like it is better to reset the state and reset everything.
-                    err.finish(Location::Undefined).into_vm_status()
-                })?;
+        let struct_name_index_map_size = match runtime_environment.struct_name_index_map_size() {
+            Err(err) => {
+                // Unlock the cache, reset all states, and return.
+                drop(previous_environment);
+                let err = self.reset_and_return_invariant_violation(&format!(
+                    "Error when getting struct name index map size: {:?}",
+                    err
+                ));
+                return Err(err);
+            },
+            Ok(size) => size,
+        };
+
         if struct_name_index_map_size > self.config.struct_name_index_map_capacity {
             flush_all_caches = true;
         }
@@ -198,6 +202,8 @@ where
         // ensure it is not possible to execute blocks concurrently.
         let mut previous_block_id = self.previous_block_id.lock();
         if self.ready_for_next_block() {
+            // This means we are executing concurrently. If so, all-but-one thread will return an
+            // error. Note that the caches are still consistent for that one thread.
             let msg = "Should not be possible to mark block execution end for execution-ready \
                              global cache, check if blocks are executed concurrently";
             return Err(invariant_violation(msg));
@@ -227,6 +233,27 @@ where
     fn mark_not_ready_for_next_block(&self) {
         self.ready_for_next_block.store(false, Ordering::SeqCst);
     }
+
+    /// 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.
+    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();
+
+        // 1. Should be ready for next execution.
+        self.mark_not_ready_for_next_block();
+        // 2. Should contain no environment.
+        *self.previous_environment.lock() = None;
+        // 3. Module cache is empty.
+        self.module_cache.flush_unchecked();
+        // 4. Block ID is unset.
+        *previous_block_id = BlockId::Unset;
+
+        // State reset, unlock.
+        drop(previous_block_id);
+
+        invariant_violation(msg)
+    }
 }
 
 /// Same as [GlobalCacheManagerInner], but uses concrete types used by execution on Aptos instead
@@ -254,6 +281,8 @@ impl GlobalCacheManager {
     ///   3. The size of the struct name re-indexing map is too large.
     ///   4. The size of the module cache is too large.
     ///
+    /// Additionally, if cache is empty, prefetches the framework code into it.
+    ///
     /// Marks [GlobalCacheManager] as not ready for next block execution. If called concurrently,
     /// only a single invocation ever succeeds and other calls return an error.
     pub fn mark_block_execution_start(
@@ -265,21 +294,45 @@ impl GlobalCacheManager {
             .mark_block_execution_start(state_view, previous_block_id)?;
 
         if self.inner.config.prefetch_framework_code && self.module_cache().size() == 0 {
-            let runtime_environment = self.environment()?;
+            let code_storage = state_view.as_aptos_code_storage(self.environment()?);
 
-            let code_storage = state_view.as_aptos_code_storage(runtime_environment);
+            // If framework code exists in storage, the transitive closure will be verified and
+            // cached.
             let result = code_storage
                 .fetch_verified_module(&AccountAddress::ONE, ident_str!("transaction_validation"));
-            // Framework must exist, prefetch.
-            if let Ok(Some(_)) = result {
-                // TODO(loader_v2): Replace with invariant violations!
-                self.inner
-                    .module_cache
-                    .insert_verified_unchecked(code_storage.into_verified_module_code_iter())
-                    .unwrap();
+
+            match result {
+                Ok(Some(_)) => {
+                    // Framework must have been loaded. Drain verified modules from local cache
+                    // into global cache.
+                    let verified_module_code_iter = code_storage
+                        .into_verified_module_code_iter()
+                        .map_err(|err| {
+                        let msg = format!(
+                            "Unable to convert cached modules into verified code: {:?}",
+                            err
+                        );
+                        self.inner.reset_and_return_invariant_violation(&msg)
+                    })?;
+                    self.inner
+                        .module_cache
+                        .insert_verified_unchecked(verified_module_code_iter)
+                        .map_err(|err| {
+                            let msg = format!("Unable to cache verified framework: {:?}", err);
+                            self.inner.reset_and_return_invariant_violation(&msg)
+                        })?;
+                },
+                Ok(None) => {
+                    // No framework in the state, do nothing.
+                },
+                Err(err) => {
+                    // There should be no errors when pre-fetching the framework, if there are, we
+                    // better return an error here.
+                    let msg = format!("Error when pre-fetching the framework: {:?}", err);
+                    return Err(self.inner.reset_and_return_invariant_violation(&msg));
+                },
             }
         }
-
         Ok(())
     }
 
@@ -301,6 +354,7 @@ impl GlobalCacheManager {
             .lock()
             .clone()
             .ok_or_else(|| {
+                // Note: we do not expect this to happen (this is really more of an unreachable).
                 invariant_violation("Environment must always be set at block execution start")
             })
     }

aptos-move/aptos-vm-types/src/module_and_script_storage/state_view_adapter.rs

@@ -4,6 +4,7 @@
 use crate::module_and_script_storage::module_storage::AptosModuleStorage;
 use ambassador::Delegate;
 use aptos_types::{
+    error::PanicError,
     state_store::{state_key::StateKey, state_value::StateValueMetadata, StateView},
     vm::modules::AptosModuleExtension,
 };
@@ -82,17 +83,27 @@ impl<'s, S: StateView, E: WithRuntimeEnvironment> AptosCodeStorageAdapter<'s, S,
         Self { storage }
     }
 
+    /// Drains cached verified modules from the code storage, transforming them into format used by
+    /// global caches (i.e., with extension and no versioning). Should only be called when the code
+    /// storage borrows [StateView].
     pub fn into_verified_module_code_iter(
         self,
-    ) -> impl Iterator<
-        Item = (
-            ModuleId,
-            Arc<ModuleCode<CompiledModule, Module, AptosModuleExtension, Option<u32>>>,
-        ),
+    ) -> Result<
+        impl Iterator<
+            Item = (
+                ModuleId,
+                Arc<ModuleCode<CompiledModule, Module, AptosModuleExtension, Option<u32>>>,
+            ),
+        >,
+        PanicError,
     > {
         let state_view = match self.storage.module_storage().byte_storage().state_view {
             BorrowedOrOwned::Borrowed(state_view) => state_view,
-            BorrowedOrOwned::Owned(_) => unreachable!(),
+            BorrowedOrOwned::Owned(_) => {
+                return Err(PanicError::CodeInvariantError(
+                    "Verified modules should only be extracted from borrowed state".to_string(),
+                ))
+            },
         };
 
         let mut modules_to_add = vec![];
@@ -101,14 +112,33 @@ impl<'s, S: StateView, E: WithRuntimeEnvironment> AptosCodeStorageAdapter<'s, S,
             .into_module_storage()
             .into_verified_modules_iter()
         {
-            let state_key = StateKey::module_id(&key);
-            // TODO(loader_v2): Replace with invariant violations!
-            let state_value = state_view.get_state_value(&state_key).unwrap().unwrap();
-            let extension = AptosModuleExtension::new(state_value);
+            // We have cached the module previously, so we must be able to find it in storage.
+            let extension = state_view
+                .get_state_value(&StateKey::module_id(&key))
+                .map_err(|err| {
+                    let msg = format!(
+                        "Failed to retrieve module {}::{} from storage {:?}",
+                        key.address(),
+                        key.name(),
+                        err
+                    );
+                    PanicError::CodeInvariantError(msg)
+                })?
+                .map(AptosModuleExtension::new)
+                .ok_or_else(|| {
+                    let msg = format!(
+                        "Module {}::{} should exist, but it does not anymore",
+                        key.address(),
+                        key.name()
+                    );
+                    PanicError::CodeInvariantError(msg)
+                })?;
+
+            // We are using storage version here.
             let module = ModuleCode::from_verified_ref(verified_code, Arc::new(extension), None);
             modules_to_add.push((key, Arc::new(module)))
         }
-        modules_to_add.into_iter()
+        Ok(modules_to_add.into_iter())
     }
 }
 

third_party/move/move-vm/runtime/src/storage/implementations/unsync_code_storage.rs

@@ -29,11 +29,12 @@ use std::sync::Arc;
 pub struct UnsyncCodeStorage<M>(UnsyncCodeStorageImpl<M>);
 
 impl<M: ModuleStorage> UnsyncCodeStorage<M> {
-    /// Returns the underlying module storage used by this code storage.
+    /// Returns the reference to the underlying module storage used by this code storage.
     pub fn module_storage(&self) -> &M {
         &self.0.module_storage
     }
 
+    /// Returns the underlying module storage used by this code storage.
     pub fn into_module_storage(self) -> M {
         self.0.module_storage
     }

third_party/move/move-vm/runtime/src/storage/implementations/unsync_module_storage.rs

@@ -167,6 +167,7 @@ impl<'s, S: ModuleBytesStorage, E: WithRuntimeEnvironment> UnsyncModuleStorage<'
         &self.0.base_storage
     }
 
+    /// Returns an iterator of all modules that have been cached and verified.
     pub fn into_verified_modules_iter(self) -> impl Iterator<Item = (ModuleId, Arc<Module>)> {
         self.0
             .module_cache