Rename context to workload in benchamrking crate (MystenLabs#3537)

crates/sui-benchmark/Cargo.toml

@@ -28,6 +28,7 @@ prometheus = "0.13.1"
 multiaddr = "0.14.0"
 crossterm = "0.23.2"
 rand = "0.8.5"
+base64 = "0.13.0"
 
 bcs = "0.1.3"
 sui-core = { path = "../sui-core" }

crates/sui-benchmark/src/lib.rs

@@ -2,4 +2,4 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub mod benchmark;
-pub mod stress;
+pub mod workloads;

crates/sui-benchmark/src/stress/mod.rs → crates/sui-benchmark/src/workloads/mod.rs

@@ -1,6 +1,6 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-pub mod context;
 pub mod shared_counter;
 pub mod transfer_object;
+pub mod workload;

crates/sui-benchmark/src/stress/shared_counter.rs → crates/sui-benchmark/src/workloads/shared_counter.rs

@@ -1,7 +1,7 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::stress::context::{get_latest, transfer_sui_for_testing, MAX_GAS_FOR_TESTING};
+use crate::workloads::workload::{get_latest, transfer_sui_for_testing, MAX_GAS_FOR_TESTING};
 use async_trait::async_trait;
 use futures::future::join_all;
 use std::{path::PathBuf, sync::Arc};
@@ -19,7 +19,7 @@ use test_utils::{
     messages::create_publish_move_package_transaction, transaction::parse_package_ref,
 };
 
-use super::context::{submit_transaction, Gas, Payload, StressTestCtx};
+use super::workload::{submit_transaction, Gas, Payload, Workload, WorkloadType};
 
 pub struct SharedCounterTestPayload {
     package_ref: ObjectRef,
@@ -30,7 +30,7 @@ pub struct SharedCounterTestPayload {
 }
 
 impl Payload for SharedCounterTestPayload {
-    fn make_new_payload(&self, _: ObjectRef, new_gas: ObjectRef) -> Box<dyn Payload> {
+    fn make_new_payload(self: Box<Self>, _: ObjectRef, new_gas: ObjectRef) -> Box<dyn Payload> {
         Box::new(SharedCounterTestPayload {
             package_ref: self.package_ref,
             counter_id: self.counter_id,
@@ -51,27 +51,30 @@ impl Payload for SharedCounterTestPayload {
     fn get_object_id(&self) -> ObjectID {
         self.counter_id
     }
+    fn get_workload_type(&self) -> WorkloadType {
+        WorkloadType::SharedCounter
+    }
 }
 
-pub struct SharedCounterTestCtx {
+pub struct SharedCounterWorkload {
     pub test_gas: ObjectID,
     pub test_gas_owner: SuiAddress,
-    pub test_gas_keypair: AccountKeyPair,
-    pub num_counters: u64,
+    pub test_gas_keypair: Arc<AccountKeyPair>,
+    pub basics_package_ref: Option<ObjectRef>,
 }
 
-impl SharedCounterTestCtx {
-    pub fn make_ctx(
-        count: u64,
+impl SharedCounterWorkload {
+    pub fn new_boxed(
         gas: ObjectID,
         owner: SuiAddress,
-        keypair: AccountKeyPair,
-    ) -> Box<dyn StressTestCtx<dyn Payload>> {
-        Box::<dyn StressTestCtx<dyn Payload>>::from(Box::new(SharedCounterTestCtx {
+        keypair: Arc<AccountKeyPair>,
+        basics_package_ref: Option<ObjectRef>,
+    ) -> Box<dyn Workload<dyn Payload>> {
+        Box::<dyn Workload<dyn Payload>>::from(Box::new(SharedCounterWorkload {
             test_gas: gas,
             test_gas_owner: owner,
             test_gas_keypair: keypair,
-            num_counters: count,
+            basics_package_ref,
         }))
     }
 }
@@ -90,17 +93,51 @@ pub async fn publish_basics_package(
 }
 
 #[async_trait]
-impl StressTestCtx<dyn Payload> for SharedCounterTestCtx {
+impl Workload<dyn Payload> for SharedCounterWorkload {
+    async fn init(&mut self, aggregator: &AuthorityAggregator<NetworkAuthorityClient>) {
+        if self.basics_package_ref.is_some() {
+            return;
+        }
+        // publish basics package
+        let primary_gas = get_latest(self.test_gas, aggregator).await.unwrap();
+        let primary_gas_ref = primary_gas.compute_object_reference();
+        let mut publish_module_gas_ref = None;
+        let (address, keypair) = get_key_pair();
+        if let Some((_updated, minted)) = transfer_sui_for_testing(
+            (primary_gas_ref, Owner::AddressOwner(self.test_gas_owner)),
+            &self.test_gas_keypair,
+            MAX_GAS_FOR_TESTING,
+            address,
+            aggregator,
+        )
+        .await
+        {
+            publish_module_gas_ref = Some((address, keypair, minted));
+        }
+        // Publish basics package
+        eprintln!("Publishing basics package");
+        let publish_module_gas = publish_module_gas_ref.unwrap();
+        self.basics_package_ref = Some(
+            publish_basics_package(
+                publish_module_gas.2,
+                aggregator,
+                publish_module_gas.0,
+                &publish_module_gas.1,
+            )
+            .await,
+        )
+    }
     async fn make_test_payloads(
         &self,
+        count: u64,
         aggregator: &AuthorityAggregator<NetworkAuthorityClient>,
     ) -> Vec<Box<dyn Payload>> {
         // Read latest test gas object
         let primary_gas = get_latest(self.test_gas, aggregator).await.unwrap();
         let mut primary_gas_ref = primary_gas.compute_object_reference();
         // Make as many gas objects as the number of counters
         let mut counters_gas = vec![];
-        for _ in 0..self.num_counters {
+        for _ in 0..count {
             let (address, keypair) = get_key_pair();
             if let Some((updated, minted)) = transfer_sui_for_testing(
                 (primary_gas_ref, Owner::AddressOwner(self.test_gas_owner)),
@@ -115,40 +152,20 @@ impl StressTestCtx<dyn Payload> for SharedCounterTestCtx {
                 counters_gas.push((address, keypair, minted));
             }
         }
-        // Make gas for publishing package
-        let mut publish_module_gas_ref = None;
-        let (address, keypair) = get_key_pair();
-        if let Some((_updated, minted)) = transfer_sui_for_testing(
-            (primary_gas_ref, Owner::AddressOwner(self.test_gas_owner)),
-            &self.test_gas_keypair,
-            MAX_GAS_FOR_TESTING,
-            address,
-            aggregator,
-        )
-        .await
-        {
-            publish_module_gas_ref = Some((address, keypair, minted));
-        }
-        // Publish basics package
-        eprintln!("Publishing basics package");
-        let publish_module_gas = publish_module_gas_ref.unwrap();
-        let package_ref = publish_basics_package(
-            publish_module_gas.2,
-            aggregator,
-            publish_module_gas.0,
-            &publish_module_gas.1,
-        )
-        .await;
         // create counters using gas objects we created above
         eprintln!("Creating shared counters, this may take a while..");
         let futures = counters_gas
             .into_iter()
             .map(|(sender, keypair, gas)| async move {
-                let transaction =
-                    make_counter_create_transaction(gas, package_ref, sender, &keypair);
+                let transaction = make_counter_create_transaction(
+                    gas,
+                    self.basics_package_ref.unwrap(),
+                    sender,
+                    &keypair,
+                );
                 if let Some(effects) = submit_transaction(transaction, aggregator).await {
                     Box::new(SharedCounterTestPayload {
-                        package_ref,
+                        package_ref: self.basics_package_ref.unwrap(),
                         counter_id: effects.created[0].0 .0,
                         gas: effects.gas_object,
                         sender,

crates/sui-benchmark/src/stress/transfer_object.rs → crates/sui-benchmark/src/workloads/transfer_object.rs

@@ -17,8 +17,8 @@ use sui_types::{
 
 use test_utils::messages::make_transfer_object_transaction;
 
-use super::context::{
-    get_latest, transfer_sui_for_testing, Gas, Payload, StressTestCtx, MAX_GAS_FOR_TESTING,
+use super::workload::{
+    get_latest, transfer_sui_for_testing, Gas, Payload, Workload, WorkloadType, MAX_GAS_FOR_TESTING,
 };
 
 pub struct TransferObjectTestPayload {
@@ -30,7 +30,11 @@ pub struct TransferObjectTestPayload {
 }
 
 impl Payload for TransferObjectTestPayload {
-    fn make_new_payload(&self, new_object: ObjectRef, new_gas: ObjectRef) -> Box<dyn Payload> {
+    fn make_new_payload(
+        self: Box<Self>,
+        new_object: ObjectRef,
+        new_gas: ObjectRef,
+    ) -> Box<dyn Payload> {
         let updated_gas: Vec<Gas> = self
             .gas
             .iter()
@@ -72,43 +76,47 @@ impl Payload for TransferObjectTestPayload {
     fn get_object_id(&self) -> ObjectID {
         self.transfer_object.0
     }
+    fn get_workload_type(&self) -> WorkloadType {
+        WorkloadType::TransferObject
+    }
 }
 
-pub struct TransferObjectTestCtx {
+pub struct TransferObjectWorkload {
     pub test_gas: ObjectID,
     pub test_gas_owner: SuiAddress,
-    pub test_gas_keypair: AccountKeyPair,
-    pub num_transfer_objects: u64,
+    pub test_gas_keypair: Arc<AccountKeyPair>,
     pub num_accounts: u64,
     pub transfer_keypairs: Arc<HashMap<SuiAddress, AccountKeyPair>>,
 }
 
-impl TransferObjectTestCtx {
-    pub fn make_ctx(
-        count: u64,
+impl TransferObjectWorkload {
+    pub fn new_boxed(
         num_accounts: u64,
         gas: ObjectID,
         owner: SuiAddress,
-        keypair: AccountKeyPair,
-    ) -> Box<dyn StressTestCtx<dyn Payload>> {
+        keypair: Arc<AccountKeyPair>,
+    ) -> Box<dyn Workload<dyn Payload>> {
         // create several accounts to transfer object between
         let keypairs: Arc<HashMap<SuiAddress, AccountKeyPair>> =
             Arc::new((0..num_accounts).map(|_| get_key_pair()).collect());
-        Box::new(TransferObjectTestCtx {
+        Box::new(TransferObjectWorkload {
             test_gas: gas,
             test_gas_owner: owner,
             test_gas_keypair: keypair,
-            num_transfer_objects: count,
             num_accounts,
             transfer_keypairs: keypairs,
         })
     }
 }
 
 #[async_trait]
-impl StressTestCtx<dyn Payload> for TransferObjectTestCtx {
+impl Workload<dyn Payload> for TransferObjectWorkload {
+    async fn init(&mut self, _aggregator: &AuthorityAggregator<NetworkAuthorityClient>) {
+        return;
+    }
     async fn make_test_payloads(
         &self,
+        count: u64,
         aggregator: &AuthorityAggregator<NetworkAuthorityClient>,
     ) -> Vec<Box<dyn Payload>> {
         // Read latest test gas object
@@ -120,8 +128,9 @@ impl StressTestCtx<dyn Payload> for TransferObjectTestCtx {
             .choose(&mut rand::thread_rng())
             .unwrap();
         // create as many gas objects as there are number of transfer objects times number of accounts
+        eprintln!("Creating enough gas to transfer objects..");
         let mut transfer_gas: Vec<Vec<Gas>> = vec![];
-        for _i in 0..self.num_transfer_objects {
+        for _i in 0..count {
             let mut account_transfer_gas = vec![];
             for (owner, _) in self.transfer_keypairs.iter() {
                 if let Some((updated, minted)) = transfer_sui_for_testing(
@@ -139,9 +148,10 @@ impl StressTestCtx<dyn Payload> for TransferObjectTestCtx {
             }
             transfer_gas.push(account_transfer_gas);
         }
+        eprintln!("Creating objects to transfer..");
         // create transfer objects with 1 SUI value each
         let mut transfer_objects: Vec<Gas> = vec![];
-        for _i in 0..self.num_transfer_objects {
+        for _i in 0..count {
             if let Some((updated, minted)) = transfer_sui_for_testing(
                 (primary_gas_ref, Owner::AddressOwner(self.test_gas_owner)),
                 &self.test_gas_keypair,

crates/sui-benchmark/src/stress/context.rs → crates/sui-benchmark/src/workloads/workload.rs

@@ -22,7 +22,6 @@ use sui_types::{
         Transaction,
     },
 };
-
 use test_utils::messages::make_transfer_sui_transaction;
 use tracing::log::error;
 
@@ -32,12 +31,6 @@ pub const MAX_GAS_FOR_TESTING: u64 = 1_000_000_000;
 
 pub type Gas = (ObjectRef, Owner);
 
-pub trait Payload: Send + Sync {
-    fn make_new_payload(&self, new_object: ObjectRef, new_gas: ObjectRef) -> Box<dyn Payload>;
-    fn make_transaction(&self) -> TransactionEnvelope<EmptySignInfo>;
-    fn get_object_id(&self) -> ObjectID;
-}
-
 pub type UpdatedAndNewlyMinted = (ObjectRef, ObjectRef);
 
 pub async fn transfer_sui_for_testing(
@@ -119,10 +112,29 @@ pub async fn submit_transaction(
     }
 }
 
+pub trait Payload: Send + Sync {
+    fn make_new_payload(
+        self: Box<Self>,
+        new_object: ObjectRef,
+        new_gas: ObjectRef,
+    ) -> Box<dyn Payload>;
+    fn make_transaction(&self) -> TransactionEnvelope<EmptySignInfo>;
+    fn get_object_id(&self) -> ObjectID;
+    fn get_workload_type(&self) -> WorkloadType;
+}
+
+#[derive(Copy, Clone, Hash, PartialEq, Eq)]
+pub enum WorkloadType {
+    SharedCounter,
+    TransferObject,
+}
+
 #[async_trait]
-pub trait StressTestCtx<T: Payload + ?Sized>: Send + Sync {
+pub trait Workload<T: Payload + ?Sized>: Send + Sync {
+    async fn init(&mut self, aggregator: &AuthorityAggregator<NetworkAuthorityClient>);
     async fn make_test_payloads(
         &self,
+        count: u64,
         client: &AuthorityAggregator<NetworkAuthorityClient>,
     ) -> Vec<Box<T>>;
 }