add code download mode for comparison testing tool

Author: rahxephon89

Cargo.lock

@@ -35,6 +35,8 @@ pub struct DataCollection {
     current_dir: PathBuf,
     batch_size: u64,
     dump_write_set: bool,
+    skip_txn_execution: bool,
+    skip_source_compilation: bool,
     filter_condition: FilterCondition,
     enable_features: Vec<FeatureFlag>,
     disable_features: Vec<FeatureFlag>,
@@ -48,6 +50,8 @@ impl DataCollection {
         skip_failed_txns: bool,
         skip_publish_txns: bool,
         dump_write_set: bool,
+        skip_txn_execution: bool,
+        skip_source_compilation: bool,
         skip_source_code: bool,
         target_account: Option<AccountAddress>,
         enable_features: Vec<FeatureFlag>,
@@ -58,6 +62,8 @@ impl DataCollection {
             current_dir,
             batch_size,
             dump_write_set,
+            skip_txn_execution,
+            skip_source_compilation,
             filter_condition: FilterCondition {
                 skip_failed_txns,
                 skip_publish_txns,
@@ -76,6 +82,8 @@ impl DataCollection {
         skip_failed_txns: bool,
         skip_publish_txns: bool,
         dump_write_set: bool,
+        skip_txn_execution: bool,
+        skip_source_compilation: bool,
         skip_source_code: bool,
         target_account: Option<AccountAddress>,
         enable_features: Vec<FeatureFlag>,
@@ -88,6 +96,8 @@ impl DataCollection {
             skip_failed_txns,
             skip_publish_txns,
             dump_write_set,
+            skip_txn_execution,
+            skip_source_compilation,
             skip_source_code,
             target_account,
             enable_features,
@@ -128,6 +138,7 @@ impl DataCollection {
         compilation_cache: &mut CompilationCache,
         current_dir: PathBuf,
         base_experiments: &[String],
+        skip_source_compilation: bool,
     ) -> Option<PackageInfo> {
         let upgrade_number = if is_aptos_package(&package_name) {
             None
@@ -160,6 +171,7 @@ impl DataCollection {
                 None,
                 base_experiments,
                 &[],
+                skip_source_compilation,
             );
             if let Err(err) = res {
                 eprintln!("{} at: {}", err, version);
@@ -205,12 +217,16 @@ impl DataCollection {
             return Err(anyhow::Error::msg("aptos packages are missing"));
         }
         let mut compiled_cache = CompilationCache::default();
-        compile_aptos_packages(
-            &aptos_commons_path,
-            &mut compiled_cache.base_compiled_package_cache,
-            &base_experiments,
-            "base",
-        )?;
+        // The compiled framework is only injected into the state view for txn
+        // execution; source-only collection does not need it.
+        if !self.skip_txn_execution {
+            compile_aptos_packages(
+                &aptos_commons_path,
+                &mut compiled_cache.base_compiled_package_cache,
+                &base_experiments,
+                "base",
+            )?;
+        }
         let compilation_cache: Arc<Mutex<CompilationCache>> = Arc::new(Mutex::new(compiled_cache));
         let data_manager = Arc::new(Mutex::new(DataManager::new_with_dir_creation(
             &self.current_dir,
@@ -253,40 +269,52 @@ impl DataCollection {
                     let data_manager = data_manager.clone();
                     let index = index_writer.clone();
                     let base_experiments = base_experiments.clone();
-                    let data_state = InMemoryStateStore::default();
-                    let features_to_enable = self.enable_features.clone();
-                    let features_to_disable = self.disable_features.clone();
-                    let cache_v1 = compilation_cache
-                        .lock()
-                        .unwrap()
-                        .base_compiled_package_cache
-                        .clone();
-                    add_aptos_packages_to_state_store(&data_state, &cache_v1);
-                    let state_view = DataStateView::new_with_data_reads_and_code(
-                        self.debugger.clone(),
-                        version,
-                        data_state,
-                        features_to_enable,
-                        features_to_disable,
-                    );
+                    let skip_source_compilation = self.skip_source_compilation;
+                    // Source-only collection skips execution, so it needs no
+                    // state view (and none of its remote state reads).
+                    let state_view_opt = if self.skip_txn_execution {
+                        None
+                    } else {
+                        let data_state = InMemoryStateStore::default();
+                        let features_to_enable = self.enable_features.clone();
+                        let features_to_disable = self.disable_features.clone();
+                        let cache_v1 = compilation_cache
+                            .lock()
+                            .unwrap()
+                            .base_compiled_package_cache
+                            .clone();
+                        add_aptos_packages_to_state_store(&data_state, &cache_v1);
+                        Some(DataStateView::new_with_data_reads_and_code(
+                            self.debugger.clone(),
+                            version,
+                            data_state,
+                            features_to_enable,
+                            features_to_disable,
+                        ))
+                    };
 
                     let txn_execution_thread = tokio::task::spawn_blocking(move || {
-                        // Use default info to improve performance
-                        let aux_info = PersistedAuxiliaryInfo::None;
+                        let epoch_result_opt = if let Some(state_view) = &state_view_opt {
+                            // Use default info to improve performance
+                            let aux_info = PersistedAuxiliaryInfo::None;
 
-                        let epoch_result_res =
-                            Self::execute_transactions_at_version_with_state_view(
+                            match Self::execute_transactions_at_version_with_state_view(
                                 vec![txn.clone()],
                                 vec![aux_info],
-                                &state_view,
-                            );
-                        if let Err(err) = epoch_result_res {
-                            println!(
-                                "execution error during transaction at version:{} :{}",
-                                version, err
-                            );
-                            return;
-                        }
+                                state_view,
+                            ) {
+                                Ok(outputs) => Some(outputs),
+                                Err(err) => {
+                                    println!(
+                                        "execution error during transaction at version:{} :{}",
+                                        version, err
+                                    );
+                                    return;
+                                },
+                            }
+                        } else {
+                            None
+                        };
 
                         let mut version_idx = TxnIndex {
                             version,
@@ -304,21 +332,26 @@ impl DataCollection {
                                 &mut compilation_cache.lock().unwrap(),
                                 current_dir.clone(),
                                 &base_experiments,
+                                skip_source_compilation,
                             );
                             if package_info_opt.is_none() {
                                 return;
                             }
                             version_idx.package_info = package_info_opt.unwrap();
                         }
 
-                        // dump through data_manager
-                        Self::dump_txn_index(
-                            &mut data_manager.lock().unwrap(),
-                            version_idx,
-                            &state_view.get_state_keys().lock().unwrap(),
-                            epoch_result_res,
-                            dump_write_set,
-                        );
+                        // dump through data_manager (skipped in source-only mode)
+                        if let (Some(state_view), Some(outputs)) =
+                            (state_view_opt, epoch_result_opt)
+                        {
+                            Self::dump_txn_index(
+                                &mut data_manager.lock().unwrap(),
+                                version_idx,
+                                &state_view.get_state_keys().lock().unwrap(),
+                                Ok(outputs),
+                                dump_write_set,
+                            );
+                        }
 
                         // Log version
                         index.lock().unwrap().add_version(version);

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

@@ -376,6 +376,9 @@ struct CompilationCache {
     failed_packages_compared: HashSet<PackageInfo>,
     base_compiled_package_cache: HashMap<PackageInfo, HashMap<ModuleId, Vec<u8>>>,
     compared_compiled_package_cache: HashMap<PackageInfo, HashMap<ModuleId, Vec<u8>>>,
+    /// Packages whose source has been written to disk; used to avoid re-dumping
+    /// when compilation (and thus `compiled_package_map`) is skipped.
+    dumped_packages: HashSet<PackageInfo>,
 }
 
 #[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Hash)]
@@ -501,7 +504,11 @@ fn dump_and_compile_from_package_metadata(
     execution_mode: Option<ExecutionMode>,
     base_experiments: &[String],
     compared_experiments: &[String],
+    skip_source_compilation: bool,
 ) -> anyhow::Result<()> {
+    if skip_source_compilation && compilation_cache.dumped_packages.contains(&package_info) {
+        return Ok(());
+    }
     let root_package_dir = root_dir.join(format!("{}", package_info,));
     if compilation_cache
         .failed_packages_base
@@ -603,6 +610,7 @@ fn dump_and_compile_from_package_metadata(
                         execution_mode,
                         base_experiments,
                         compared_experiments,
+                        skip_source_compilation,
                     )?;
                 }
                 break;
@@ -613,8 +621,14 @@ fn dump_and_compile_from_package_metadata(
     // step 4: dump the fixed manifest file
     let toml_path = root_package_dir.join("Move.toml");
     std::fs::write(toml_path, manifest.to_string()).unwrap();
+    compilation_cache
+        .dumped_packages
+        .insert(package_info.clone());
 
     // step 5: test whether the code can be compiled
+    if skip_source_compilation {
+        return Ok(());
+    }
     if !compilation_cache
         .compiled_package_map
         .contains_key(&package_info)

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

@@ -38,6 +38,14 @@ pub enum Cmd {
         /// Dump the write set of txns
         #[clap(long, default_value_t = false)]
         dump_write_set: bool,
+        /// Only collect source code: skip txn execution and do not dump
+        /// state data, write sets or the txn index
+        #[clap(long, default_value_t = false)]
+        skip_txn_execution: bool,
+        /// Do not compile the dumped source code to validate it; sources are
+        /// still written to disk. Speeds up source-only collection considerably
+        #[clap(long, default_value_t = false)]
+        skip_source_compilation: bool,
         /// With this set, only dump transactions that are sent to this account
         #[clap(long)]
         target_account: Option<AccountAddress>,
@@ -161,6 +169,8 @@ async fn main() -> Result<()> {
             skip_publish_txns,
             skip_source_code_check: skip_source_code,
             dump_write_set,
+            skip_txn_execution,
+            skip_source_compilation,
             target_account,
         } => {
             let batch_size = BATCH_SIZE;
@@ -187,6 +197,8 @@ async fn main() -> Result<()> {
                 skip_failed_txns,
                 skip_publish_txns,
                 dump_write_set,
+                skip_txn_execution,
+                skip_source_compilation,
                 skip_source_code,
                 target_account,
                 args.enable_features,

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

@@ -125,6 +125,7 @@ impl OnlineExecutor {
                 execution_mode,
                 base_experiments,
                 compared_experiments,
+                false,
             );
             if let Err(err) = res {
                 eprintln!("{} at:{}", err, version);

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

@@ -25,6 +25,7 @@ async-recursion = { workspace = true }
 async-trait = { workspace = true }
 bcs = { workspace = true }
 bytes = { workspace = true }
+futures = { workspace = true }
 lru = { workspace = true }
 move-core-types = { workspace = true }
 tokio = { workspace = true }

aptos-move/aptos-validator-interface/Cargo.toml

@@ -30,6 +30,54 @@ impl RestDebuggerInterface {
     pub fn new(client: Client) -> Self {
         Self(client)
     }
+
+    /// Fetches transactions and their infos for `[start, start + limit)`, requesting the
+    /// 100-txn API pages concurrently. Unlike [`AptosValidatorInterface::get_committed_transactions`],
+    /// this skips fetching persisted auxiliary infos, which the filtering path does not use.
+    async fn get_committed_transactions_no_aux(
+        &self,
+        start: Version,
+        limit: u64,
+    ) -> Result<(Vec<Transaction>, Vec<TransactionInfo>)> {
+        const PAGE_SIZE: u64 = 100;
+        let mut chunk_futures = vec![];
+        let mut cursor = start;
+        while cursor < start + limit {
+            let count = PAGE_SIZE.min(start + limit - cursor);
+            let client = self.0.clone();
+            chunk_futures.push(async move {
+                let mut chunk = Vec::with_capacity(count as usize);
+                while (chunk.len() as u64) < count {
+                    let page = client
+                        .get_transactions_bcs(
+                            Some(cursor + chunk.len() as u64),
+                            Some((count - chunk.len() as u64) as u16),
+                        )
+                        .await?
+                        .into_inner();
+                    if page.is_empty() {
+                        return Err(anyhow!(
+                            "no transactions returned at version {}",
+                            cursor + chunk.len() as u64
+                        ));
+                    }
+                    chunk.extend(page);
+                }
+                Ok(chunk)
+            });
+            cursor += count;
+        }
+        let mut txns = Vec::with_capacity(limit as usize);
+        let mut txn_infos = Vec::with_capacity(limit as usize);
+        for chunk in futures::future::join_all(chunk_futures).await {
+            for txn in chunk? {
+                txns.push(txn.transaction);
+                txn_infos.push(txn.info);
+            }
+        }
+        println!("Got {}/{} txns from RestApi.", txns.len(), limit);
+        Ok((txns, txn_infos))
+    }
 }
 
 #[async_recursion]
@@ -276,7 +324,7 @@ impl AptosValidatorInterface for RestDebuggerInterface {
         )>,
     > {
         let mut txns = Vec::with_capacity(limit as usize);
-        let (tns, infos, _auxiliary_infos) = self.get_committed_transactions(start, limit).await?;
+        let (tns, infos) = self.get_committed_transactions_no_aux(start, limit).await?;
         let temp_txns = tns
             .iter()
             .zip(infos)