[benchmark] Add option for package overrides

aptos-move/replay-benchmark/Cargo.toml

@@ -15,6 +15,7 @@ rust-version = { workspace = true }
 [dependencies]
 anyhow = { workspace = true }
 aptos-block-executor = { workspace = true }
+aptos-framework = { workspace = true }
 aptos-gas-schedule = { workspace = true }
 aptos-logger = { workspace = true }
 aptos-move-debugger = { workspace = true }

aptos-move/replay-benchmark/src/commands/initialize.rs

@@ -4,10 +4,11 @@
 use crate::{
     commands::{build_debugger, init_logger_and_metrics, RestAPI},
     generator::InputOutputDiffGenerator,
-    overrides::OverrideConfig,
+    overrides::{OverrideConfig, PackageOverride},
     workload::TransactionBlock,
 };
 use anyhow::{anyhow, bail};
+use aptos_framework::{BuildOptions, BuiltPackage};
 use aptos_gas_schedule::LATEST_GAS_FEATURE_VERSION;
 use aptos_logger::Level;
 use aptos_types::on_chain_config::FeatureFlag;
@@ -55,6 +56,14 @@ pub struct InitializeCommand {
         help = "If set, overrides the gas feature version used by the gas schedule"
     )]
     gas_feature_version: Option<u64>,
+
+    #[clap(
+        long,
+        num_args = 1..,
+        value_delimiter = ' ',
+        help = "List of space-separated paths to compiled / built packages with Move code"
+    )]
+    override_packages: Vec<String>,
 }
 
 impl InitializeCommand {
@@ -84,15 +93,17 @@ impl InitializeCommand {
         })?;
 
         // TODO:
-        //  Right now, only features can be overridden. In the future, we may want to support:
-        //      1. Framework code, e.g., to test performance of new natives or compiler,
-        //      2. Gas schedule, to track the costs of charging gas or tracking limits.
-        //      3. BlockExecutorConfigFromOnchain to experiment with different block cutting based
-        //         on gas limits.
+        //   1. Override gas schedule, to track the costs of charging gas or tracking limits.
+        //   2. BlockExecutorConfigFromOnchain to experiment with different block cutting based
+        //      on gas limits?.
+        //   3. Build options for package overrides.
+        let build_options = BuildOptions::move_2();
+        let package_override = PackageOverride::new(self.override_packages, build_options)?;
         let override_config = OverrideConfig::new(
             self.enable_features,
             self.disable_features,
             self.gas_feature_version,
+            package_override,
         );
 
         let debugger = build_debugger(self.rest_api.rest_endpoint, self.rest_api.api_key)?;

aptos-move/replay-benchmark/src/overrides.rs

@@ -5,31 +5,60 @@
 //! transactions can be replayed on top of a modified state, and we can evaluate how it impacts
 //! performance or other things.
 
+use aptos_framework::{natives::code::PackageRegistry, BuildOptions, BuiltPackage};
 use aptos_logger::error;
 use aptos_types::{
     on_chain_config::{FeatureFlag, Features, GasScheduleV2, OnChainConfig},
     state_store::{state_key::StateKey, state_value::StateValue, StateView},
 };
 use serde::Serialize;
-use std::collections::HashMap;
+use std::{collections::HashMap, path::PathBuf};
 
