Initial attempt to benchmark proof generation using the prover (#16)

* Initial attempt to benchmark proof generation using the prover

prover/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "status_rln_prover"
+name = "prover"
 version = "0.1.0"
 edition = "2024"
 
@@ -13,7 +13,6 @@ prost = "0.13"
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 tracing = "0.1.41"
-tracing-test = "0.2.5"
 alloy.workspace = true
 thiserror = "2.0"
 futures = "0.3"
@@ -50,7 +49,8 @@ tonic-build = "*"
 criterion.workspace = true
 ark-groth16.workspace = true
 tempfile = "3.20"
+tracing-test = "0.2.5"
 
-# [[bench]]
-# name = "user_db_heavy_write"
-# harness = false
+[[bench]]
+name = "prover_bench"
+harness = false

prover/benches/prover_bench.rs

@@ -0,0 +1,214 @@
+use criterion::BenchmarkId;
+use criterion::Criterion;
+use criterion::{criterion_group, criterion_main};
+
+// std
+use std::net::{IpAddr, Ipv4Addr};
+use std::sync::Arc;
+use std::time::Duration;
+use std::str::FromStr;
+// third-party
+use alloy::{
+    primitives::{Address, U256},
+};
+use parking_lot::RwLock;
+use tokio::sync::Notify;
+use tokio::task::JoinSet;
+use tonic::Response;
+use futures::FutureExt;
+// internal
+use prover::{
+    AppArgs,
+    run_prover
+};
+
+// grpc
+pub mod prover_proto {
+    // Include generated code (see build.rs)
+    tonic::include_proto!("prover");
+}
+use prover_proto::{
+    Address as GrpcAddress,
+    U256 as GrpcU256,
+    Wei as GrpcWei,
+    RegisterUserReply, RegisterUserRequest, RegistrationStatus,
+    SendTransactionRequest, SendTransactionReply,
+    RlnProofFilter, RlnProofReply,
+    rln_prover_client::RlnProverClient
+};
+
+async fn register_users(port: u16, addresses: Vec<Address>) {
+
+    let url = format!("http://127.0.0.1:{}", port);
+    let mut client = RlnProverClient::connect(url).await.unwrap();
+
+    for address in addresses {
+
+        let addr = GrpcAddress {
+            value: address.to_vec(),
+        };
+
+        let request_0 = RegisterUserRequest {
+            user: Some(addr),
+        };
+        let request = tonic::Request::new(request_0);
+        let response: Response<RegisterUserReply> = client.register_user(request).await.unwrap();
+
+        assert_eq!(
+            RegistrationStatus::try_from(response.into_inner().status).unwrap(),
+            RegistrationStatus::Success);
+    }
+}
+
+async fn proof_sender(port: u16, addresses: Vec<Address>, proof_count: usize) {
+
+    let chain_id = GrpcU256 {
+        // FIXME: LE or BE?
+        value: U256::from(1).to_le_bytes::<32>().to_vec(),
+    };
+
+    let url = format!("http://127.0.0.1:{}", port);
+    let mut client = RlnProverClient::connect(url).await.unwrap();
+
+    let addr = GrpcAddress {
+        value: addresses[0].to_vec(),
+    };
+    let wei = GrpcWei {
+        // FIXME: LE or BE?
+        value: U256::from(1000).to_le_bytes::<32>().to_vec(),
+    };
+
+    for i in 0..proof_count {
+        let tx_hash = U256::from(42+i).to_le_bytes::<32>().to_vec();
+
+        let request_0 = SendTransactionRequest {
+            gas_price: Some(wei.clone()),
+            sender: Some(addr.clone()),
+            chain_id: Some(chain_id.clone()),
+            transaction_hash: tx_hash,
+        };
+
+        let request = tonic::Request::new(request_0);
+        let response: Response<SendTransactionReply> = client.send_transaction(request).await.unwrap();
+        assert_eq!(response.into_inner().result, true);
+    }
+}
+
+async fn proof_collector(port: u16, proof_count: usize) -> Vec<RlnProofReply> {
+
+    let result= Arc::new(RwLock::new(vec![]));
+
+    let url = format!("http://127.0.0.1:{}", port);
+    let mut client = RlnProverClient::connect(url).await.unwrap();
+
+    let request_0 = RlnProofFilter {
+        address: None,
+    };
+
+    let request = tonic::Request::new(request_0);
+    let stream_ = client.get_proofs(request).await.unwrap();
+    let mut stream = stream_.into_inner();
+    let result_2 = result.clone();
+    let mut proof_received = 0;
+    while let Some(response) = stream.message().await.unwrap() {
+        result_2.write().push(response);
+        proof_received += 1;
+        if proof_received >= proof_count {
+            break;
+        }
+    }
+
+    let res = std::mem::take(&mut *result.write());
+
+    // println!("[Proof collector] Received {} proofs", res.len());
+    res
+}
+
+fn proof_generation_bench(c: &mut Criterion) {
+
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+
+    let port = 50051;
+    let temp_folder = tempfile::tempdir().unwrap();
+    let temp_folder_tree = tempfile::tempdir().unwrap();
+    let app_args = AppArgs {
+        ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
+        port,
+        ws_rpc_url: None,
+        db_path: temp_folder.path().to_path_buf(),
+        merkle_tree_path: temp_folder_tree.path().to_path_buf(),
+        ksc_address: None,
+        rlnsc_address: None,
+        tsc_address: None,
+        mock_sc: Some(true),
+        mock_user: None,
+        config_path: Default::default(),
+        no_config: Some(true),
+        broadcast_channel_size: 100,
+        proof_service_count: 32,
+        transaction_channel_size: 100,
+        proof_sender_channel_size: 100,
+    };
+
+    let addresses = vec![
+        Address::from_str("0xd8da6bf26964af9d7eed9e03e53415d37aa96045").unwrap(),
+        Address::from_str("0xb20a608c624Ca5003905aA834De7156C68b2E1d0").unwrap(),
+    ];
+    // Tokio notify - wait for some time after spawning run_prover then notify it's ready to accept
+    // connections
+    let notify_start = Arc::new(Notify::new());
+
+    // Spawn prover
+    let notify_start_1 = notify_start.clone();
+    rt.spawn(
+        async move {
+            tokio::spawn(run_prover(app_args));
+            tokio::time::sleep(Duration::from_secs(10)).await;
+            println!("Prover is ready, notifying it...");
+            notify_start_1.clone().notify_one();
+        }
+    );
+
+    let notify_start_2 = notify_start.clone();
+    let addresses_0 = addresses.clone();
+
+    // Wait for proof_collector to be connected and waiting for some proofs
+    let _res = rt.block_on(
+        async move {
+            notify_start_2.notified().await;
+            println!("Prover is ready, registering users...");
+            register_users(port, addresses_0).await;
+        }
+    );
+
+    println!("Starting benchmark...");
+    let size: usize = 1024;
+    let proof_count = 100;
+    c.bench_with_input(BenchmarkId::new("input_example", size), &size, |b, &_s| {
+        b.to_async(&rt).iter(|| {
+            async {
+                let mut set = JoinSet::new();
+                set.spawn(proof_collector(port, proof_count));
+                set.spawn(
+                    proof_sender(port, addresses.clone(), proof_count).map(|_r| vec![])
+                );
+                // Wait for proof_sender + proof_collector to complete
+                let res = set.join_all().await;
+                // Check we receive enough proof
+                assert_eq!(res[1].len(), proof_count);
+            }
+        });
+    });
+}
+
+criterion_group!(
+    name = benches;
+    config = Criterion::default()
+        .sample_size(10)
+        .measurement_time(Duration::from_secs(150));
+    targets = proof_generation_bench
+);
+criterion_main!(benches);

prover/src/args.rs

@@ -33,48 +33,48 @@ const ARGS_DEFAULT_PROOF_SENDER_CHANNEL_SIZE: &str = "100";
 #[command(about = "RLN prover service", long_about = None)]
 pub struct AppArgs {
     #[arg(short = 'i', long = "ip", default_value = "::1", help = "Service ip")]
-    pub(crate) ip: IpAddr,
+    pub ip: IpAddr,
     #[arg(
         short = 'p',
         long = "port",
         default_value = "50051",
         help = "Service port"
     )]
