[comparison-testing-tool] Add more features to the comparison testing tool (aptos-labs#11890)

* add features

* refactor

* handling review comments

* fix

* always get aptos framework from git

* simplify mismatch info

* add failed cache

* remove clean session

* revert clean session

* handle review comments

* handle comments

* comment followup

* add gas meter to avoid hanging in loop

* update package publish handling

Author: rahxephon89

Cargo.lock

@@ -33,6 +33,7 @@ bcs = { workspace = true }
 clap = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
+lru = { workspace = true }
 move-binary-format = { workspace = true }
 move-cli = { workspace = true }
 move-compiler = { workspace = true }

aptos-move/aptos-e2e-comparison-testing/Cargo.toml

@@ -2,8 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    dump_and_compile_from_package_metadata, is_aptos_package, CompilationCache, DataManager,
-    IndexWriter, PackageInfo, TxnIndex,
+    data_state_view::DataStateView, dump_and_compile_from_package_metadata, is_aptos_package,
+    CompilationCache, DataManager, IndexWriter, PackageInfo, TxnIndex,
 };
 use anyhow::{format_err, Result};
 use aptos_framework::natives::code::PackageMetadata;
@@ -16,9 +16,7 @@ use aptos_types::{
     },
     write_set::TOTAL_SUPPLY_STATE_KEY,
 };
