[benchmark] Add concept of workload (#14198)

lxfind · web-flow · commit 6f1533718130 · 2023-10-11T15:07:48.000-07:00
## Description 

This PR introduces the concept of Workload, and wraps all workload
related code/data into it. It cleans up the command interface to make it
more extensible.

## Test Plan 

CI

---
If your changes are not user-facing and not a breaking change, you can
skip the following section. Otherwise, please indicate what changed, and
then add to the Release Notes section as highlighted during the release
process.

### Type of Change (Check all that apply)

- [ ] protocol change
- [ ] user-visible impact
- [ ] breaking change for a client SDKs
- [ ] breaking change for FNs (FN binary must upgrade)
- [ ] breaking change for validators or node operators (must upgrade
binaries)
- [ ] breaking change for on-chain data layout
- [ ] necessitate either a data wipe or data migration

### Release notes
diff --git a/crates/sui-single-node-benchmark/src/benchmark_context.rs b/crates/sui-single-node-benchmark/src/benchmark_context.rs
@@ -4,15 +4,15 @@
 use crate::command::Component;
 use crate::mock_consensus::ConsensusMode;
 use crate::single_node::SingleValidator;
-use crate::tx_generator::TxGenerator;
+use crate::tx_generator::{RootObjectCreateTxGenerator, TxGenerator};
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use sui_types::base_types::{ObjectID, ObjectRef, SuiAddress, SUI_ADDRESS_LENGTH};
 use sui_types::crypto::{get_account_key_pair, AccountKeyPair};
-use sui_types::effects::TransactionEffects;
+use sui_types::effects::{TransactionEffects, TransactionEffectsAPI};
 use sui_types::messages_grpc::HandleTransactionResponse;
 use sui_types::object::Object;
 use sui_types::transaction::{CertifiedTransaction, SignedTransaction, Transaction};
@@ -133,6 +133,50 @@ impl BenchmarkContext {
             .await
     }
 
+    /// In order to benchmark transactions that can read dynamic fields, we must first create
+    /// a root object with dynamic fields for each account address.
+    pub(crate) async fn preparing_dynamic_fields(
+        &mut self,
+        move_package: ObjectID,
+        num_dynamic_fields: u64,
+    ) -> HashMap<SuiAddress, ObjectRef> {
+        if num_dynamic_fields == 0 {
+            return HashMap::new();
+        }
+
+        info!("Preparing root object with dynamic fields");
+        let root_object_create_transactions = self
+            .generate_transactions(Arc::new(RootObjectCreateTxGenerator::new(
+                move_package,
+                num_dynamic_fields,
+            )))
+            .await;
+        let results = self
+            .execute_transactions_immediately(root_object_create_transactions)
+            .await;
+        let mut root_objects = HashMap::new();
+        let mut new_gas_objects = HashMap::new();
+        for effects in results {
+            let (owner, root_object) = effects
+                .created()
+                .into_iter()
+                .filter_map(|(oref, owner)| {
+                    owner
+                        .get_address_owner_address()
+                        .ok()
+                        .map(|owner| (owner, oref))
+                })
+                .next()
+                .unwrap();
+            let gas_object = effects.gas_object().0;
+            root_objects.insert(owner, root_object);
+            new_gas_objects.insert(gas_object.0, gas_object);
+        }
+        self.refresh_gas_objects(new_gas_objects);
+        info!("Finished preparing root object with dynamic fields");
+        root_objects
+    }
+
     pub(crate) async fn generate_transactions(
         &self,
         tx_generator: Arc<dyn TxGenerator>,
@@ -199,7 +243,7 @@ impl BenchmarkContext {
         results.into_iter().map(|r| r.unwrap()).collect()
     }
 
-    pub(crate) async fn execute_transactions_immediately(
+    async fn execute_transactions_immediately(
         &self,
         transactions: Vec<Transaction>,
     ) -> Vec<TransactionEffects> {
@@ -214,10 +258,7 @@ impl BenchmarkContext {
         results.into_iter().map(|r| r.unwrap()).collect()
     }
 
-    pub(crate) fn refresh_gas_objects(
-        &mut self,
-        mut new_gas_objects: HashMap<ObjectID, ObjectRef>,
-    ) {
+    fn refresh_gas_objects(&mut self, mut new_gas_objects: HashMap<ObjectID, ObjectRef>) {
         info!("Refreshing gas objects");
         for gas_objects in self.gas_object_refs.iter_mut() {
             let refreshed_gas_objects: Vec<_> = gas_objects
diff --git a/crates/sui-single-node-benchmark/src/command.rs b/crates/sui-single-node-benchmark/src/command.rs
@@ -1,49 +1,58 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use clap::{Parser, ValueEnum};
+use clap::{Parser, Subcommand, ValueEnum};
 use strum_macros::EnumIter;
 
 #[derive(Parser)]
-#[command(
+#[clap(
     name = "sui-single-node-benchmark",
     about = "Benchmark a single validator node",
     rename_all = "kebab-case",
     author,
     version
 )]
-pub enum Command {
-    #[command(name = "no-move")]
-    NoMove {
-        #[arg(
-            long,
-            default_value_t = 500000,
-            help = "Number of transactions to submit"
-        )]
-        tx_count: u64,
-        #[arg(
-            long,
-            default_value = "baseline",
-            ignore_case = true,
-            help = "Which component to benchmark"
-        )]
-        component: Component,
-    },
-    #[command(name = "move")]
+pub struct Command {
+    #[arg(
+        long,
+        default_value_t = 500000,
+        help = "Number of transactions to submit"
+    )]
+    pub tx_count: u64,
+    #[arg(
+        long,
+        default_value = "baseline",
+        ignore_case = true,
+        help = "Which component to benchmark"
+    )]
+    pub component: Component,
+    #[clap(subcommand)]
+    pub workload: WorkloadKind,
+}
+
+#[derive(Copy, Clone, EnumIter, ValueEnum)]
+pub enum Component {
+    /// Baseline includes the execution and storage layer only.
+    Baseline,
+    /// On top of Baseline, this schedules transactions through the transaction manager.
+    WithTxManager,
+    /// This goes through the `handle_certificate` entry point on authority_server, which includes
+    /// certificate verification, transaction manager, as well as a noop consensus layer. The noop
+    /// consensus layer does absolutely nothing when receiving a transaction in consensus.
+    ValidatorWithoutConsensus,
+    /// Similar to ValidatorWithNoopConsensus, but the consensus layer contains a fake consensus
+    /// protocol that basically sequences transactions in order. It then verify the transaction
+    /// and store the sequenced transactions into the store. It covers the consensus-independent
+    /// portion of the code in consensus handler.
+    ValidatorWithFakeConsensus,
+    /// Benchmark only validator signing component: `handle_transaction`.
+    TxnSigning,
+}
+
+#[derive(Subcommand)]
+pub enum WorkloadKind {
+    NoMove,
     Move {
-        #[arg(
-            long,
-            default_value_t = 500000,
-            help = "Number of transactions to submit"
-        )]
-        tx_count: u64,
-        #[arg(
-            long,
-            default_value = "baseline",
-            ignore_case = true,
-            help = "Which component to benchmark"
-        )]
-        component: Component,
         #[arg(
             long,
             default_value_t = 2,
