sui_rpc_loadgen/payload/
mod.rs

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

sui_rpc_loadgen/payload/
mod.rs
