MystenLabs · gegaowp · Jan 17, 2025 · Jan 8, 2025 · Jan 15, 2025 · Jan 17, 2025
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/sui-rpc-benchmark/Cargo.toml b/crates/sui-rpc-benchmark/Cargo.toml
@@ -12,6 +12,9 @@ clap = { workspace = true, features = ["derive"] }
 dashmap.workspace = true
 futures.workspace = true
 rand.workspace = true
+reqwest.workspace = true
+serde.workspace = true
+serde_json.workspace = true
 sui-indexer-alt-framework.workspace = true
 telemetry-subscribers.workspace = true
 tracing.workspace = true

diff --git a/crates/sui-rpc-benchmark/src/bin/pull_grafana_logs.rs b/crates/sui-rpc-benchmark/src/bin/pull_grafana_logs.rs
@@ -0,0 +1,129 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+/// This script pulls JSON RPC read logs from Grafana, extracts JSON bodies,
+/// and groups them by RPC "method" for later replay and analysis.
+use reqwest::header::{ACCEPT, AUTHORIZATION};
+use serde::Deserialize;
+use serde_json::Value;
+use std::collections::HashMap;
+use std::env;
+use std::error::Error;
+use std::fs::File;
+use std::io::{BufWriter, Write};
+use std::process;
+use tracing::{debug, error, info, warn};
+
+#[derive(Debug, Deserialize)]
+struct GrafanaLog {
+    message: String,
+}
+
+fn extract_body_from_message(message: &str) -> Option<String> {
+    if let Some(body_start) = message.find("body=") {
+        if let Some(peer_type_start) = message.find(" peer_type=") {
+            let raw_body = &message[(body_start + 5)..peer_type_start].trim();
+            if raw_body.starts_with('b') {
+                let trimmed = raw_body.trim_start_matches('b').trim_matches('"');
+                let unescaped = trimmed.replace("\\\"", "\"");
+                return Some(unescaped);
+            }
+        }
+    }
+    None
+}
+
+#[tokio::main]
+async fn main() {
+    let _guard = telemetry_subscribers::TelemetryConfig::new()
+        .with_env()
+        .init();
+    if let Err(e) = run().await {
+        error!("Error: {}", e);
+        process::exit(1);
+    }
+}
+
+async fn run() -> Result<(), Box<dyn Error>> {
+    let grafana_url = env::var("GRAFANA_LOGS_URL")
+        .unwrap_or_else(|_| "https://metrics.sui.io/loki/api/v1/query_range".to_string());
+    let grafana_token = env::var("GRAFANA_API_TOKEN").unwrap_or_else(|_| "".to_string());
+
+    let net = env::var("NET").unwrap_or_else(|_| "mainnet".to_string());
+    let namespace = if net == "testnet" {
+        "rpc-testnet".to_string()
+    } else if net == "mainnet" {
+        "rpc-mainnet".to_string()
+    } else {
+        "UNKNOWN_NET".to_string()
+    };
+    let substring = env::var("SUBSTRING").unwrap_or_else(|_| "Sampled read request".to_string());
+    let query = format!(
+        r#"{{namespace="{}", container="sui-edge-proxy-mysten"}} |= "{}""#,
+        namespace, substring
+    );
+    debug!("Query: {}", query);
+
+    let start = env::var("START").unwrap_or_else(|_| "now-1h".to_string());
+    let end = env::var("END").unwrap_or_else(|_| "now".to_string());
+
+    let client = reqwest::Client::new();
+    let mut query_params = vec![
+        ("query", query.as_str()),
+        ("start", start.as_str()),
+        ("end", end.as_str()),
+    ];
+    let limit = env::var("LIMIT").ok();
+    if let Some(ref l) = limit {
+        query_params.push(("limit", l));
+    }
+
+    let resp = client
+        .get(&grafana_url)
+        .header(ACCEPT, "application/json")
+        .header(AUTHORIZATION, format!("Bearer {}", grafana_token))
+        .query(&query_params)
+        .send()
+        .await?;
+
+    if !resp.status().is_success() {
+        warn!("Request failed with status: {}", resp.status());
+        return Ok(());
+    } else {
+        info!("Request succeeded with status: {}", resp.status());
+        debug!("Response: {:?}", resp);
+    }
+
+    let logs: Vec<GrafanaLog> = resp.json().await?;
+    info!("Found {} logs.", logs.len());
+
+    let mut method_map: HashMap<String, Vec<String>> = HashMap::new();
+    for log_entry in logs {
+        if let Some(body_content) = extract_body_from_message(&log_entry.message) {
+            if let Ok(parsed) = serde_json::from_str::<Value>(&body_content) {
+                let method = parsed
+                    .get("method")
+                    .and_then(|m| m.as_str())
+                    .unwrap_or("unknown_method")
+                    .to_string();
+                method_map.entry(method).or_default().push(body_content);
+            }
+        }
+    }
+
+    let file = File::create("sampled_read_requests.jsonl")?;
+    let mut writer = BufWriter::new(file);
+
+    for (method, bodies) in method_map {
+        info!("Writing {} logs for method: {}", bodies.len(), method);
+        for body in bodies {
+            let line = format!(r#"{{"method":"{}", "body":{}}}"#, method, body);
+            writer.write_all(line.as_bytes())?;
+            writer.write_all(b"\n")?;
+        }
+    }
+
+    writer.flush()?;
+    info!("Done! Wrote grouped logs to sampled_read_requests.jsonl");
+    Ok(())
+}
diff --git a/...-benchmark/src/direct/benchmark_config.rs → crates/sui-rpc-benchmark/src/config.rs b/...-benchmark/src/direct/benchmark_config.rs → crates/sui-rpc-benchmark/src/config.rs
@@ -3,18 +3,24 @@
 
 use std::time::Duration;
 
+#[derive(Debug, Clone)]
 pub struct BenchmarkConfig {
     /// Number of concurrent clients
     pub concurrency: usize,
     /// Duration to run the benchmark in seconds
     pub duration: Duration,
+    /// Optional path to a jsonl file for JSON RPC benchmarks.
+    /// The file contains a list of JSON RPC requests that are collected from Grafana,
+    /// and will be run concurrently by the JSON RPC benchmark runner.
+    pub json_rpc_file_path: Option<String>,
 }
 
 impl Default for BenchmarkConfig {
     fn default() -> Self {
         Self {
             concurrency: 50,
             duration: Duration::from_secs(30),
+            json_rpc_file_path: None,
         }
     }
 }
diff --git a/crates/sui-rpc-benchmark/src/direct/mod.rs b/crates/sui-rpc-benchmark/src/direct/mod.rs
@@ -1,7 +1,6 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-pub mod benchmark_config;
 pub mod metrics;
 pub mod query_executor;
 pub mod query_generator;
diff --git a/crates/sui-rpc-benchmark/src/direct/query_executor.rs b/crates/sui-rpc-benchmark/src/direct/query_executor.rs
@@ -11,7 +11,7 @@ use sui_indexer_alt_framework::task::TrySpawnStreamExt;
 use tokio_postgres::{types::ToSql, types::Type, NoTls, Row};
 use tracing::info;
 
-use crate::direct::benchmark_config::BenchmarkConfig;
+use crate::config::BenchmarkConfig;
 use crate::direct::metrics::{BenchmarkResult, MetricsCollector};
 use crate::direct::query_generator::BenchmarkQuery;
 

diff --git a/crates/sui-rpc-benchmark/src/json_rpc/mod.rs b/crates/sui-rpc-benchmark/src/json_rpc/mod.rs
@@ -0,0 +1,59 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::config::BenchmarkConfig;
+use anyhow::Result;
+use request_loader::load_json_rpc_requests;
+use runner::run_queries;
+use std::time::Duration;
+use tracing::info;
+
+pub mod request_loader;
+pub mod runner;
+
+pub async fn run_benchmark(
+    endpoint: &str,
+    file_path: &str,
+    concurrency: usize,
+    duration_secs: u64,
+) -> Result<()> {
+    let config = BenchmarkConfig {
+        concurrency,
+        duration: Duration::from_secs(duration_secs),
+        json_rpc_file_path: Some(file_path.to_string()),
+    };
+
+    info!("Loading JSON RPC requests from {}", file_path);
+    let requests = load_json_rpc_requests(file_path)?;
+    info!("Loaded {} requests", requests.len());
+
+    let metrics = run_queries(endpoint, &requests, &config).await?;
+    info!("Benchmark results:");
+    info!("=== Overall Statistics ===");
+    info!("Total requests sent: {}", metrics.total_sent);
+    info!("Total errors: {}", metrics.total_errors);
+    if metrics.total_sent > 0 {
+        let avg_latency = metrics.total_latency_ms / metrics.total_sent as f64;
+        info!("Average latency: {:.2}ms", avg_latency);
+        let success_rate = ((metrics.total_sent - metrics.total_errors) as f64
+            / metrics.total_sent as f64)
+            * 100.0;
+        info!("Success rate: {:.1}%", success_rate);
+    }
+    info!("=== Per-Method Statistics ===");
+    let mut methods: Vec<_> = metrics.per_method.iter().collect();
+    methods.sort_by_key(|(method, _)| *method);
+    for (method, stats) in methods {
+        info!("Method: {}", method);
+        info!("  Requests: {}", stats.total_sent);
+        info!("  Errors: {}", stats.total_errors);
+        if stats.total_sent > 0 {
+            let method_avg_latency = stats.total_latency_ms / stats.total_sent as f64;
+            let method_success_rate =
+                ((stats.total_sent - stats.total_errors) as f64 / stats.total_sent as f64) * 100.0;
+            info!("  Avg latency: {:.2}ms", method_avg_latency);
+            info!("  Success rate: {:.1}%", method_success_rate);
+        }
+    }
+    Ok(())
+}
diff --git a/crates/sui-rpc-benchmark/src/json_rpc/request_loader.rs b/crates/sui-rpc-benchmark/src/json_rpc/request_loader.rs
@@ -0,0 +1,34 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+/// This module implements the request loader, which is used to load
+/// the JSON RPC requests from a jsonl file.
+use anyhow::{Context, Result};
+use serde::Deserialize;
+use std::{
+    fs::File,
+    io::{BufRead, BufReader},
+};
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct JsonRpcRequestLine {
+    pub method: String,
+    #[serde(rename = "body")]
+    pub body_json: serde_json::Value,
+}
+
+pub fn load_json_rpc_requests(file_path: &str) -> Result<Vec<JsonRpcRequestLine>> {
+    let file = File::open(file_path)
+        .with_context(|| format!("Could not open JSON RPC file at {}", file_path))?;
+    let reader = BufReader::new(file);
+
+    let mut requests = Vec::new();
+    for line in reader.lines() {
+        let line = line?;
+        let request_line: JsonRpcRequestLine =
+            serde_json::from_str(&line).with_context(|| "Failed to parse JSON RPC line")?;
+        requests.push(request_line);
+    }
+
+    Ok(requests)
+}
diff --git a/crates/sui-rpc-benchmark/src/json_rpc/runner.rs b/crates/sui-rpc-benchmark/src/json_rpc/runner.rs
@@ -0,0 +1,104 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+/// This module implements the JSON RPC benchmark runner.
+/// The main function is `run_queries`, which runs the queries concurrently
+/// and records the overall and per-method stats.
+use anyhow::Result;
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+    time::Instant,
+};
+use sui_indexer_alt_framework::task::TrySpawnStreamExt;
+use tokio::time::timeout;
+
+use super::request_loader::JsonRpcRequestLine;
+use crate::config::BenchmarkConfig;
+
+#[derive(Clone, Default)]
+pub struct PerMethodStats {
+    pub total_sent: usize,
+    pub total_errors: usize,
+    // record total latency and calculate average latency later to avoid duplicate calculations
+    pub total_latency_ms: f64,
+}
+
+#[derive(Clone, Default)]
+pub struct JsonRpcStats {
+    pub total_sent: usize,
+    pub total_errors: usize,
+    // record total latency and calculate average latency to avoid duplicate calculations
+    pub total_latency_ms: f64,
+    pub per_method: HashMap<String, PerMethodStats>,
+}
+
+impl JsonRpcStats {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    fn record_request(&mut self, method: &str, latency_ms: f64, is_error: bool) {
+        self.total_sent += 1;
+        self.total_latency_ms += latency_ms;
+        if is_error {
+            self.total_errors += 1;
+        }
+
+        let method_stats = self.per_method.entry(method.to_string()).or_default();
+        method_stats.total_sent += 1;
+        method_stats.total_latency_ms += latency_ms;
+        if is_error {
+            method_stats.total_errors += 1;
+        }
+    }
+}
+
+pub async fn run_queries(
+    endpoint: &str,
+    requests: &[JsonRpcRequestLine],
+    config: &BenchmarkConfig,
+) -> Result<JsonRpcStats> {
+    let concurrency = config.concurrency;
+    let shared_stats = Arc::new(Mutex::new(JsonRpcStats::new()));
+    let client = reqwest::Client::new();
+    let endpoint = endpoint.to_owned();
+    let requests = requests.to_vec();
+    let stats = shared_stats.clone();
+
+    let stream = futures::stream::iter(requests.into_iter().map(move |request_line| {
+        let task_stats = stats.clone();
+        let client = client.clone();
+        let endpoint = endpoint.clone();
+        async move {
+            let now = Instant::now();
+            let res = timeout(
+                std::time::Duration::from_secs(10),
+                client.post(&endpoint).json(&request_line.body_json).send(),
+            )
+            .await;
+
+            let elapsed_ms = now.elapsed().as_millis() as f64;
+            let is_error = !matches!(res, Ok(Ok(ref resp)) if resp.status().is_success());
+
+            let mut stats = task_stats
+                .lock()
+                .map_err(|e| anyhow::anyhow!("Failed to acquire stats lock: {}", e))?;
+            stats.record_request(&request_line.method, elapsed_ms, is_error);
+            Ok::<(), anyhow::Error>(())
+        }
+    }));
+
+    timeout(
+        config.duration,
+        stream.try_for_each_spawned(concurrency, |fut| fut),
+    )
+    .await
+    .unwrap_or(Ok(()))?;
+
+    let final_stats = shared_stats
+        .lock()
+        .map_err(|e| anyhow::anyhow!("Failed to acquire stats lock for final results: {}", e))?
+        .clone();
+    Ok(final_stats)
+}