Skip to content

feat: support recover in faucet #18

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
+178 −215
Closed

feat: support recover in faucet #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6a30a3b
recover
keanji-x Nov 10, 2025
62fc170
recover
keanji-x Nov 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
230 changes: 88 additions & 142 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,14 @@ use alloy::{
};
use anyhow::Result;
use clap::Parser;
use futures::future::join_all;
use std::{
collections::HashMap,
process::{Command, Output},
str::FromStr,
sync::Arc,
time::{Duration, Instant},
};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader as TokioBufReader};
use tracing::{info, Level};

use crate::{
Expand All @@ -24,15 +24,12 @@ use crate::{
faucet_txn_builder::{Erc20FaucetTxnBuilder, EthFaucetTxnBuilder, FaucetTxnBuilder},
PlanBuilder, TxnPlan,
},
util::gen_account::gen_account,
util::account_generator::AccountGenerator,
};

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
#[arg(long, default_value_t = false)]
recover: bool,

#[arg(long, default_value = "bench_config.toml")]
config: String,
}
Expand All @@ -47,27 +44,6 @@ mod config;
mod eth;
mod txn_plan;

async fn load_accounts_from_file(
path: &str,
) -> Result<HashMap<Arc<Address>, Arc<PrivateKeySigner>>> {
let file = tokio::fs::File::open(path).await?;
let reader = TokioBufReader::new(file);
let mut lines = reader.lines();
let mut accounts = HashMap::new();

while let Some(line) = lines.next_line().await? {
let parts: Vec<&str> = line.split(", ").collect();
if parts.len() == 2 {
let signer = PrivateKeySigner::from_str(parts[1])?;
let address = signer.address();
accounts.insert(Arc::new(address), Arc::new(signer));
} else {
return Err(anyhow::anyhow!("Invalid line in accounts file: {}", line));
}
}
Ok(accounts)
}

