[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/README.md

@@ -98,13 +98,15 @@ the state. Currently, the only supported overrides are the following:
   1. Forcefully enable a feature flag (`--enable-features F1 F2 ...`).
   2. Forcefully disable a feature flag (`--disable-features F1 F2 ...`).
   3. Forcefully override the gas feature version (`--gas-feature-version V`).
+  4. Override existing on-chain packages (`--override-packages P1 P2 P3`). The paths to the
+     packages must be the path to the source directories.
 
 Feature flags should be spelled in capital letters, e.g., `ENABLE_LOADER_V2`. For the full list of
 available features, see [here](../../types/src/on_chain_config/aptos_features.rs).
 
-Overriding the feature flags can be very useful if you want to experiment with a new feature and
-check its performance as well as the gas usage. For example, if there is a new feature that makes
-MoveVM faster, overriding it for past transactions it is possible to see the execution performance
+Overriding the state can be very useful if you want to experiment with a new feature or Move code,
+and check its performance as well as the gas usage. For example, if there is a new feature that
+makes MoveVM faster, overriding it for past transactions it is possible to see the execution performance
 on historical workloads.
 
 #### Example

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

@@ -7,8 +7,7 @@ use crate::{
     overrides::OverrideConfig,
     workload::TransactionBlock,
 };
-use anyhow::{anyhow, bail};
-use aptos_gas_schedule::LATEST_GAS_FEATURE_VERSION;
+use anyhow::anyhow;
 use aptos_logger::Level;
 use aptos_types::on_chain_config::FeatureFlag;
 use clap::Parser;
@@ -55,26 +54,20 @@ 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 {
     pub async fn initialize_inputs(self) -> anyhow::Result<()> {
         init_logger_and_metrics(self.log_level);
 
-        if !self
-            .enable_features
-            .iter()
-            .all(|f| !self.disable_features.contains(f))
-        {
-            bail!("Enabled and disabled feature flags cannot overlap")
-        }
-        if matches!(self.gas_feature_version, Some(v) if v > LATEST_GAS_FEATURE_VERSION) {
-            bail!(
-                "Gas feature version must be at most the latest one: {}",
-                LATEST_GAS_FEATURE_VERSION
-            );
-        }
-
         let bytes = fs::read(PathBuf::from(&self.transactions_file)).await?;
         let txn_blocks: Vec<TransactionBlock> = bcs::from_bytes(&bytes).map_err(|err| {
             anyhow!(
@@ -84,16 +77,16 @@ 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 override_config = OverrideConfig::new(
             self.enable_features,
             self.disable_features,
             self.gas_feature_version,
-        );
+            self.override_packages,
+        )?;
 
         let debugger = build_debugger(self.rest_api.rest_endpoint, self.rest_api.api_key)?;
         let inputs =

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

@@ -3,34 +3,89 @@
 
 //! Defines different overrides for on-chain state used for benchmarking. With overrides, past
 //! transactions can be replayed on top of a modified state, and we can evaluate how it impacts
-//! performance or other things.
+//! performance or other things. Supported overrides include:
+//!   1. enabling feature flags,
+//!   2. disabling feature flags,
+//!   3. overriding gas feature version,
+//!   4. changing modules (bytecode, metadata, etc.) and package information.
 
+use anyhow::bail;
+use aptos_framework::{natives::code::PackageRegistry, BuildOptions, BuiltPackage};
+use aptos_gas_schedule::LATEST_GAS_FEATURE_VERSION;
 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::{BTreeSet, HashMap},
+    path::PathBuf,
+};
+
+/// Stores information about compiled Move packages and the build options used to create them. Used
+/// by the override configuration to shadow existing on-chain modules with modules defined in these
+/// packages.
+struct PackageOverride {
+    packages: Vec<BuiltPackage>,
+    build_options: BuildOptions,
+}
+
+impl PackageOverride {
+    /// Uses the provided build options to build multiple packages from the specified paths.
+    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 feature flags to enable/disable, essentially overriding on-chain state.
+/// Stores all state overrides.
 pub struct OverrideConfig {
+    /// Feature flags to enable. Invariant: does not overlap with disabled features.
     additional_enabled_features: Vec<FeatureFlag>,
+    /// Feature flags to disable. Invariant: does not overlap with enabled features.
     additional_disabled_features: Vec<FeatureFlag>,
+    /// Gas feature version to use. Invariant: must be at most the latest version.
     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>,
-    ) -> Self {
-        Self {
+        override_packages: Vec<String>,
+    ) -> anyhow::Result<Self> {
+        let build_options = BuildOptions::move_2();
+        let package_override = PackageOverride::new(override_packages, build_options)?;
+
+        if !additional_enabled_features
+            .iter()
+            .all(|f| !additional_disabled_features.contains(f))
+        {
+            bail!("Enabled and disabled feature flags cannot overlap")
+        }
+        if matches!(gas_feature_version, Some(v) if v > LATEST_GAS_FEATURE_VERSION) {
+            bail!(
+                "Gas feature version must be at most the latest one: {}",
+                LATEST_GAS_FEATURE_VERSION
+            );
+        }
+
+        Ok(Self {
             additional_enabled_features,
             additional_disabled_features,
             gas_feature_version,
-        }
+            package_override,
+        })
     }
 
     pub(crate) fn get_state_override(
@@ -73,6 +128,123 @@ 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 addresses = package
+                .modules()
+                .map(|m| *m.self_addr())
+                .collect::<BTreeSet<_>>();
+            assert_eq!(
+                addresses.len(),
+                1,
+                "Modules in the same package must have the same address"
+            );
+
+            let package_address = addresses
+                .last()
+                .expect("Package must contain at least one module");
+            let package_registry_state_key =
+                StateKey::resource(package_address, &PackageRegistry::struct_tag())
+                    .expect("Should always be able to create state key for package registry");
+
+            let old_package_registry_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_registry_state_value = old_package_registry_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())
+                })
+                .expect("Modifying package never returns an error");
+
+            overridden_package_registries
+                .insert(package_registry_state_key, new_package_registry_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
     }
 }