From 99fdd00ee0a5c30ff13ebc517d5074b8280da551 Mon Sep 17 00:00:00 2001 From: George Mitenkov Date: Wed, 29 Jan 2025 14:29:09 +0000 Subject: [PATCH] [benchmark] Add option for package overrides --- Cargo.lock | 1 + aptos-move/replay-benchmark/Cargo.toml | 1 + .../src/commands/initialize.rs | 23 ++- aptos-move/replay-benchmark/src/overrides.rs | 140 +++++++++++++++++- types/src/transaction/mod.rs | 4 + 5 files changed, 161 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7b76ccacbd579d..9221e29c83283d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3584,6 +3584,7 @@ dependencies = [ "anyhow", "aptos-block-executor", "aptos-crypto", + "aptos-framework", "aptos-gas-schedule", "aptos-logger", "aptos-move-debugger", diff --git a/aptos-move/replay-benchmark/Cargo.toml b/aptos-move/replay-benchmark/Cargo.toml index 5ef890ad7c54ba..5fff4bf29a667c 100644 --- a/aptos-move/replay-benchmark/Cargo.toml +++ b/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 } diff --git a/aptos-move/replay-benchmark/src/commands/initialize.rs b/aptos-move/replay-benchmark/src/commands/initialize.rs index 9461aeff87d179..196ec69faa3e20 100644 --- a/aptos-move/replay-benchmark/src/commands/initialize.rs +++ b/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, + + #[clap( + long, + num_args = 1.., + value_delimiter = ' ', + help = "List of space-separated paths to compiled / built packages with Move code" + )] + override_packages: Vec, } 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)?; diff --git a/aptos-move/replay-benchmark/src/overrides.rs b/aptos-move/replay-benchmark/src/overrides.rs index b2dc32f1d63846..6433ff00ad6849 100644 --- a/aptos-move/replay-benchmark/src/overrides.rs +++ b/aptos-move/replay-benchmark/src/overrides.rs @@ -5,19 +5,46 @@ //! 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, + build_options: BuildOptions, +} + +impl PackageOverride { + pub(crate) fn new( + package_paths: Vec, + build_options: BuildOptions, + ) -> anyhow::Result { + let packages = package_paths + .into_iter() + .map(|path| BuiltPackage::build(PathBuf::from(&path), build_options.clone())) + .collect::>()?; + Ok(Self { + packages, + build_options, + }) + } +} + +/// Stores all state overrides. pub struct OverrideConfig { + /// Feature flags to enable. additional_enabled_features: Vec, + /// Feature flags to disable. additional_disabled_features: Vec, + /// Gas feature version to use. gas_feature_version: Option, + /// Information about overridden packages. + package_override: PackageOverride, } impl OverrideConfig { @@ -25,11 +52,13 @@ impl OverrideConfig { additional_enabled_features: Vec, additional_disabled_features: Vec, gas_feature_version: Option, + 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::(&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 } } diff --git a/types/src/transaction/mod.rs b/types/src/transaction/mod.rs index 2581e09e0d3461..90e076c09e564a 100644 --- a/types/src/transaction/mod.rs +++ b/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 }