-use aptos_validator_interface::{
-    AptosValidatorInterface, DebuggerStateView, FilterCondition, RestDebuggerInterface,
-};
+use aptos_validator_interface::{AptosValidatorInterface, FilterCondition, RestDebuggerInterface};
 use aptos_vm::{AptosVM, VMExecutor};
 use move_core_types::account_address::AccountAddress;
 use std::{
@@ -45,6 +43,7 @@ impl DataCollection {
         skip_publish_txns: bool,
         dump_write_set: bool,
         skip_source_code: bool,
+        target_account: Option<AccountAddress>,
     ) -> Self {
         Self {
             debugger,
@@ -55,6 +54,7 @@ impl DataCollection {
                 skip_failed_txns,
                 skip_publish_txns,
                 check_source_code: !skip_source_code,
+                target_account,
             },
         }
     }
@@ -67,6 +67,7 @@ impl DataCollection {
         skip_publish_txns: bool,
         dump_write_set: bool,
         skip_source_code: bool,
+        target_account: Option<AccountAddress>,
     ) -> Result<Self> {
         Ok(Self::new(
             Arc::new(RestDebuggerInterface::new(rest_client)),
@@ -76,21 +77,22 @@ impl DataCollection {
             skip_publish_txns,
             dump_write_set,
             skip_source_code,
+            target_account,
         ))
     }
 
     fn execute_transactions_at_version_with_state_view(
         txns: Vec<Transaction>,
-        debugger_stateview: &DebuggerStateView,
+        debugger_state_view: &DataStateView,
     ) -> Result<Vec<TransactionOutput>> {
         let sig_verified_txns: Vec<SignatureVerifiedTransaction> =
             txns.into_iter().map(|x| x.into()).collect::<Vec<_>>();
         // check whether total supply can be retrieved
         // used for debugging the aggregator panic issue, will be removed later
         // FIXME(#10412): remove the assert
-        let val = debugger_stateview.get_state_value(TOTAL_SUPPLY_STATE_KEY.deref());
+        let val = debugger_state_view.get_state_value(TOTAL_SUPPLY_STATE_KEY.deref());
         assert!(val.is_ok() && val.unwrap().is_some());
-        AptosVM::execute_block_no_limit(&sig_verified_txns, debugger_stateview)
+        AptosVM::execute_block_no_limit(&sig_verified_txns, debugger_state_view)
             .map_err(|err| format_err!("Unexpected VM Error: {:?}", err))
     }
 
@@ -114,15 +116,14 @@ impl DataCollection {
             package_name: package_name.clone(),
             upgrade_number,
         };
-
+        if compilation_cache.failed_packages_v1.contains(&package_info) {
+            return None;
+        }
         if !is_aptos_package(&package_name)
             && !compilation_cache
                 .compiled_package_map
                 .contains_key(&package_info)
         {
-            if compilation_cache.failed_packages.contains(&package_info) {
-                return None;
-            }
             let res = dump_and_compile_from_package_metadata(
                 package_info.clone(),
                 current_dir,
@@ -131,7 +132,7 @@ impl DataCollection {
                 None,
             );
             if res.is_err() {
-                println!("compile package failed at:{}", version);
+                eprintln!("{} at: {}", res.unwrap_err(), version);
                 return None;
             }
         }
@@ -167,7 +168,7 @@ impl DataCollection {
         let index_writer = Arc::new(Mutex::new(IndexWriter::new(&self.current_dir)));
 
         let mut cur_version = begin;
-
+        let mut module_registry_map = HashMap::new();
         while cur_version < begin + limit {
             let batch = if cur_version + self.batch_size <= begin + limit {
                 self.batch_size
@@ -176,7 +177,12 @@ impl DataCollection {
             };
             let res_txns = self
                 .debugger
-                .get_and_filter_committed_transactions(cur_version, batch, self.filter_condition)
+                .get_and_filter_committed_transactions(
+                    cur_version,
+                    batch,
+                    self.filter_condition,
+                    &mut module_registry_map,
+                )
                 .await;
             // if error happens when collecting txns, log the version range
             if res_txns.is_err() {
@@ -198,7 +204,7 @@ impl DataCollection {
                     let index = index_writer.clone();
 
                     let state_view =
-                        DebuggerStateView::new_with_data_reads(self.debugger.clone(), version);
+                        DataStateView::new_with_data_reads(self.debugger.clone(), version);
 
                     let txn_execution_thread = tokio::task::spawn_blocking(move || {
                         let epoch_result_res =

aptos-move/aptos-e2e-comparison-testing/src/data_collection.rs

@@ -0,0 +1,86 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use aptos_language_e2e_tests::data_store::FakeDataStore;
+use aptos_types::{
+    state_store::{
+        state_key::StateKey, state_storage_usage::StateStorageUsage, state_value::StateValue,
+        Result as StateViewResult, TStateView,
+    },
+    transaction::Version,
+};
+use aptos_validator_interface::{AptosValidatorInterface, DebuggerStateView};
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+pub struct DataStateView {
+    debugger_view: DebuggerStateView,
+    code_data: Option<FakeDataStore>,
+    data_read_state_keys: Option<Arc<Mutex<HashMap<StateKey, StateValue>>>>,
+}
+use std::ops::DerefMut;
+
+impl DataStateView {
+    pub fn new(
+        db: Arc<dyn AptosValidatorInterface + Send>,
+        version: Version,
+        code_data: FakeDataStore,
+    ) -> Self {
+        Self {
+            debugger_view: DebuggerStateView::new(db, version),
+            code_data: Some(code_data),
+            data_read_state_keys: None,
+        }
+    }
+
+    pub fn new_with_data_reads(
+        db: Arc<dyn AptosValidatorInterface + Send>,
+        version: Version,
+    ) -> Self {
+        Self {
+            debugger_view: DebuggerStateView::new(db, version),
+            code_data: None,
+            data_read_state_keys: Some(Arc::new(Mutex::new(HashMap::new()))),
+        }
+    }
+
+    pub fn get_state_keys(self) -> Arc<Mutex<HashMap<StateKey, StateValue>>> {
+        self.data_read_state_keys.unwrap()
+    }
+}
+
+impl TStateView for DataStateView {
+    type Key = StateKey;
+
+    fn get_state_value(&self, state_key: &StateKey) -> StateViewResult<Option<StateValue>> {
+        if let Some(code) = &self.code_data {
+            if code.contains_key(state_key) {
+                return code.get_state_value(state_key).map_err(Into::into);
+            }
+        }
+        let ret = self.debugger_view.get_state_value(state_key);
+        if let Some(reads) = &self.data_read_state_keys {
+            if !state_key.is_aptos_code()
+                && !reads.lock().unwrap().contains_key(state_key)
+                && ret.is_ok()
+            {
+                let val = ret?;
+                if val.is_some() {
+                    reads
+                        .lock()
+                        .unwrap()
+                        .deref_mut()
+                        .insert(state_key.clone(), val.clone().unwrap());
+                };
+                return Ok(val);
+            }
+        }
+        ret
+    }
+
+    fn get_usage(&self) -> StateViewResult<StateStorageUsage> {
+        unreachable!()
+    }
+}