iota_rpc_loadgen/payload/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5mod checkpoint_utils;
6mod get_all_balances;
7mod get_checkpoints;
8mod get_object;
9mod get_reference_gas_price;
10mod multi_get_objects;
11mod multi_get_transaction_blocks;
12mod pay_iota;
13mod query_transactions;
14mod rpc_command_processor;
15mod validation;
16
17use core::default::Default;
18use std::time::Duration;
19
20use anyhow::Result;
21use async_trait::async_trait;
22use iota_sdk_types::{Address, ObjectId, TransactionDigest};
23use iota_types::messages_checkpoint::CheckpointSequenceNumber;
24pub use rpc_command_processor::{
25    RpcCommandProcessor, load_addresses_from_file, load_digests_from_file, load_objects_from_file,
26};
27use strum::EnumString;
28
29use crate::load_test::LoadTestConfig;
30
31#[derive(Default, Clone)]
32pub struct SignerInfo {
33    pub encoded_keypair: String,
34    /// Different thread should use different gas_payment to avoid equivocation
35    pub gas_payment: Option<Vec<ObjectId>>,
36    pub gas_budget: Option<u64>,
37}
38
39impl SignerInfo {
40    pub fn new(encoded_keypair: String) -> Self {
41        Self {
42            encoded_keypair,
43            gas_payment: None,
44            gas_budget: None,
45        }
46    }
47}
48
49#[derive(Clone, Default)]
50pub struct Payload {
51    pub commands: Vec<Command>,
52    pub signer_info: Option<SignerInfo>,
53}
54
55#[derive(Default, Clone)]
56pub struct Command {
57    pub data: CommandData,
58    /// 0 means the command will be run once. Default to be 0
59    pub repeat_n_times: usize,
60    /// how long to wait between the start of two subsequent repeats
61    /// If the previous command takes longer than `repeat_interval` to finish,
62    /// the next command will run as soon as the previous command finishes
63    /// Default to be 0
64    pub repeat_interval: Duration,
65}
66
67impl Command {
68    pub fn new_dry_run() -> Self {
69        Self {
70            data: CommandData::DryRun(DryRun {}),
71            ..Default::default()
72        }
73    }
74
75    pub fn new_pay_iota() -> Self {
76        Self {
77            data: CommandData::PayIota(PayIota {}),
78            ..Default::default()
79        }
80    }
81
82    pub fn new_get_checkpoints(
83        start: CheckpointSequenceNumber,
84        end: Option<CheckpointSequenceNumber>,
85        verify_transactions: bool,
86        verify_objects: bool,
87        record: bool,
88    ) -> Self {
89        Self {
90            data: CommandData::GetCheckpoints(GetCheckpoints {
91                start,
92                end,
93                verify_transactions,
94                verify_objects,
95                record,
96            }),
97            ..Default::default()
98        }
99    }
100
101    pub fn new_query_transaction_blocks(
102        address_type: AddressQueryType,
103        addresses: Vec<Address>,
104    ) -> Self {
105        let query_transactions = QueryTransactionBlocks {
106            address_type,
107            addresses,
108        };
109        Self {
110            data: CommandData::QueryTransactionBlocks(query_transactions),
111            ..Default::default()
112        }
113    }
114
115    pub fn new_multi_get_transaction_blocks(digests: Vec<TransactionDigest>) -> Self {
116        let multi_get_transaction_blocks = MultiGetTransactionBlocks { digests };
117        Self {
118            data: CommandData::MultiGetTransactionBlocks(multi_get_transaction_blocks),
119            ..Default::default()
120        }
121    }
122
123    pub fn new_multi_get_objects(object_ids: Vec<ObjectId>) -> Self {
124        let multi_get_objects = MultiGetObjects { object_ids };
125        Self {
126            data: CommandData::MultiGetObjects(multi_get_objects),
127            ..Default::default()
128        }
129    }
130
131    pub fn new_get_object(object_ids: Vec<ObjectId>, chunk_size: usize) -> Self {
132        let get_object = GetObject {
133            object_ids,
134            chunk_size,
135        };
136        Self {
137            data: CommandData::GetObject(get_object),
138            ..Default::default()
139        }
140    }
141
142    pub fn new_get_all_balances(addresses: Vec<Address>, chunk_size: usize) -> Self {
143        let get_all_balances = GetAllBalances {
144            addresses,
145            chunk_size,
146        };
147        Self {
148            data: CommandData::GetAllBalances(get_all_balances),
149            ..Default::default()
150        }
151    }
152
153    pub fn new_get_reference_gas_price(num_repeats: usize) -> Self {
154        let get_reference_gas_price = GetReferenceGasPrice { num_repeats };
155        Self {
156            data: CommandData::GetReferenceGasPrice(get_reference_gas_price),
157            ..Default::default()
158        }
159    }
160
161    pub fn with_repeat_n_times(mut self, num: usize) -> Self {
162        self.repeat_n_times = num;
163        self
164    }
165
166    pub fn with_repeat_interval(mut self, duration: Duration) -> Self {
167        self.repeat_interval = duration;
168        self
169    }
170}
171
172#[derive(Clone)]
173pub enum CommandData {
174    DryRun(DryRun),
175    GetCheckpoints(GetCheckpoints),
176    PayIota(PayIota),
177    QueryTransactionBlocks(QueryTransactionBlocks),
178    MultiGetTransactionBlocks(MultiGetTransactionBlocks),
179    MultiGetObjects(MultiGetObjects),
180    GetObject(GetObject),
181    GetAllBalances(GetAllBalances),
182    GetReferenceGasPrice(GetReferenceGasPrice),
183}
184
185impl Default for CommandData {
186    fn default() -> Self {
187        CommandData::DryRun(DryRun {})
188    }
189}
190
191#[derive(Clone)]
192pub struct DryRun {}
193
194#[derive(Clone, Default)]
195pub struct GetCheckpoints {
196    /// Default to start from 0
197    pub start: CheckpointSequenceNumber,
198    /// If None, use `getLatestCheckpointSequenceNumber`
199    pub end: Option<CheckpointSequenceNumber>,
200    pub verify_transactions: bool,
201    pub verify_objects: bool,
202    pub record: bool,
203}
204
205#[derive(Clone)]
206pub struct PayIota {}
207
208#[derive(Clone, Default)]
209pub struct QueryTransactionBlocks {
210    pub address_type: AddressQueryType,
211    pub addresses: Vec<Address>,
212}
213
214#[derive(Clone)]
215pub struct MultiGetTransactionBlocks {
216    pub digests: Vec<TransactionDigest>,
217}
218
219#[derive(Clone, EnumString, Default)]
220#[strum(serialize_all = "lowercase")]
221pub enum AddressQueryType {
222    #[default]
223    From,
224    To,
225    Both,
226}
227
228#[derive(Clone)]
229pub struct MultiGetObjects {
230    pub object_ids: Vec<ObjectId>,
231}
232
233#[derive(Clone)]
234pub struct GetObject {
235    pub object_ids: Vec<ObjectId>,
236    pub chunk_size: usize,
237}
238
239#[derive(Clone)]
240pub struct GetAllBalances {
241    pub addresses: Vec<Address>,
242    pub chunk_size: usize,
243}
244
245#[derive(Clone)]
246pub struct GetReferenceGasPrice {
247    num_repeats: usize,
248}
249
250#[async_trait]
251pub trait Processor {
252    /// process commands in order
253    async fn apply(&self, payload: &Payload) -> Result<()>;
254
255    /// prepare payload for each thread according to LoadTestConfig
256    async fn prepare(&self, config: &LoadTestConfig) -> Result<Vec<Payload>>;
257
258    /// write results to file based on LoadTestConfig
259    fn dump_cache_to_file(&self, config: &LoadTestConfig);
260}
261
262/// all payload should implement this trait
263#[async_trait]
264pub trait ProcessPayload<'a, T> {
265    async fn process(&'a self, op: T, signer_info: &Option<SignerInfo>) -> Result<()>;
266}
iota_rpc_loadgen/payload/mod.rs

iota_rpc_loadgen/payload/
mod.rs