-    pub(crate) port: u16,
+    pub port: u16,
     #[arg(
         short = 'u',
         long = "ws-rpc-url",
         help = "Websocket rpc url (e.g. wss://eth-mainnet.g.alchemy.com/v2/your-api-key)"
     )]
-    pub(crate) ws_rpc_url: Option<Url>,
+    pub ws_rpc_url: Option<Url>,
     #[arg(long = "db", help = "Db path", default_value = "./storage/db")]
-    pub(crate) db_path: PathBuf,
+    pub db_path: PathBuf,
     #[arg(
         long = "tree",
         help = "Merkle tree path",
         default_value = "./storage/tree"
     )]
-    pub(crate) merkle_tree_path: PathBuf,
+    pub merkle_tree_path: PathBuf,
     #[arg(short = 'k', long = "ksc", help = "Karma smart contract address")]
-    pub(crate) ksc_address: Option<Address>,
+    pub ksc_address: Option<Address>,
     #[arg(short = 'r', long = "rlnsc", help = "RLN smart contract address")]
-    pub(crate) rlnsc_address: Option<Address>,
+    pub rlnsc_address: Option<Address>,
     #[arg(short = 't', long = "tsc", help = "KarmaTiers smart contract address")]
-    pub(crate) tsc_address: Option<Address>,
+    pub tsc_address: Option<Address>,
     #[arg(
         help_heading = "mock",
         long = "mock-sc",
         help = "Test only - mock smart contracts",
         action
     )]