-/// Stores feature flags to enable/disable, essentially overriding on-chain state.
+pub(crate) struct PackageOverride {
+    packages: Vec<BuiltPackage>,
+    build_options: BuildOptions,
+}
+
+impl PackageOverride {
+    pub(crate) fn new(
+        package_paths: Vec<String>,
+        build_options: BuildOptions,
+    ) -> anyhow::Result<Self> {
+        let packages = package_paths
+            .into_iter()
+            .map(|path| BuiltPackage::build(PathBuf::from(&path), build_options.clone()))
+            .collect::<anyhow::Result<_>>()?;
+        Ok(Self {
+            packages,
+            build_options,
+        })
+    }
+}
+
+/// Stores all state overrides.
 pub struct OverrideConfig {
+    /// Feature flags to enable.
     additional_enabled_features: Vec<FeatureFlag>,
+    /// Feature flags to disable.
     additional_disabled_features: Vec<FeatureFlag>,
+    /// Gas feature version to use.
     gas_feature_version: Option<u64>,
+    /// Information about overridden packages.
+    package_override: PackageOverride,
 }
 
 impl OverrideConfig {
     pub fn new(
         additional_enabled_features: Vec<FeatureFlag>,
         additional_disabled_features: Vec<FeatureFlag>,
         gas_feature_version: Option<u64>,
+        package_override: PackageOverride,
     ) -> Self {
         Self {
             additional_enabled_features,
             additional_disabled_features,
             gas_feature_version,
+            package_override,
         }
     }
 
@@ -73,6 +102,113 @@ impl OverrideConfig {
             state_override.insert(gas_schedule_state_key, gas_schedule_state_value);
         }
 
+        // Override packages.
+        let mut overridden_package_registries = HashMap::new();
+        for package in &self.package_override.packages {
+            // Modify existing package metadata or add new one.
+            let package_address = package
+                .modules()
+                .map(|m| m.self_addr())
+                .last()
+                .expect("Package must contain at least one module");
+            let package_registry_state_key =
+                StateKey::resource(package_address, &PackageRegistry::struct_tag()).unwrap();
+
+            let old_package_state_value =
+                match overridden_package_registries.remove(&package_registry_state_key) {
+                    Some(state_value) => state_value,
+                    None => state_view
+                        .get_state_value(&package_registry_state_key)
+                        .unwrap_or_else(|err| {
+                            panic!(
+                                "Failed to fetch package registry at {}: {:?}",
+                                package_address, err
+                            )
+                        })
+                        .expect("Package registry for override must always exist"),
+                };
+
+            let metadata = package.extract_metadata().unwrap_or_else(|err| {
+                panic!(
+                    "Failed to extract metadata for package {}: {:?}",
+                    package.name(),
+                    err
+                )
+            });
+            let new_package_state_value = old_package_state_value
+                .map_bytes(|bytes| {
+                    let mut package_registry = bcs::from_bytes::<PackageRegistry>(&bytes)
+                        .expect("Package registry should deserialize");
+
+                    let mut metadata_idx = None;
+                    for (idx, package_metadata) in package_registry.packages.iter().enumerate() {
+                        if package_metadata.name == metadata.name {
+                            metadata_idx = Some(idx);
+                            break;
+                        }
+                    }
+                    match metadata_idx {
+                        Some(idx) => {
+                            package_registry.packages[idx] = metadata;
+                        },
+                        None => {
+                            package_registry.packages.push(metadata);
+                        },
+                    }
+
+                    let bytes = bcs::to_bytes(&package_registry)
+                        .expect("Package registry should serialize");
+                    Ok(bytes.into())
+                })
+                .unwrap();
+
+            overridden_package_registries
+                .insert(package_registry_state_key, new_package_state_value);
+
+            // Modify all existing modules or add new ones.
+            let bytecode_version = self.package_override.build_options.bytecode_version;
+            for module in package.modules() {
+                let mut module_bytes = vec![];
+                module
+                    .serialize_for_version(bytecode_version, &mut module_bytes)
+                    .unwrap_or_else(|err| {
+                        panic!(
+                            "Failed to serialize module {}::{}: {:?}",
+                            module.self_addr(),
+                            module.self_name(),
+                            err
+                        )
+                    });
+
+                let state_key = StateKey::module(module.self_addr(), module.self_name());
+                let onchain_state_value =
+                    state_view
+                        .get_state_value(&state_key)
+                        .unwrap_or_else(|err| {
+                            panic!(
+                                "Failed to fetch module {}::{}: {:?}",
+                                module.self_addr(),
+                                module.self_name(),
+                                err
+                            )
+                        });
+                let state_value = match onchain_state_value {
+                    Some(state_value) => {
+                        state_value.map_bytes(|_| Ok(module_bytes.into())).unwrap()
+                    },
+                    None => StateValue::new_legacy(module_bytes.into()),
+                };
+                if state_override.insert(state_key, state_value).is_some() {
+                    panic!(
+                        "Overriding module {}::{} more than once",
+                        module.self_addr(),
+                        module.self_name()
+                    );
+                }
+            }
+        }
+        state_override.extend(overridden_package_registries);
+
         state_override
     }
 }

types/src/transaction/mod.rs

@@ -773,6 +773,10 @@ impl SignedTransaction {
         self.raw_txn.max_gas_amount
     }
 
+    pub fn increase_max_gas_amount_by(&mut self, amount: u64) {
+        self.raw_txn.max_gas_amount += amount;
+    }
+
     pub fn gas_unit_price(&self) -> u64 {
         self.raw_txn.gas_unit_price
     }