iota_single_node_benchmark/
workload.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{path::PathBuf, sync::Arc};
6
7use iota_test_transaction_builder::PublishData;
8
9use crate::{
10    benchmark_context::BenchmarkContext,
11    command::WorkloadKind,
12    tx_generator::{MoveTxGenerator, PackagePublishTxGenerator, TxGenerator},
13};
14
15#[derive(Clone)]
16pub struct Workload {
17    pub tx_count: u64,
18    pub workload_kind: WorkloadKind,
19}
20
21impl Workload {
22    pub fn new(tx_count: u64, workload_kind: WorkloadKind) -> Self {
23        Self {
24            tx_count,
25            workload_kind,
26        }
27    }
28
29    pub(crate) fn num_accounts(&self) -> u64 {
30        self.tx_count
31    }
32
33    pub(crate) fn gas_object_num_per_account(&self) -> u64 {
34        self.workload_kind.gas_object_num_per_account()
35    }
36
37    pub(crate) async fn create_tx_generator(
38        &self,
39        ctx: &mut BenchmarkContext,
40    ) -> Arc<dyn TxGenerator> {
41        match &self.workload_kind {
42            WorkloadKind::PTB {
43                num_transfers,
44                use_native_transfer,
45                num_dynamic_fields,
46                computation,
47                num_shared_objects,
48                num_mints,
49                nft_size,
50                use_batch_mint,
51            } => {
52                let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
53                path.extend(["move_package"]);
54                let move_package = ctx.publish_package(PublishData::Source(path, false)).await;
55                let root_objects = ctx
56                    .preparing_dynamic_fields(move_package.0, *num_dynamic_fields)
57                    .await;
58                let shared_objects = ctx
59                    .prepare_shared_objects(move_package.0, *num_shared_objects)
60                    .await;
61                Arc::new(MoveTxGenerator::new(
62                    move_package.0,
63                    *num_transfers,
64                    *use_native_transfer,
65                    *computation,
66                    root_objects,
67                    shared_objects,
68                    *num_mints,
69                    *nft_size,
70                    *use_batch_mint,
71                ))
72            }
73            WorkloadKind::Publish {
74                manifest_file: manifest_path,
75            } => Arc::new(PackagePublishTxGenerator::new(ctx, manifest_path.clone()).await),
76        }
77    }
78}