-    pub(crate) mock_sc: Option<bool>,
+    pub mock_sc: Option<bool>,
     #[arg(
         help_heading = "mock",
         long = "mock-user",
         help = "Test only - register user (requite --mock-sc to be enabled)",
         action
     )]
-    pub(crate) mock_user: Option<PathBuf>,
+    pub mock_user: Option<PathBuf>,
     #[arg(
         short = 'c',
         long = "config",
@@ -82,7 +82,7 @@ pub struct AppArgs {
         default_value = "./config.toml",
         help_heading = "config"
     )]
-    pub(crate) config_path: PathBuf,
+    pub config_path: PathBuf,
     #[arg(
         long = "no-config",
         help = "Dont read a config file",
@@ -90,7 +90,7 @@ pub struct AppArgs {
         action = SetTrue,
         help_heading = "config"
     )]
-    pub(crate) no_config: Option<bool>,
+    pub no_config: Option<bool>,
     // Hidden option - expect user set it via a config file
     #[arg(
         long = "broadcast-channel-size",
@@ -98,28 +98,28 @@ pub struct AppArgs {
         default_value = ARGS_DEFAULT_BROADCAST_CHANNEL_SIZE,
         hide = true,
     )] // see const doc for more info
-    pub(crate) broadcast_channel_size: usize,
+    pub broadcast_channel_size: usize,
     #[arg(
         long = "proof-service",
         help = "Number of proof service (tasks) to generate proof",
         default_value = ARGS_DEFAULT_PROOF_SERVICE_COUNT,
         hide = true,
     )] // see const doc for more info
-    pub(crate) proof_service_count: u16,
+    pub proof_service_count: u16,
     #[arg(
         long = "transaction-channel-size",
         help = "Proof bounded channel size",
         default_value = ARGS_DEFAULT_TRANSACTION_CHANNEL_SIZE,
         hide = true,
     )] // see const doc for more info
-    pub(crate) transaction_channel_size: usize,
+    pub transaction_channel_size: usize,
     #[arg(
         long = "proof-sender-channel-size",
         help = "Proof bounded sender channel size",
         default_value = ARGS_DEFAULT_PROOF_SENDER_CHANNEL_SIZE,
         hide = true,
     )] // see const doc for more info
