Skip to content

Commit

Permalink
Refactor sui benchmark code for simplicity (MystenLabs#9261)
Browse files Browse the repository at this point in the history
## Description 

This PR does a few things to simplify adding new workloads:

1. Simplifications in workload_configuration.rs: There is no need to
handle individual workload in an error prone way to configure workload
weight, tps, etc. It only requires one call to the workload to configure
it now. For example, there is only one call to shared_counter workload
to configure it which looks like this:
```
let shared_workload = SharedCounterWorkloadBuilder::from(
            shared_counter_weight as f32 / total_weight as f32,
            target_qps,
            num_workers,
            in_flight_ratio,
            shared_counter_hotness_factor,
        );
```
2. Simplification in generating gas: There is no need to handle
individual workload while generating gas. The call to generate gas is
handled by a trait now in:
```
pub async fn generate(
        &mut self,
        builders: Vec<Box<dyn WorkloadBuilder<dyn Payload>>>,
        gas_price: u64,
        chunk_size: u64,
    ) -> Result<Vec<Box<dyn Workload<dyn Payload>>>> {
```
3. Remove the concept of combination workload as it was hard to maintain
while it was not being used in production anymore. This helped remove a
lot of code which simplified the crate a lot.
4. Structs like `WorkloadInitGas`, `WorkloadPayloadGas`, etc are all
removed now
5. `Payload` and `Workload` abstractions are much cleaner
## Test Plan 

Existing tests
  • Loading branch information
sadhansood authored Mar 14, 2023
1 parent 17c7ed9 commit 9e7d916
Show file tree
Hide file tree
Showing 17 changed files with 771 additions and 1,313 deletions.
200 changes: 200 additions & 0 deletions crates/sui-benchmark/src/bank.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use crate::util::{make_pay_tx, UpdatedAndNewlyMintedGasCoins};
use crate::workloads::payload::Payload;
use crate::workloads::workload::{Workload, WorkloadBuilder};
use crate::workloads::{Gas, GasCoinConfig};
use crate::ValidatorProxy;
use anyhow::{Error, Result};
use itertools::Itertools;
use move_core_types::language_storage::TypeTag;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use sui_types::base_types::{ObjectRef, SuiAddress};
use sui_types::crypto::AccountKeyPair;
use sui_types::gas_coin::GAS;
use sui_types::messages::{
CallArg, ObjectArg, TransactionData, VerifiedTransaction, DUMMY_GAS_PRICE,
};
use sui_types::utils::to_sender_signed_transaction;
use sui_types::{coin, SUI_FRAMEWORK_OBJECT_ID};

/// Bank is used for generating gas for running the benchmark. It is initialized with two gas coins i.e.
/// `pay_coin` which is split into smaller gas coins and `primary_gas` which is the gas coin used
/// for executing coin split transactions
#[derive(Clone)]
pub struct BenchmarkBank {
pub proxy: Arc<dyn ValidatorProxy + Send + Sync>,
// Gas to use for execution of gas generation transaction
pub primary_gas: Gas,
// Coin to use for splitting and generating small gas coins
pub pay_coin: Gas,
}

impl BenchmarkBank {
pub fn new(
proxy: Arc<dyn ValidatorProxy + Send + Sync>,
primary_gas: Gas,
pay_coin: Gas,
) -> Self {
BenchmarkBank {
proxy,
primary_gas,
pay_coin,
}
}
pub async fn generate(
&mut self,
builders: Vec<Box<dyn WorkloadBuilder<dyn Payload>>>,
gas_price: u64,
chunk_size: u64,
) -> Result<Vec<Box<dyn Workload<dyn Payload>>>> {
let mut coin_configs = VecDeque::new();
for builder in builders.iter() {
let init_gas_config = builder.generate_coin_config_for_init().await;
let payload_gas_config = builder.generate_coin_config_for_payloads().await;
coin_configs.push_back(init_gas_config);
coin_configs.push_back(payload_gas_config);
}
let mut all_coin_configs = vec![];
coin_configs
.iter()
.for_each(|v| all_coin_configs.extend(v.clone()));

let mut new_gas_coins: Vec<Gas> = vec![];
let chunked_coin_configs = all_coin_configs.chunks(chunk_size as usize);
eprintln!("Number of gas requests = {}", chunked_coin_configs.len());
for chunk in chunked_coin_configs {
let (updated_primary_gas, updated_coin, gas_coins) =
self.split_coin_and_pay(chunk, gas_price).await?;
self.primary_gas = updated_primary_gas;
self.pay_coin = updated_coin;
new_gas_coins.extend(gas_coins);
}
let mut workloads = vec![];
for builder in builders.iter() {
let init_gas_config = coin_configs.pop_front().unwrap();
let payload_gas_config = coin_configs.pop_front().unwrap();
let init_gas: Vec<Gas> = init_gas_config
.iter()
.map(|c| {
let (index, _) = new_gas_coins
.iter()
.find_position(|g| g.1 == c.address)
.unwrap();
new_gas_coins.remove(index)
})
.collect();
let payload_gas: Vec<Gas> = payload_gas_config
.iter()
.map(|c| {
let (index, _) = new_gas_coins
.iter()
.find_position(|g| g.1 == c.address)
.unwrap();
new_gas_coins.remove(index)
})
.collect();
workloads.push(builder.build(init_gas, payload_gas).await);
}
Ok(workloads)
}
fn make_split_coin_tx(
&self,
split_amounts: Vec<u64>,
gas_price: Option<u64>,
keypair: &AccountKeyPair,
) -> Result<VerifiedTransaction> {
let split_coin = TransactionData::new_move_call(
self.primary_gas.1,
SUI_FRAMEWORK_OBJECT_ID,
coin::PAY_MODULE_NAME.to_owned(),
coin::PAY_SPLIT_VEC_FUNC_NAME.to_owned(),
vec![TypeTag::Struct(Box::new(GAS::type_()))],
self.primary_gas.0,
vec![
CallArg::Object(ObjectArg::ImmOrOwnedObject(self.pay_coin.0)),
CallArg::Pure(bcs::to_bytes(&split_amounts).unwrap()),
],
1000000,
gas_price.unwrap_or(DUMMY_GAS_PRICE),
)?;
let verified_tx = to_sender_signed_transaction(split_coin, keypair);
Ok(verified_tx)
}
async fn split_coin_and_pay(
&mut self,
coin_configs: &[GasCoinConfig],
gas_price: u64,
) -> Result<UpdatedAndNewlyMintedGasCoins> {
// split one coin into smaller coins of different amounts and send them to recipients
let split_amounts: Vec<u64> = coin_configs.iter().map(|c| c.amount).collect();
// TODO: Instead of splitting the coin and then using pay tx to transfer it to recipients,
// we can do both in one tx with pay_sui which will split the coin out for us before
// transferring it to recipients
let verified_tx =
self.make_split_coin_tx(split_amounts.clone(), Some(gas_price), &self.primary_gas.2)?;
let effects = self.proxy.execute_transaction(verified_tx.into()).await?;
let updated_gas = effects
.mutated()
.into_iter()
.find(|(k, _)| k.0 == self.primary_gas.0 .0)
.ok_or("Input gas missing in the effects")
.map_err(Error::msg)?;
let created_coins: Vec<ObjectRef> = effects.created().into_iter().map(|c| c.0).collect();
assert_eq!(created_coins.len(), split_amounts.len());
let updated_coin = effects
.mutated()
.into_iter()
.find(|(k, _)| k.0 == self.pay_coin.0 .0)
.ok_or("Input gas missing in the effects")
.map_err(Error::msg)?;
let recipient_addresses: Vec<SuiAddress> = coin_configs.iter().map(|g| g.address).collect();
let verified_tx = make_pay_tx(
created_coins,
self.primary_gas.1,
recipient_addresses,
split_amounts,
updated_gas.0,
&self.primary_gas.2,
Some(gas_price),
)?;
let effects = self.proxy.execute_transaction(verified_tx.into()).await?;
let address_map: HashMap<SuiAddress, Arc<AccountKeyPair>> = coin_configs
.iter()
.map(|c| (c.address, c.keypair.clone()))
.collect();
let transferred_coins: Result<Vec<Gas>> = effects
.created()
.into_iter()
.map(|c| {
let address = c.1.get_owner_address()?;
let keypair = address_map
.get(&address)
.ok_or("Owner address missing in the address map")
.map_err(Error::msg)?;
Ok((c.0, address, keypair.clone()))
})
.collect();
let updated_gas = effects
.mutated()
.into_iter()
.find(|(k, _)| k.0 == self.primary_gas.0 .0)
.ok_or("Input gas missing in the effects")
.map_err(Error::msg)?;
Ok((
(
updated_gas.0,
updated_gas.1.get_owner_address()?,
self.primary_gas.2.clone(),
),
(
updated_coin.0,
updated_coin.1.get_owner_address()?,
self.primary_gas.2.clone(),
),
transferred_coins?,
))
}
}
Loading

0 comments on commit 9e7d916

Please sign in to comment.