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

iota_rpc_loadgen/payload/
mod.rs