async fn run_plan(
plan: Box<dyn TxnPlan>,
producer: &Addr<Producer>,
Expand Down Expand Up @@ -214,36 +190,30 @@ async fn main() -> Result<()> {
.with_thread_ids(false)
.init();

let (contract_config, accounts) = if args.recover {
info!("Starting in recovery mode...");
let contract_config =
ContractConfig::load_from_file(&benchmark_config.contract_config_path).unwrap();
let accounts = load_accounts_from_file("accounts.txt").await?;
info!("Loaded {} accounts from accounts.txt", accounts.len());
(contract_config, accounts)
} else {
info!("Starting in normal mode...");
let mut command = format!(
"python scripts/deploy.py --private-key \"{}\" --num-tokens {} --output-file \"{}\" --rpc-url \"{}\"",
benchmark_config.faucet.private_key,
benchmark_config.num_tokens,
benchmark_config.contract_config_path,
benchmark_config.nodes[0].rpc_url
);
if benchmark_config.enable_swap_token {
command.push_str(" --enable-swap-token");
}
let res = run_command(&command).unwrap();
info!("{}", String::from_utf8_lossy(&res.stdout));
let contract_config = ContractConfig::load_from_file(
&benchmark_config.contract_config_path,
)
.unwrap_or_else(|e| {
let mut generator = AccountGenerator::new();

info!("Starting contract deployment and account generation...");
let mut command = format!(
"python scripts/deploy.py --private-key \"{}\" --num-tokens {} --output-file \"{}\" --rpc-url \"{}\"",
benchmark_config.faucet.private_key,
benchmark_config.num_tokens,
benchmark_config.contract_config_path,
benchmark_config.nodes[0].rpc_url
);
if benchmark_config.enable_swap_token {
command.push_str(" --enable-swap-token");
}
let res = run_command(&command).unwrap();
info!("{}", String::from_utf8_lossy(&res.stdout));

let contract_config =
ContractConfig::load_from_file(&benchmark_config.contract_config_path).unwrap_or_else(|e| {
panic!("Contract config file not found {}", e);
});
let accounts = gen_account(benchmark_config.accounts.num_accounts).unwrap();
(contract_config, accounts)
};

let accounts = generator
.gen(benchmark_config.accounts.num_accounts)
.unwrap();

let account_addresses = Arc::new(
accounts
Expand Down Expand Up @@ -281,99 +251,78 @@ async fn main() -> Result<()> {
Some(benchmark_config.target_tps as u32),
)
.start();
let nonce_map = init_nonce(&accounts, eth_clients[0].clone(), args.recover).await;

info!("Initializing nonces for {} accounts...", accounts.len());
let nonce_map = init_nonce(&accounts, eth_clients[0].clone()).await;
let producer = Producer::new(accounts.clone(), nonce_map, consumer, monitor)
.unwrap()
.start();
let chain_id = benchmark_config.nodes[0].chain_id;

if !args.recover {
let faucet_address =
PrivateKeySigner::from_str(&benchmark_config.faucet.private_key).unwrap();
let faucet_start_nonce = eth_clients[0]
.get_transaction_count(faucet_address.address())
.await
.unwrap();
let faucet_balance = eth_clients[0]
.get_balance(&faucet_address.address())
// Faucet distribution: always run this, the builder will handle recovery internally
let faucet_signer =
PrivateKeySigner::from_str(&benchmark_config.faucet.private_key).unwrap();
let faucet_balance = eth_clients[0]
.get_balance(&faucet_signer.address())
.await
.unwrap();

info!("Initializing ETH Faucet distribution...");
let eth_faucet_builder = PlanBuilder::create_faucet_tree_plan_builder(
benchmark_config.faucet.faucet_level as usize,
faucet_balance,
&benchmark_config.faucet.private_key,
account_addresses.clone(),
Arc::new(EthFaucetTxnBuilder),
U256::from(benchmark_config.num_tokens)
* U256::from(21000)
* U256::from(1000_000_000_000u64),
eth_clients[0].clone(),
&mut generator,
)
.await?;
execute_faucet_distribution(
eth_faucet_builder,
chain_id,
&producer,
"ETH",
benchmark_config.faucet.wait_duration_secs,
)
.await?;

let all_token_addresses = contract_config.get_all_token_addresses();

for token in &all_token_addresses {
info!("distributing token: {}", token);

let balance = IERC20::new(*token, eth_clients[0].provider())
.balanceOf(faucet_signer.address())
.call()
.await
.unwrap();

info!("Initializing Faucet constructor...");
let eth_faucet_builder = PlanBuilder::create_faucet_tree_plan_builder(
info!("balance of token: {}", balance);

let token_faucet_builder = PlanBuilder::create_faucet_tree_plan_builder(
benchmark_config.faucet.faucet_level as usize,
faucet_balance,
balance,
&benchmark_config.faucet.private_key,
faucet_start_nonce,
account_addresses.clone(),
Arc::new(EthFaucetTxnBuilder),
U256::from(benchmark_config.num_tokens)
* U256::from(21000)
* U256::from(1000_000_000_000u64),
Arc::new(Erc20FaucetTxnBuilder::new(*token)),
U256::ZERO,
eth_clients[0].clone(),
&mut generator,
)
.unwrap();
.await?;

execute_faucet_distribution(
eth_faucet_builder,
token_faucet_builder,
chain_id,
&producer,
"ETH",
&format!("Token {}", token),
benchmark_config.faucet.wait_duration_secs,
)
.await?;

let all_token_addresses = contract_config.get_all_token_addresses();
let faucet_signer_for_token =
PrivateKeySigner::from_str(&benchmark_config.faucet.private_key).unwrap();

for token in &all_token_addresses {
info!("distributing token: {}", token);

let faucet_current_nonce = eth_clients[0]
.get_transaction_count(faucet_signer_for_token.address())
.await
.unwrap();
let balance = IERC20::new(*token, eth_clients[0].provider())
.balanceOf(faucet_signer_for_token.address())
.call()
.await
.unwrap();

info!("balance of token: {}", balance);

let token_faucet_builder = PlanBuilder::create_faucet_tree_plan_builder(
benchmark_config.faucet.faucet_level as usize,
balance,
&benchmark_config.faucet.private_key,
faucet_current_nonce,
account_addresses.clone(),
Arc::new(Erc20FaucetTxnBuilder::new(*token)),
U256::ZERO,
)
.unwrap();

execute_faucet_distribution(
token_faucet_builder,
chain_id,
&producer,
&format!("Token {}", token),
benchmark_config.faucet.wait_duration_secs,
)
.await?;
}

let mut file = tokio::fs::File::create("accounts.txt").await.unwrap();
for account in accounts.iter() {
file.write(
format!(
"{}, {}\n",
account.0.to_string(),
hex::encode(account.1.to_bytes())
)
.as_bytes(),
)
.await
.unwrap();
}
}

let tps = benchmark_config.target_tps as usize;
Expand Down Expand Up @@ -407,21 +356,18 @@ async fn main() -> Result<()> {
async fn init_nonce(
accounts: &HashMap<Arc<Address>, Arc<PrivateKeySigner>>,
eth_client: Arc<EthHttpCli>,
recover: bool,
) -> HashMap<Arc<Address>, u32> {
let mut nonce_map = HashMap::with_capacity(accounts.len());
if recover {
for account in accounts.iter() {
let nonce = eth_client
.get_transaction_count(*account.0.clone())
let nonce_futs = accounts.keys().map(|account| {
let client = eth_client.clone();
async move {
let nonce = client
.get_transaction_count(*account.clone())
.await
.unwrap();
nonce_map.insert(account.0.clone(), nonce as u32);
}
} else {
for account in accounts.iter() {
nonce_map.insert(account.0.clone(), 0);
(account.clone(), nonce as u32)
}
}
nonce_map
});

let results = join_all(nonce_futs).await;
results.into_iter().collect()
}
Loading
Loading