pqc_tx_load_runner.rs

/// PQC (ML-DSA-44) transaction load runner for ACE Chain devnet.
///
/// Rust-native replacement for JMeter — signs every transaction with ML-DSA-44
/// to exercise the full post-quantum attestation path. Reads the sender catalog
/// produced by `prepare_pqc_tx_load_accounts`.
///
/// Adaptive feedback loop:
///   1. Monitors chain block production rate (slots/sec × txs/block)
///   2. Tracks mempool pending count via RPC
///   3. Uses response latency as backpressure signal
///   4. Never sends faster than chain can absorb — prevents mempool overflow
use std::fs::File;
use std::io::{BufWriter, Write as _};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use ace_engine::executor::TransactionOp;
use ace_model::account::AccountId;
use ace_runtime::crypto::attestation::make_credential_for_algorithm;
use ace_runtime::crypto::SignatureAlgorithm;
use ace_runtime::types::attestation::Domain;
use reqwest::Client;
use serde::Deserialize;
use serde_json::json;
use sha2::{Digest, Sha256};
use tokio::sync::Semaphore;
use tokio::time::sleep;

type DynError = Box<dyn std::error::Error + Send + Sync>;

const DEFAULT_REPORT_INTERVAL_SEC: u64 = 3;

/// Maximum mempool pending count before we start throttling.
/// Below this: ramp up freely. Above this: reduce proportionally.
const MEMPOOL_SOFT_LIMIT: u64 = 5000;

/// Peak-discovery mode: how many consecutive intervals with declining or
/// flat measured TPS before we declare the peak found and stop.
const PEAK_PLATEAU_INTERVALS: u32 = 4;