-    pub(crate) proof_sender_channel_size: usize,
+    pub proof_sender_channel_size: usize,
 }
 
 #[cfg(test)]

prover/src/lib.rs

@@ -18,25 +18,21 @@ mod user_db_types;
 
 // std
 use std::net::SocketAddr;
-use std::path::PathBuf;
 use std::time::Duration;
 // third-party
 use alloy::primitives::U256;
 use chrono::{DateTime, Utc};
-use clap::CommandFactory;
 use tokio::task::JoinSet;
-use tracing::level_filters::LevelFilter;
 use tracing::{
     debug,
     // error,
     // info
 };
-use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};
 // internal
 use rln_proof::RlnIdentifier;
 use smart_contract::KarmaTiersSC::KarmaTiersSCInstance;
 use smart_contract::TIER_LIMITS;
-use crate::args::{AppArgs, AppArgsConfig};
+pub use crate::args::{AppArgs, AppArgsConfig};
 use crate::epoch_service::EpochService;
 use crate::grpc_service::GrpcProverService;
 use crate::mock::read_mock_user;
@@ -53,54 +49,7 @@ const GENESIS: DateTime<Utc> = DateTime::from_timestamp(1431648000, 0).unwrap();
 const PROVER_MINIMAL_AMOUNT_FOR_REGISTRATION: U256 =
     U256::from_le_slice(10u64.to_le_bytes().as_slice());
 
-#[tokio::main]
-async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
-    // debug!("Args: {:?}", std::env::args());
-
-    let filter = EnvFilter::builder()
-        .with_default_directive(LevelFilter::INFO.into())
-        .from_env_lossy();
-    tracing_subscriber::registry()
-        .with(fmt::layer())
-        .with(filter)
-        .init();
-
-    // let app_args = AppArgs::parse();
-    let app_args = <AppArgs as CommandFactory>::command().get_matches();
-    debug!("Arguments: {:?}", app_args);
-
-    let app_ars_config = if !app_args.get_flag("no_config") {
-        // Unwrap safe - default value provided
-        let config_path = app_args.get_one::<PathBuf>("config_path").unwrap();
-        debug!("Reading config path: {:?}...", config_path);
-        let config_str = std::fs::read_to_string(config_path)?;
-        let config: AppArgsConfig = toml::from_str(config_str.as_str())?;
-        debug!("Config: {:?}", config);
-        config
-    } else {
-        AppArgsConfig::default()
-    };
-
-    // Merge command line args & config
-    let app_args = AppArgs::from_merged(app_args, Some(app_ars_config));
-    debug!("Arguments (merged with config): {:?}", app_args);
-
-    // Application cli arguments checks
-    if app_args.ws_rpc_url.is_some() {
-        if app_args.ksc_address.is_none()
-            || app_args.ksc_address.is_none()
-            || app_args.tsc_address.is_none()
-        {
-            return Err("Please provide smart contract addresses".into());
-        }
-    } else if app_args.mock_sc.is_none() {
-        return Err("Please provide rpc url (--ws-rpc-url) or mock (--mock-sc)".into());
-    }
-
-    run_prover(app_args).await
-}
-
-async fn run_prover(app_args: AppArgs) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+pub async fn run_prover(app_args: AppArgs) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
 
     // Epoch
     let epoch_service = EpochService::try_from((Duration::from_secs(60 * 2), GENESIS))

prover/src/tiers_listener.rs

@@ -70,7 +70,7 @@ impl TiersListener {
                         };
                     }
                     Err(e) => {
-                        eprintln!("Error decoding log data: {:?}", e);
+                        eprintln!("Error decoding log data: {e:?}");
                         // It's also useful to print the raw log data for debugging
                         eprintln!("Raw log topics: {:?}", log.topics());
                         eprintln!("Raw log data: {:?}", log.data());