@@ -68,22 +77,3 @@ pub enum Command {
         computation: u8,
     },
 }
-
-#[derive(Copy, Clone, EnumIter, ValueEnum)]
-pub enum Component {
-    /// Baseline includes the execution and storage layer only.
-    Baseline,
-    /// On top of Baseline, this schedules transactions through the transaction manager.
-    WithTxManager,
-    /// This goes through the `handle_certificate` entry point on authority_server, which includes
-    /// certificate verification, transaction manager, as well as a noop consensus layer. The noop
-    /// consensus layer does absolutely nothing when receiving a transaction in consensus.
-    ValidatorWithoutConsensus,
-    /// Similar to ValidatorWithNoopConsensus, but the consensus layer contains a fake consensus
-    /// protocol that basically sequences transactions in order. It then verify the transaction
-    /// and store the sequenced transactions into the store. It covers the consensus-independent
-    /// portion of the code in consensus handler.
-    ValidatorWithFakeConsensus,
-    /// Benchmark only validator signing compoment: `handle_transaction`.
-    TxnSigning,
-}
diff --git a/crates/sui-single-node-benchmark/src/execution.rs b/crates/sui-single-node-benchmark/src/execution.rs
@@ -3,107 +3,25 @@
 
 use crate::benchmark_context::BenchmarkContext;
 use crate::command::Component;
-use crate::tx_generator::{MoveTxGenerator, NonMoveTxGenerator, RootObjectCreateTxGenerator};
-use std::collections::HashMap;
-use std::sync::Arc;
-use sui_types::base_types::{ObjectID, ObjectRef, SuiAddress};
+use crate::workload::Workload;
 use sui_types::effects::TransactionEffectsAPI;
 use sui_types::transaction::Transaction;
 use tracing::info;
 
