sui_rpc_benchmark/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4pub mod config;
5pub mod direct;
6pub mod json_rpc;
7
8use std::collections::HashSet;
9use std::time::Duration;
10
11use anyhow::Result;
12use clap::{Parser, Subcommand, value_parser};
13use tracing::info;
14use url::Url;
15
16use crate::config::BenchmarkConfig;
17use crate::direct::query_enricher::QueryEnricher;
18use crate::direct::query_executor::QueryExecutor;
19use crate::direct::query_template_generator::QueryTemplateGenerator;
20
21#[derive(Parser)]
22#[clap(
23    name = "sui-rpc-benchmark",
24    about = "Benchmark tool for comparing Sui RPC access methods"
25)]
26pub struct Opts {
27    #[clap(subcommand)]
28    pub command: Command,
29}
30
31#[derive(Subcommand)]
32pub enum Command {
33    /// Benchmark direct database queries
34    #[clap(name = "direct")]
35    DirectQuery {
36        #[clap(
37            long,
38            default_value = "postgres://postgres:postgres@localhost:5432/sui",
39            value_parser = value_parser!(Url)
40        )]
41        db_url: Url,
42        #[clap(long, default_value = "50")]
43        concurrency: usize,
44        #[clap(long, default_value = "30")]
45        duration_secs: u64,
46    },
47    /// Benchmark JSON RPC endpoints
48    #[clap(name = "jsonrpc")]
49    JsonRpc {
50        #[clap(long, default_value = "http://127.0.0.1:9000")]
51        endpoint: String,
52        #[clap(long, default_value = "50")]
53        concurrency: usize,
54        #[clap(long)]
55        duration_secs: Option<u64>,
56        #[clap(long, default_value = "requests.jsonl")]
57        requests_file: String,
58        #[clap(long, value_delimiter = ',')]
59        methods_to_skip: Vec<String>,
60    },
61    /// Benchmark GraphQL queries
62    #[clap(name = "graphql")]
63    GraphQL {
64        #[clap(long, default_value = "http://127.0.0.1:9000/graphql")]
65        endpoint: String,
66    },
67}
68
69pub async fn run_benchmarks() -> Result<(), anyhow::Error> {
70    let opts: Opts = Opts::parse();
71
72    match opts.command {
73        Command::DirectQuery {
74            db_url,
75            concurrency,
76            duration_secs,
77        } => {
78            info!("Running direct query benchmark against DB {}", db_url);
79
80            let template_generator = QueryTemplateGenerator::new(db_url.clone());
81            let query_templates = template_generator.generate_query_templates().await?;
82            info!("Generated {} query templates", query_templates.len());
83
84            let query_enricher = QueryEnricher::new(&db_url).await?;
85            let enriched_queries = query_enricher.enrich_queries(query_templates).await?;
86            info!(
87                "Enriched {} queries with sample data",
88                enriched_queries.len()
89            );
90
91            let config = BenchmarkConfig {
92                concurrency,
93                duration: Some(Duration::from_secs(duration_secs)),
94                json_rpc_file_path: None,
95                json_rpc_methods_to_skip: HashSet::new(),
96            };
97            let query_executor = QueryExecutor::new(&db_url, enriched_queries, config).await?;
98            let result = query_executor.run().await?;
99
100            info!("Total queries: {}", result.total_queries);
101            info!("Total errors: {}", result.total_errors);
102            info!("Average latency: {:.2}ms", result.avg_latency_ms);
103            info!("Per-table statistics:");
104            for stat in &result.table_stats {
105                info!(
106                    "  {:<30} queries: {:<8} errors: {:<8} avg latency: {:.2}ms",
107                    stat.table_name, stat.queries, stat.errors, stat.avg_latency_ms
108                );
109            }
110            Ok(())
111        }
112        Command::JsonRpc {
113            endpoint,
114            concurrency,
115            duration_secs,
116            requests_file,
117            methods_to_skip,
118        } => {
119            info!(
120                concurrency,
121                ?duration_secs,
122                requests_file,
123                "Running JSON RPC benchmark against {endpoint}"
124            );
125            json_rpc::run_benchmark(
126                &endpoint,
127                &requests_file,
128                concurrency,
129                duration_secs,
130                methods_to_skip.into_iter().collect(),
131            )
132            .await?;
133            Ok(())
134        }
135        Command::GraphQL { endpoint } => {
136            info!("Running GraphQL benchmark against {}", endpoint);
137            todo!()
138        }
139    }
140}