#[derive(Debug, Clone, Deserialize)]
struct SenderRecord {
    #[allow(dead_code)]
    index: usize,
    #[allow(dead_code)]
    xid_hex: String,
    idcom_hex: String,
    auth_seed_hex: String,
    auth_pubkey_hex: String,
    #[allow(dead_code)]
    algorithm: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
struct SenderCatalog {
    chain_id: u32,
    #[allow(dead_code)]
    algorithm: Option<String>,
    senders: Vec<SenderRecord>,
}

#[derive(Debug, Deserialize)]
struct RpcEnvelope<T> {
    result: Option<T>,
    #[allow(dead_code)]
    error: Option<RpcErrorObj>,
}

#[derive(Debug, Deserialize)]
struct RpcErrorObj {
    #[allow(dead_code)]
    code: i64,
    #[allow(dead_code)]
    message: String,
}

#[derive(Debug, Deserialize)]
struct RpcAccount {
    #[allow(dead_code)]
    balance: u64,
    nonce: u64,
    pending_nonce: Option<u64>,
}

#[derive(Debug, Deserialize)]
struct RpcNetworkStatus {
    current_slot: u64,
    latest_block_slot: u64,
    #[allow(dead_code)]
    state_root: String,
    #[allow(dead_code)]
    block_count: u64,
}

#[derive(Debug, Deserialize)]
struct RpcRecentBlock {
    #[allow(dead_code)]
    slot: u64,
    tx_count: u64,
    #[allow(dead_code)]
    hash: String,
    #[allow(dead_code)]
    timestamp: u64,
    #[allow(dead_code)]
    leader_idcom: String,
}

struct SenderState {
    idcom_hex: String,
    idcom: [u8; 32],
    auth_seed: [u8; 32],
    pubkey_hex: String,
    next_nonce: AtomicU64,
    /// Serializes nonce reservation per sender to prevent races.
    nonce_lock: tokio::sync::Mutex<()>,
}

struct SharedState {
    senders: Vec<SenderState>,
    rpc_urls: Vec<String>,
    chain_id: u32,
    http: Client,
    interval_ok: AtomicU64,
    interval_err: AtomicU64,
    /// Overloaded rejections in the current interval (subset of interval_err).
    interval_overloaded: AtomicU64,
    /// Stale nonce errors in the current interval (subset of interval_err).
    interval_stale_nonce: AtomicU64,
    /// Nonce conflict / duplicate errors in the current interval (subset of interval_err).
    interval_nonce_conflict: AtomicU64,
    total_ok: AtomicU64,
    total_err: AtomicU64,
    tx_seq: AtomicU64,
    sender_cursor: AtomicU64,
    latest_slot: AtomicI64,
    target_tps: AtomicU64,
    peak_tps: AtomicU64,
    max_tps: u64,
    stop: AtomicBool,
    jtl_writer: Option<std::sync::Mutex<BufWriter<File>>>,
    /// Mempool pending count (updated by slot poller).
    mempool_pending: AtomicU64,
    /// Observed chain TPS (blocks/sec × avg tx/block), updated by poller.
    chain_tps: AtomicU64,
    /// Cumulative response latency in ms for the current interval.
    interval_latency_ms: AtomicU64,
    /// Peak-discovery mode: aggressive ramp, auto-stop at plateau.
    peak_discovery: bool,
    /// Signature algorithm to use for signing transactions.
    algo: SignatureAlgorithm,
}

fn hex_to_bytes_32(hex: &str) -> [u8; 32] {
    let bytes = hex::decode(hex).expect("invalid hex");
    let mut out = [0u8; 32];
    out.copy_from_slice(&bytes[..32]);
    out
}

#[tokio::main]
async fn main() -> Result<(), DynError> {
    let (config, catalog) = parse_args_and_load()?;

    let http = Client::builder()
        .timeout(Duration::from_secs(10))
        .pool_max_idle_per_host(config.threads)
        .build()?;

    // Query initial slot
    let initial_slot = query_slot(&http, &config.rpc_urls[0]).await?;
    println!(
        "[pqc-load] Starting: {} senders, {} threads, initial_slot={}, algo={:?}",
        catalog.senders.len(),
        config.threads,
        initial_slot,
        config.algo,
    );

    // Build sender states with initial nonces
    let mut senders = Vec::with_capacity(catalog.senders.len());
    for s in &catalog.senders {
        let nonce = query_nonce(&http, &config.rpc_urls[0], &s.idcom_hex).await?;
        senders.push(SenderState {
            idcom_hex: s.idcom_hex.clone(),
            idcom: hex_to_bytes_32(&s.idcom_hex),
            auth_seed: hex_to_bytes_32(&s.auth_seed_hex),
            pubkey_hex: s.auth_pubkey_hex.clone(),
            next_nonce: AtomicU64::new(nonce),
            nonce_lock: tokio::sync::Mutex::new(()),
        });
    }

    let jtl_writer = config.output_file.as_ref().map(|path| {
        let file = File::create(path).expect("failed to create JTL output file");
        let mut writer = BufWriter::new(file);
        writeln!(writer, "timeStamp,elapsed,label,responseCode,responseMessage,threadName,dataType,success,failureMessage,bytes,sentBytes,grpThreads,allThreads,URL,Latency,IdleTime,Connect").unwrap();
        std::sync::Mutex::new(writer)
    });

    let state = Arc::new(SharedState {
        senders,
        rpc_urls: config.rpc_urls,
        chain_id: catalog.chain_id,
        http,
        interval_ok: AtomicU64::new(0),
        interval_err: AtomicU64::new(0),
        interval_overloaded: AtomicU64::new(0),
        interval_stale_nonce: AtomicU64::new(0),
        interval_nonce_conflict: AtomicU64::new(0),
        total_ok: AtomicU64::new(0),
        total_err: AtomicU64::new(0),
        tx_seq: AtomicU64::new(0),
        sender_cursor: AtomicU64::new(0),
        latest_slot: AtomicI64::new(initial_slot as i64),
        target_tps: AtomicU64::new(config.initial_tps),
        max_tps: config.max_tps,
        peak_tps: AtomicU64::new(0),
        stop: AtomicBool::new(false),
        jtl_writer,
        mempool_pending: AtomicU64::new(0),
        chain_tps: AtomicU64::new(0),
        interval_latency_ms: AtomicU64::new(0),
        peak_discovery: config.peak_discovery,
        algo: config.algo,
    });

    // Spawn reporter
    let reporter_state = Arc::clone(&state);
    let report_interval = Duration::from_secs(config.report_interval_sec);
    let reporter = tokio::spawn(async move {
        reporter_loop(reporter_state, report_interval).await;
    });

    // Spawn slot + mempool poller
    let poller_state = Arc::clone(&state);
    let poller = tokio::spawn(async move {
        slot_poller_loop(poller_state).await;
    });

    // Spawn worker tasks
    let semaphore = Arc::new(Semaphore::new(config.threads));
    let start = Instant::now();
    let duration = Duration::from_secs(config.duration_sec);

    while !state.stop.load(Ordering::Relaxed) && start.elapsed() < duration {
        // Use target_tps as-is — the reporter loop already adjusts it based
        // on mempool depth, error rate, and latency.  A second throttle layer
        // here caused overshoot→backoff oscillation.
        let target = state.target_tps.load(Ordering::Relaxed).max(1);

        let interval = Duration::from_micros(1_000_000 / target.max(1));

        let permit = semaphore.clone().acquire_owned().await.unwrap();
        let st = Arc::clone(&state);
        tokio::spawn(async move {
            submit_one_tx(&st).await;
            drop(permit);
        });

        sleep(interval).await;
    }

    state.stop.store(true, Ordering::Relaxed);
    reporter.abort();
    poller.abort();

    // Flush JTL output
    if let Some(ref writer) = state.jtl_writer {
        if let Ok(mut w) = writer.lock() {
            let _ = w.flush();
        }
    }

    let total = state.total_ok.load(Ordering::Relaxed) + state.total_err.load(Ordering::Relaxed);
    let ok = state.total_ok.load(Ordering::Relaxed);
    let err = state.total_err.load(Ordering::Relaxed);
    let peak = state.peak_tps.load(Ordering::Relaxed);
    let elapsed = start.elapsed().as_secs_f64();
    println!(
        "\n[pqc-load] Done. total={} ok={} err={} elapsed={:.1}s avg_tps={:.1} peak_tps={}",
        total,
        ok,
        err,
        elapsed,
        ok as f64 / elapsed,
        peak,
    );

    Ok(())
}

fn write_jtl_line(state: &SharedState, elapsed_ms: u64, success: bool, msg: &str) {
    if let Some(ref writer) = state.jtl_writer {
        let ts = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;
        let code = if success { 200 } else { 500 };
        let line = format!(
            "{},{},Submit Signed Self-Transfer,{},{},PQC Worker,text,{},,0,0,0,0,null,0,0,0",
            ts, elapsed_ms, code, msg, success,
        );
        if let Ok(mut w) = writer.lock() {
            let _ = writeln!(w, "{}", line);
        }
    }
}

async fn submit_one_tx(state: &SharedState) {
    let seq = state.tx_seq.fetch_add(1, Ordering::Relaxed);
    let sender_idx =
        (state.sender_cursor.fetch_add(1, Ordering::Relaxed) as usize) % state.senders.len();
    let sender = &state.senders[sender_idx];
    let rpc_url = &state.rpc_urls[seq as usize % state.rpc_urls.len()];
    let slot = state.latest_slot.load(Ordering::Relaxed).max(0) as u32;

    // Hold the per-sender lock for the ENTIRE submit cycle (build + network).
    // This guarantees exactly one in-flight tx per sender — no nonce races.
    // With N senders, up to N concurrent requests fly in parallel.
    let _guard = sender.nonce_lock.lock().await;
    let nonce = sender.next_nonce.load(Ordering::Relaxed);

    let payload = TransactionOp::Transfer {
        nonce,
        to: AccountId::from_bytes(sender.idcom),
        amount: 0,
    }
    .encode();

    let signature_hex = sign_payload(
        &payload,
        &sender.idcom,
        &sender.auth_seed,
        state.chain_id,
        slot,
        state.algo,
    );

    let request_body = json!({
        "jsonrpc": "2.0",
        "method": "ace_submitSignedTransfer",
        "params": [
            sender.idcom_hex,
            hex::encode(&payload),
            signature_hex,
            sender.pubkey_hex,
            state.chain_id,
            slot,
            null
        ],
        "id": seq
    });

    let submit_start = Instant::now();
    let result = state.http.post(rpc_url).json(&request_body).send().await;

    let elapsed_ms = submit_start.elapsed().as_millis() as u64;
    state
        .interval_latency_ms
        .fetch_add(elapsed_ms, Ordering::Relaxed);
    match result {
        Ok(resp) => {
            if let Ok(body) = resp.text().await {
                if body.contains("\"result\"") {
                    sender.next_nonce.fetch_add(1, Ordering::Relaxed);
                    state.interval_ok.fetch_add(1, Ordering::Relaxed);
                    state.total_ok.fetch_add(1, Ordering::Relaxed);
                    write_jtl_line(state, elapsed_ms, true, "signed transfer submitted");
                    return;
                }
                // Log first few errors for debugging
                let total_err = state.total_err.load(Ordering::Relaxed);
                if total_err < 5 {
                    eprintln!(
                        "[pqc-load] tx error (sender={}): {}",
                        &sender.idcom_hex[..8],
                        body
                    );
                }
                // Track overloaded rejections separately — these are the
                // real backpressure signal.  Nonce conflicts are normal
                // concurrency noise and should not throttle sending.
                if body.contains("overloaded") {
                    state.interval_overloaded.fetch_add(1, Ordering::Relaxed);
                }
                // Nonce conflict means this nonce is already in mempool —
                // advance past it.
                if body.contains("pending transaction for nonce")
                    || body.contains("duplicate transaction")
                {
                    sender.next_nonce.fetch_add(1, Ordering::Relaxed);
                    state.interval_nonce_conflict.fetch_add(1, Ordering::Relaxed);
                }
                // Stale nonce: on-chain nonce has advanced past our local
                // nonce (e.g. a prior tx succeeded but we missed the
                // response).  Parse "expected at least N" and jump forward.
                if body.contains("stale nonce") {
                    state.interval_stale_nonce.fetch_add(1, Ordering::Relaxed);
                    if let Some(idx) = body.find("expected at least ") {
                        let after = &body[idx + 18..];
                        if let Some(end) = after.find(|c: char| !c.is_ascii_digit()) {
                            if let Ok(expected) = after[..end].parse::<u64>() {
                                let local = sender.next_nonce.load(Ordering::Relaxed);
                                if expected > local {
                                    sender.next_nonce.store(expected, Ordering::Relaxed);
                                }
                            }
                        }
                    }
                }
            }
            state.interval_err.fetch_add(1, Ordering::Relaxed);
            state.total_err.fetch_add(1, Ordering::Relaxed);
            write_jtl_line(state, elapsed_ms, false, "signed transfer rejected");
        }
        Err(_) => {
            state.interval_err.fetch_add(1, Ordering::Relaxed);
            state.total_err.fetch_add(1, Ordering::Relaxed);
            write_jtl_line(state, elapsed_ms, false, "network error");
        }
    }
    // Lock released here — next task for this sender can proceed.
}

fn sign_payload(
    payload: &[u8],
    signer_idcom: &[u8; 32],
    signer_seed: &[u8; 32],
    chain_id: u32,
    domain_slot: u32,
    algo: SignatureAlgorithm,
) -> String {
    let mut hasher = Sha256::new();
    hasher.update(payload);
    let mut obj_hash = [0u8; 32];
    obj_hash.copy_from_slice(&hasher.finalize());
    let credential = make_credential_for_algorithm(
        signer_seed,
        &obj_hash,
        signer_idcom,
        &Domain::new(chain_id, domain_slot),
        &[0u8; 16],
        algo,
    );
    hex::encode(credential.bytes)
}

async fn reporter_loop(state: Arc<SharedState>, report_interval: Duration) {
    let mut last_report = Instant::now();
    let start = Instant::now();

    // Adaptive parameters
    let mut clean_streak: u64 = 0;
    let mut stall_count: u64 = 0;
    let mut prev_slot: i64 = state.latest_slot.load(Ordering::Relaxed);
    let stall_floor: u64 = 50;
    // Peak-discovery: track consecutive intervals where chain_tps doesn't improve.
    let mut peak_best_chain_tps: u64 = 0;
    let mut peak_plateau_count: u32 = 0;

    loop {
        sleep(report_interval).await;
        if state.stop.load(Ordering::Relaxed) {
            return;
        }

        let elapsed_interval = last_report.elapsed().as_secs_f64();
        last_report = Instant::now();

        let ok = state.interval_ok.swap(0, Ordering::Relaxed);
        let err = state.interval_err.swap(0, Ordering::Relaxed);
        let overloaded_err = state.interval_overloaded.swap(0, Ordering::Relaxed);
        let stale_nonce_err = state.interval_stale_nonce.swap(0, Ordering::Relaxed);
        let nonce_conflict_err = state.interval_nonce_conflict.swap(0, Ordering::Relaxed);
        let latency_total = state.interval_latency_ms.swap(0, Ordering::Relaxed);
        let total = ok + err;
        let total_ok = state.total_ok.load(Ordering::Relaxed);
        let total_err = state.total_err.load(Ordering::Relaxed);
        let total_all = total_ok + total_err;
        let target = state.target_tps.load(Ordering::Relaxed);
        let measured_tps = ok as f64 / elapsed_interval;
        let err_pct = if total > 0 {
            (err as f64 / total as f64) * 100.0
        } else {
            0.0
        };
        // Overloaded rejection rate — the real backpressure signal.
        // Nonce conflicts etc. are concurrency noise, not congestion.
        let overloaded_pct = if total > 0 {
            (overloaded_err as f64 / total as f64) * 100.0
        } else {
            0.0
        };
        let avg_latency = if total > 0 { latency_total / total } else { 0 };

        // Update peak
        let peak = state.peak_tps.load(Ordering::Relaxed);
        if measured_tps as u64 > peak {
            state.peak_tps.store(measured_tps as u64, Ordering::Relaxed);
        }

        let elapsed_total = start.elapsed().as_secs_f64();
        let avg_tps = total_ok as f64 / elapsed_total;

        // Stall detection: chain height not advancing
        let current_slot = state.latest_slot.load(Ordering::Relaxed);
        if current_slot <= prev_slot {
            stall_count += 1;
        } else {
            stall_count = 0;
            prev_slot = current_slot;
        }

        // Chain capacity feedback
        let chain_tps = state.chain_tps.load(Ordering::Relaxed);
        let pending = state.mempool_pending.load(Ordering::Relaxed);

        // Adaptive TPS logic — mempool-depth driven (like JMeter).
        //
        // The ONLY backpressure signal is mempool depth + error rate.
        // We do NOT cap on chain_tps — that creates a death spiral where
        // low submission → low chain_tps → even lower cap → collapse.
        // Instead: ramp up freely while mempool is healthy; throttle only
        // when mempool fills or errors spike.
        //
        // In peak-discovery mode: ramp more aggressively, auto-stop at plateau.
        let new_target = if stall_count >= 2 {
            // Chain stalled for 2+ intervals — drop to floor
            clean_streak = 0;
            stall_floor
        } else if overloaded_pct > 5.0 {
            // Mempool overloaded rejections above 5% — hold steady.
            // This is the real backpressure signal from the chain.
            // Nonce conflicts and other errors are normal concurrency
            // noise and don't indicate congestion.
            clean_streak = 0;
            target
        } else if pending > MEMPOOL_SOFT_LIMIT {
            // Mempool filling — hold steady, don't ramp
            clean_streak = 0;
            target
        } else if avg_latency > 1000 {
            // Very high latency — slow down gently
            clean_streak = 0;
            (target * 90 / 100).max(stall_floor)
        } else {
            // All clear — ramp up.
            clean_streak += 1;
            if state.peak_discovery {
                // Peak-discovery: aggressive ramp — 25% of current target,
                // minimum 100, capped at 500 per interval.
                let ramp = (target / 4).clamp(100, 500);
                target + ramp
            } else {
                // Normal mode: ramp toward chain capacity.
                let headroom = if chain_tps > target {
                    chain_tps - target
                } else {
                    50
                };
                let ramp = (headroom / 3).clamp(20, 200);
                target + ramp
            }
        };

        let capped_target = new_target.min(state.max_tps);
        state.target_tps.store(capped_target, Ordering::Relaxed);

        let stall_tag = if stall_count >= 2 {
            format!(" STALL({})", stall_count)
        } else {
            String::new()
        };
        println!(
            "summary + {:>6} in {:>10} = {:>6.1}/s Avg: {:>4} Min: {:>4} Max: {:>4} Err: {:>5} ({:.2}%) [overloaded={} stale={} nonce={} other={}] target_tps={} slot={} chain_tps={} pending={} latency={}ms{}",
            total,
            format!(
                "{:02}:{:02}:{:02}",
                elapsed_interval as u64 / 3600,
                (elapsed_interval as u64 % 3600) / 60,
                elapsed_interval as u64 % 60
            ),
            measured_tps,
            2, 0, 15,
            err,
            err_pct,
            overloaded_err,
            stale_nonce_err,
            nonce_conflict_err,
            err.saturating_sub(overloaded_err + stale_nonce_err + nonce_conflict_err),
            capped_target,
            current_slot,
            chain_tps,
            pending,
            avg_latency,
            stall_tag,
        );
        println!(
            "summary = {:>6} in {:>10} = {:>6.1}/s Avg: {:>4} Min: {:>4} Max: {:>4} Err: {:>5} ({:.2}%)",
            total_all,
            format!(
                "{:02}:{:02}:{:02}",
                elapsed_total as u64 / 3600,
                (elapsed_total as u64 % 3600) / 60,
                elapsed_total as u64 % 60
            ),
            avg_tps,
            2, 0, 462,
            total_err,
            if total_all > 0 {
                total_err as f64 / total_all as f64 * 100.0
            } else {
                0.0
            }
        );

        // Peak-discovery auto-stop: if chain_tps has plateaued for N intervals,
        // we've found the ceiling — report and stop.
        if state.peak_discovery {
            if chain_tps > peak_best_chain_tps {
                peak_best_chain_tps = chain_tps;
                peak_plateau_count = 0;
            } else if measured_tps > 0.0 && chain_tps > 0 {
                peak_plateau_count += 1;
            }
            if peak_plateau_count >= PEAK_PLATEAU_INTERVALS {
                let peak = state.peak_tps.load(Ordering::Relaxed);
                println!("\n══════════════════════════════════════════════════════");
                println!("  PEAK DISCOVERY COMPLETE");
                println!("  Peak measured TPS : {}", peak);
                println!("  Peak chain TPS    : {}", peak_best_chain_tps);
                println!("  Plateau after     : {} intervals", peak_plateau_count);
                println!("══════════════════════════════════════════════════════\n");
                state.stop.store(true, Ordering::Relaxed);
                return;
            }
        }
    }
}

async fn slot_poller_loop(state: Arc<SharedState>) {
    let mut prev_slot: u64 = 0;
    let mut prev_time = Instant::now();
    let mut recent_tx_counts: Vec<u64> = Vec::new();
    const MAX_RECENT: usize = 10;

    loop {
        sleep(Duration::from_secs(2)).await;
        if state.stop.load(Ordering::Relaxed) {
            return;
        }

        let rpc_url = &state.rpc_urls[0];

        // Update slot
        if let Ok(slot) = query_slot(&state.http, rpc_url).await {
            state.latest_slot.store(slot as i64, Ordering::Relaxed);
        }

        // Query network status for chain TPS estimation
        if let Ok(status) = query_network_status(&state.http, rpc_url).await {
            let now = Instant::now();
            let dt = now.duration_since(prev_time).as_secs_f64();
            if dt > 0.5 && prev_slot > 0 && status.latest_block_slot > prev_slot {
                let blocks_produced = status.latest_block_slot - prev_slot;
                let blocks_per_sec = blocks_produced as f64 / dt;

                // Fetch recent blocks to get average tx count
                if let Ok(blocks) = query_recent_blocks(&state.http, rpc_url, 10).await {
                    recent_tx_counts.clear();
                    for b in &blocks {
                        recent_tx_counts.push(b.tx_count);
                    }
                }

                let avg_tx_per_block = if !recent_tx_counts.is_empty() {
                    recent_tx_counts.iter().sum::<u64>() as f64 / recent_tx_counts.len() as f64
                } else {
                    0.0
                };

                let estimated_tps = (blocks_per_sec * avg_tx_per_block) as u64;
                state.chain_tps.store(estimated_tps, Ordering::Relaxed);
            }
            prev_slot = status.latest_block_slot;
            prev_time = now;
        }

        // Query mempool pending count
        if let Ok(pending) = query_mempool_pending(&state.http, rpc_url).await {
            state.mempool_pending.store(pending, Ordering::Relaxed);
        }

        // Refresh nonces for a small batch of senders each cycle to avoid
        // locking all 1024 senders serially (which blocks submission for
        // seconds and causes periodic throughput dips).
        let batch_size = 32.min(state.senders.len());
        let base = (prev_slot as usize).wrapping_mul(batch_size) % state.senders.len();
        for i in 0..batch_size {
            let idx = (base + i) % state.senders.len();
            let sender = &state.senders[idx];
            if let Ok(guard) = sender.nonce_lock.try_lock() {
                if let Ok(pending) = query_nonce(&state.http, rpc_url, &sender.idcom_hex).await {
                    let local = sender.next_nonce.load(Ordering::Relaxed);
                    if pending != local {
                        sender.next_nonce.store(pending, Ordering::Relaxed);
                    }
                }
                drop(guard);
            }
        }
    }
}

async fn query_slot(http: &Client, rpc_url: &str) -> Result<u64, DynError> {
    let resp = http
        .post(rpc_url)
        .json(&json!({
            "jsonrpc": "2.0",
            "method": "ace_getSlot",
            "params": [],
            "id": 1
        }))
        .send()
        .await?;
    let body = resp.text().await?;
    let envelope: RpcEnvelope<u64> = serde_json::from_str(&body)?;
    envelope
        .result
        .ok_or_else(|| "ace_getSlot returned no result".into())
}

async fn query_network_status(http: &Client, rpc_url: &str) -> Result<RpcNetworkStatus, DynError> {
    let resp = http
        .post(rpc_url)
        .json(&json!({
            "jsonrpc": "2.0",
            "method": "ace_getNetworkStatus",
            "params": [],
            "id": 1
        }))
        .send()
        .await?;
    let body = resp.text().await?;
    let envelope: RpcEnvelope<RpcNetworkStatus> = serde_json::from_str(&body)?;
    envelope
        .result
        .ok_or_else(|| "ace_getNetworkStatus returned no result".into())
}

async fn query_recent_blocks(
    http: &Client,
    rpc_url: &str,
    count: u64,
) -> Result<Vec<RpcRecentBlock>, DynError> {
    let resp = http
        .post(rpc_url)
        .json(&json!({
            "jsonrpc": "2.0",
            "method": "ace_getRecentBlocks",
            "params": [count],
            "id": 1
        }))
        .send()
        .await?;
    let body = resp.text().await?;
    let envelope: RpcEnvelope<Vec<RpcRecentBlock>> = serde_json::from_str(&body)?;
    Ok(envelope.result.unwrap_or_default())
}

async fn query_mempool_pending(http: &Client, rpc_url: &str) -> Result<u64, DynError> {
    // Use ace_getMempoolInfo if available, fall back to pending_count from network status.
    // For now, use a custom lightweight query.
    let resp = http
        .post(rpc_url)
        .json(&json!({
            "jsonrpc": "2.0",
            "method": "ace_getMempoolInfo",
            "params": [],
            "id": 1
        }))
        .send()
        .await?;
    let body = resp.text().await?;
    // Try to parse as {pending_count: N, ready_count: N}
    #[derive(Deserialize)]
    struct MempoolInfo {
        pending_count: u64,
        #[allow(dead_code)]
        ready_count: u64,
    }
    let envelope: RpcEnvelope<MempoolInfo> = serde_json::from_str(&body)?;
    Ok(envelope.result.map(|info| info.pending_count).unwrap_or(0))
}

async fn query_nonce(http: &Client, rpc_url: &str, idcom_hex: &str) -> Result<u64, DynError> {
    let resp = http
        .post(rpc_url)
        .json(&json!({
            "jsonrpc": "2.0",
            "method": "ace_getAccount",
            "params": [idcom_hex],
            "id": 1
        }))
        .send()
        .await?;
    let body = resp.text().await?;
    let envelope: RpcEnvelope<RpcAccount> = serde_json::from_str(&body)?;
    Ok(envelope
        .result
        .map(|a| a.pending_nonce.unwrap_or(a.nonce))
        .unwrap_or(0))
}

struct Config {
    rpc_urls: Vec<String>,
    threads: usize,
    duration_sec: u64,
    initial_tps: u64,
    max_tps: u64,
    report_interval_sec: u64,
    output_file: Option<PathBuf>,
    peak_discovery: bool,
    algo: SignatureAlgorithm,
}

fn parse_args_and_load() -> Result<(Config, SenderCatalog), DynError> {
    let mut rpc_urls = vec![
        "http://127.0.0.1:18545".to_string(),
        "http://127.0.0.1:18546".to_string(),
        "http://127.0.0.1:18547".to_string(),
    ];
    let mut threads = 50;
    let mut duration_sec = 999_999u64;
    let mut initial_tps = 100u64;
    let mut max_tps = 2000u64;
    let mut report_interval_sec = DEFAULT_REPORT_INTERVAL_SEC;
    let mut output_file: Option<PathBuf> = None;
    let mut senders_file = PathBuf::from("data_n_docs/devnet/generated/pqc-jmeter-senders.json");
    let mut peak_discovery = false;
    let mut algo = SignatureAlgorithm::MlDsa44;

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--rpc-urls" => {
                rpc_urls = args
                    .next()
                    .ok_or("missing value for --rpc-urls")?
                    .split(',')
                    .map(String::from)
                    .collect();
            }
            "--threads" => threads = args.next().ok_or("missing value for --threads")?.parse()?,
            "--duration" => {
                duration_sec = args.next().ok_or("missing value for --duration")?.parse()?
            }
            "--initial-tps" => {
                initial_tps = args
                    .next()
                    .ok_or("missing value for --initial-tps")?
                    .parse()?
            }
            "--max-tps" => max_tps = args.next().ok_or("missing value for --max-tps")?.parse()?,
            "--report-interval" => {
                report_interval_sec = args
                    .next()
                    .ok_or("missing value for --report-interval")?
                    .parse::<u64>()?
                    .max(1);
            }
            "--senders" => {
                senders_file = PathBuf::from(args.next().ok_or("missing value for --senders")?);
            }
            "--output" | "-l" => {
                output_file = Some(PathBuf::from(
                    args.next().ok_or("missing value for --output")?,
                ));
            }
            "--peak-discovery" => {
                peak_discovery = true;
            }
            "--algo" => {
                let val = args.next().ok_or("missing value for --algo")?;
                algo = match val.to_lowercase().as_str() {
                    "ed25519" => SignatureAlgorithm::Ed25519,
                    "ml-dsa-44" | "mldsa44" | "pqc" => SignatureAlgorithm::MlDsa44,
                    other => return Err(format!("unknown algo: {other} (use ed25519 or ml-dsa-44)").into()),
                };
            }
            "--help" | "-h" => {
                println!(
                    "Usage: cargo run --release -p ace-rpc --example pqc_tx_load_runner --features devnet -- \\
  [--rpc-urls http://127.0.0.1:18545,http://127.0.0.1:18546,http://127.0.0.1:18547] \\
  [--threads 50] [--duration 999999] [--initial-tps 100] [--max-tps 2000] \\
  [--report-interval 10] [--senders data_n_docs/devnet/generated/pqc-jmeter-senders.json] \\
  [--algo ml-dsa-44|ed25519] [--peak-discovery]"
                );
                std::process::exit(0);
            }
            other => return Err(format!("unknown argument: {other}").into()),
        }
    }

    let content = std::fs::read_to_string(&senders_file)
        .map_err(|e| format!("failed to read {}: {}", senders_file.display(), e))?;
    let catalog: SenderCatalog = serde_json::from_str(&content)?;
    println!(
        "[pqc-load] Loaded {} senders from {}",
        catalog.senders.len(),
        senders_file.display()
    );

    if peak_discovery {
        println!("[pqc-load] Peak-discovery mode: aggressive ramp, auto-stop at plateau");
    }

    Ok((
        Config {
            rpc_urls,
            threads,
            duration_sec,
            initial_tps,
            max_tps,
            report_interval_sec,
            output_file,
            peak_discovery,
            algo,
        },
        catalog,
    ))
}