iota_rpc_loadgen/
main.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5mod load_test;
6mod payload;
7
8use std::{
9    error::Error,
10    path::PathBuf,
11    time::{Duration, SystemTime, UNIX_EPOCH},
12};
13
14use anyhow::Result;
15use clap::Parser;
16use iota_keys::keystore::{AccountKeystore, FileBasedKeystore, Keystore};
17use iota_types::crypto::{EncodeDecodeBase64, IotaKeyPair};
18use payload::AddressQueryType;
19use tracing::info;
20
21use crate::{
22    load_test::{LoadTest, LoadTestConfig},
23    payload::{
24        Command, RpcCommandProcessor, SignerInfo, load_addresses_from_file, load_digests_from_file,
25        load_objects_from_file,
26    },
27};
28
29#[derive(Parser)]
30#[command(
31    name = "IOTA RPC Load Generator",
32    version = "0.1",
33    about = "A load test application for IOTA RPC"
34)]
35struct Opts {
36    // TODO(chris): support running multiple commands at once
37    #[command(subcommand)]
38    pub command: ClapCommand,
39    #[arg(long, default_value_t = 1)]
40    pub num_threads: usize,
41    #[arg(long, default_value_t = true)]
42    pub cross_validate: bool,
43    #[arg(long, num_args(1..), default_value = "http://127.0.0.1:9000")]
44    pub urls: Vec<String>,
45    /// the path to log file directory
46    #[arg(long, default_value = "~/.iota/iota_config/logs")]
47    logs_directory: String,
48
49    #[arg(long, default_value = "~/.iota/loadgen/data")]
50    data_directory: String,
51}
52
53#[derive(Parser)]
54pub struct CommonOptions {
55    #[arg(short, long, default_value_t = 0)]
56    pub repeat: usize,
57
58    #[arg(short, long, default_value_t = 0)]
59    pub interval_in_ms: u64,
60
61    /// different chunks will be executed concurrently on the same thread
62    #[arg(long, default_value_t = 1)]
63    num_chunks_per_thread: usize,
64}
65
66#[derive(Parser)]
67pub enum ClapCommand {
68    DryRun {
69        #[command(flatten)]
70        common: CommonOptions,
71    },
72    GetCheckpoints {
73        /// Default to start from checkpoint 0
74        #[arg(short, long, default_value_t = 0)]
75        start: u64,
76
77        /// inclusive, uses `getLatestCheckpointSequenceNumber` if `None`
78        #[arg(short, long)]
79        end: Option<u64>,
80
81        #[arg(long)]
82        skip_verify_transactions: bool,
83
84        #[arg(long)]
85        skip_verify_objects: bool,
86
87        // Whether to record data from checkpoint
88        #[arg(long)]
89        skip_record: bool,
90
91        #[command(flatten)]
92        common: CommonOptions,
93    },
94    PayIota {
95        // TODO(chris) customize recipients and amounts
96        #[command(flatten)]
97        common: CommonOptions,
98    },
99    QueryTransactionBlocks {
100        #[arg(long, ignore_case = true)]
101        address_type: AddressQueryType,
102
103        #[command(flatten)]
104        common: CommonOptions,
105    },
106    MultiGetTransactionBlocks {
107        #[command(flatten)]
108        common: CommonOptions,
109    },
110    MultiGetObjects {
111        #[command(flatten)]
112        common: CommonOptions,
113    },
114    GetObject {
115        #[arg(long)]
116        chunk_size: usize,
117
118        #[command(flatten)]
119        common: CommonOptions,
120    },
121    GetAllBalances {
122        #[arg(long)]
123        chunk_size: usize,
124
125        #[command(flatten)]
126        common: CommonOptions,
127    },
128    GetReferenceGasPrice {
129        #[command(flatten)]
130        common: CommonOptions,
131    },
132}
133
134fn get_keypair() -> Result<SignerInfo> {
135    // TODO(chris) allow pass in custom path for keystore
136    // Load keystore from ~/.iota/iota_config/iota.keystore
137    let keystore_path = get_iota_config_directory().join("iota.keystore");
138    let keystore = Keystore::from(FileBasedKeystore::new(&keystore_path)?);
139    let active_address = keystore.addresses().pop().unwrap();
140    let keypair: &IotaKeyPair = keystore.get_key(&active_address)?;
141    println!("using address {active_address} for signing");
142    Ok(SignerInfo::new(keypair.encode_base64()))
143}
144
145fn get_iota_config_directory() -> PathBuf {
146    match dirs::home_dir() {
147        Some(v) => v.join(".iota").join("iota_config"),
148        None => panic!("Cannot obtain home directory path"),
149    }
150}
151
152pub fn expand_path(dir_path: &str) -> String {
153    shellexpand::full(&dir_path)
154        .map(|v| v.into_owned())
155        .unwrap_or_else(|e| panic!("Failed to expand directory '{:?}': {}", dir_path, e))
156}
157
158fn get_log_file_path(dir_path: String) -> String {
159    let current_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
160    let timestamp = current_time.as_secs();
161    // use timestamp to signify which file is newer
162    let log_filename = format!("iota-rpc-loadgen.{}.log", timestamp);
163
164    let dir_path = expand_path(&dir_path);
165    format!("{dir_path}/{log_filename}")
166}
167
168#[tokio::main]
169async fn main() -> Result<(), Box<dyn Error>> {
170    let tracing_level = "debug";
171    let network_tracing_level = "info";
172    let log_filter = format!(
173        "{tracing_level},h2={network_tracing_level},tower={network_tracing_level},hyper={network_tracing_level},tonic::transport={network_tracing_level}"
174    );
175    let opts = Opts::parse();
176
177    let log_filename = get_log_file_path(opts.logs_directory);
178
179    // Initialize logger
180    let (_guard, _filter_handle) = telemetry_subscribers::TelemetryConfig::new()
181        .with_env()
182        .with_log_level(&log_filter)
183        .with_log_file(&log_filename)
184        .init();
185
186    println!("Logging to {}", &log_filename);
187    info!("Running Load Gen with following urls {:?}", opts.urls);
188
189    let (command, common, need_keystore) = match opts.command {
190        ClapCommand::DryRun { common } => (Command::new_dry_run(), common, false),
191        ClapCommand::PayIota { common } => (Command::new_pay_iota(), common, true),
192        ClapCommand::GetCheckpoints {
193            common,
194            start,
195            end,
196            skip_verify_transactions,
197            skip_verify_objects,
198            skip_record,
199        } => (
200            Command::new_get_checkpoints(
201                start,
202                end,
203                !skip_verify_transactions,
204                !skip_verify_objects,
205                !skip_record,
206            ),
207            common,
208            false,
209        ),
210        ClapCommand::QueryTransactionBlocks {
211            common,
212            address_type,
213        } => {
214            let addresses = load_addresses_from_file(expand_path(&opts.data_directory));
215            (
216                Command::new_query_transaction_blocks(address_type, addresses),
217                common,
218                false,
219            )
220        }
221        ClapCommand::MultiGetTransactionBlocks { common } => {
222            let digests = load_digests_from_file(expand_path(&opts.data_directory));
223            (
224                Command::new_multi_get_transaction_blocks(digests),
225                common,
226                false,
227            )
228        }
229        ClapCommand::GetAllBalances { common, chunk_size } => {
230            let addresses = load_addresses_from_file(expand_path(&opts.data_directory));
231            (
232                Command::new_get_all_balances(addresses, chunk_size),
233                common,
234                false,
235            )
236        }
237        ClapCommand::MultiGetObjects { common } => {
238            let objects = load_objects_from_file(expand_path(&opts.data_directory));
239            (Command::new_multi_get_objects(objects), common, false)
240        }
241        ClapCommand::GetReferenceGasPrice { common } => {
242            let num_repeats = common.num_chunks_per_thread;
243            (
244                Command::new_get_reference_gas_price(num_repeats),
245                common,
246                false,
247            )
248        }
249        ClapCommand::GetObject { common, chunk_size } => {
250            let objects = load_objects_from_file(expand_path(&opts.data_directory));
251            (Command::new_get_object(objects, chunk_size), common, false)
252        }
253    };
254
255    let signer_info = need_keystore.then_some(get_keypair()?);
256
257    let command = command
258        .with_repeat_interval(Duration::from_millis(common.interval_in_ms))
259        .with_repeat_n_times(common.repeat);
260
261    let processor = RpcCommandProcessor::new(&opts.urls, expand_path(&opts.data_directory)).await;
262
263    let load_test = LoadTest {
264        processor,
265        config: LoadTestConfig {
266            command,
267            num_threads: opts.num_threads,
268            // TODO: pass in from config
269            divide_tasks: true,
270            signer_info,
271            num_chunks_per_thread: common.num_chunks_per_thread,
272            max_repeat: common.repeat,
273        },
274    };
275    load_test.run().await?;
276
277    Ok(())
278}
iota_rpc_loadgen/main.rs

iota_rpc_loadgen/
main.rs