-/// Benchmark simple transfer transactions.
-/// Each transaction transfers the gas object from an address to itself.
-/// The execution does not invoke Move VM, and is considered the cheapest kind of transaction.
-///
-/// \tx_count: the number of transactions to execute.
-/// \component: The component to benchmark.
-pub async fn benchmark_simple_transfer(tx_count: u64, component: Component) {
-    let ctx = BenchmarkContext::new(tx_count, 1, component).await;
-    let transactions = ctx
-        .generate_transactions(Arc::new(NonMoveTxGenerator::new()))
-        .await;
-    benchmark_transactions(&ctx, transactions, component).await;
-}
-
-/// Benchmark Move transactions.
-/// Each transaction is a programmable transaction that performs a series of operations specified
-/// by the parameters.
-///
-/// \tx_count: the number of transactions to execute.
-/// \component: The component to benchmark.
-/// \num_input_objects: the number of address owned input coin objects for each transaction.
-/// These objects will be read during input checking, and merged during execution.
-/// \num_dynamic_fields: the number of dynamic fields read during execution of each transaction.
-/// \computation: Computation intensity for each transaction.
-pub async fn benchmark_move_transactions(
-    tx_count: u64,
-    component: Component,
-    num_input_objects: u8,
-    num_dynamic_fields: u64,
-    computation: u8,
-) {
-    assert!(
-        num_input_objects >= 1,
-        "Each transaction requires at least 1 input object"
-    );
-    let mut ctx = BenchmarkContext::new(tx_count, num_input_objects as u64, component).await;
-    let move_package = ctx.publish_package().await;
-    let root_objects = preparing_dynamic_fields(&mut ctx, move_package.0, num_dynamic_fields).await;
-    let transactions = ctx
-        .generate_transactions(Arc::new(MoveTxGenerator::new(
-            move_package.0,
-            num_input_objects,
-            computation,
-            root_objects,
-        )))
-        .await;
+/// Benchmark a given workload on a specified component.
+/// The different kinds of workloads and components can be found in command.rs.
+pub async fn run_benchmark(workload: Workload, component: Component) {
+    let mut ctx = BenchmarkContext::new(
+        workload.num_accounts(),
+        workload.gas_object_num_per_account(),
+        component,
+    )
+    .await;
+    let tx_generator = workload.create_tx_generator(&mut ctx).await;
+    let transactions = ctx.generate_transactions(tx_generator).await;
     benchmark_transactions(&ctx, transactions, component).await;
 }
 
-/// In order to benchmark transactions that can read dynamic fields, we must first create
-/// a root object with dynamic fields for each account address.
-async fn preparing_dynamic_fields(
-    ctx: &mut BenchmarkContext,
-    move_package: ObjectID,
-    num_dynamic_fields: u64,
-) -> HashMap<SuiAddress, ObjectRef> {
-    if num_dynamic_fields == 0 {
-        return HashMap::new();
-    }
-
-    info!("Preparing root object with dynamic fields");
-    let root_object_create_transactions = ctx
-        .generate_transactions(Arc::new(RootObjectCreateTxGenerator::new(
-            move_package,
-            num_dynamic_fields,
-        )))
-        .await;
-    let results = ctx
-        .execute_transactions_immediately(root_object_create_transactions)
-        .await;
-    let mut root_objects = HashMap::new();
-    let mut new_gas_objects = HashMap::new();
-    for effects in results {
-        let (owner, root_object) = effects
-            .created()
-            .into_iter()
-            .filter_map(|(oref, owner)| {
-                owner
-                    .get_address_owner_address()
-                    .ok()
-                    .map(|owner| (owner, oref))
-            })
-            .next()
-            .unwrap();
-        let gas_object = effects.gas_object().0;
-        root_objects.insert(owner, root_object);
-        new_gas_objects.insert(gas_object.0, gas_object);
-    }
-    ctx.refresh_gas_objects(new_gas_objects);
-    info!("Finished preparing root object with dynamic fields");
-    root_objects
-}
-
 async fn benchmark_transactions(
     ctx: &BenchmarkContext,
     transactions: Vec<Transaction>,
diff --git a/crates/sui-single-node-benchmark/src/lib.rs b/crates/sui-single-node-benchmark/src/lib.rs
@@ -7,3 +7,4 @@ pub mod execution;
 pub(crate) mod mock_consensus;
 pub(crate) mod single_node;
 pub(crate) mod tx_generator;
+pub mod workload;
diff --git a/crates/sui-single-node-benchmark/src/main.rs b/crates/sui-single-node-benchmark/src/main.rs
@@ -3,9 +3,8 @@
 
 use clap::Parser;
 use sui_single_node_benchmark::command::Command;
-use sui_single_node_benchmark::execution::{
-    benchmark_move_transactions, benchmark_simple_transfer,
-};
+use sui_single_node_benchmark::execution::run_benchmark;
+use sui_single_node_benchmark::workload::Workload;
 
 #[tokio::main]
 async fn main() {
@@ -15,26 +14,5 @@ async fn main() {
         .init();
 
     let args = Command::parse();
-    match args {
-        Command::NoMove {
-            tx_count,
-            component,
-        } => benchmark_simple_transfer(tx_count, component).await,
-        Command::Move {
-            tx_count,
-            component,
-            num_input_objects,
-            num_dynamic_fields,
-            computation,
-        } => {
-            benchmark_move_transactions(
-                tx_count,
-                component,
-                num_input_objects,
-                num_dynamic_fields,
-                computation,
-            )
-            .await
-        }
-    }
+    run_benchmark(Workload::new(args.tx_count, args.workload), args.component).await;
 }
diff --git a/crates/sui-single-node-benchmark/src/workload.rs b/crates/sui-single-node-benchmark/src/workload.rs
diff --git a/crates/sui-single-node-benchmark/tests/smoke_tests.rs b/crates/sui-single-node-benchmark/tests/smoke_tests.rs
