Fix the Bitcoin reorg bug and Add BtcTx bench (#1990)

crates/rooch-benchmarks/Cargo.toml

@@ -50,6 +50,7 @@ parking_lot = { workspace = true }
 proptest = { workspace = true }
 rand_core = { default-features = false, workspace = true }
 bitcoincore-rpc-json = { workspace = true }
+bitcoincore-rpc = { workspace = true }
 bitcoin = { workspace = true }
 toml = { workspace = true }
 

crates/rooch-benchmarks/README.md

@@ -43,12 +43,25 @@ is `rooch-benchmarks/config/bench_tx.toml`.
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize, Parser, Eq)]
 pub struct BenchTxConfig {
     pub tx_type: Option<TxType>,    // empty(default)/transfer/btc-block
-    pub data_import_flag: bool,
-    pub btc_block_dir: Option<String>, // btc block dir, file name: <height>.hex
+    pub btc_block_dir: Option<String>, // btc block dir, default: target/btc_blocks, file name: <height>.hex
+    pub btc_block_start_height: Option<u64>, // btc block start height, default: 820000
+    pub btc_rpc_url: Option<String>,
+    pub btc_rpc_username: Option<String>,
+    pub btc_rpc_password: Option<String>,
     pub pprof_output: Option<PProfOutput>, // flamegraph(default)/proto
 }
 ```
 
+The env var has higher priority than the config file.
+
+* `ROOCH_BENCH_TX_TYPE`: override `tx_type` in config file.
+* `ROOCH_BENCH_BTC_BLOCK_DIR`: override `btc_block_dir` in config file.
+* `ROOCH_BENCH_BTC_BLOCK_START_HEIGHT`: override `btc_block_start_height` in config file.
+* `ROOCH_BENCH_BTC_RPC_URL`: override `btc_rpc_url` in config file.
+* `ROOCH_BENCH_BTC_RPC_USERNAME`: override `btc_rpc_username` in config file.
+* `ROOCH_BENCH_BTC_RPC_PASSWORD`: override `btc_rpc_password` in config file.
+* `ROOCH_BENCH_PPROF_OUTPUT`: override `pprof_output` in config file.
+
 ## Profiling
 
 When your run bench with `-- --profile-time=<seconds>` option, it will generate a flamegraph file
@@ -69,6 +82,7 @@ for PPROF_OUT output location:
 
 1. `l2_tx_<transfer/empty>`
 2. `btc_block`
+3. `btc_tx`
 
 for proto, run these to get svg:
 
@@ -81,3 +95,11 @@ pprof -svg profile.pb
 ### Why not run in CI pipeline?
 
 Coming soon...
+
+### How to prepare the Bitcoin blocks
+
+Run the benchmark with Bitcoin RPC config, it will download the blocks from Bitcoin network and save them in `target/btc_blocks` dir.
+
+```shell
+ROOCH_BENCH_TX_TYPE=btc_tx ROOCH_BENCH_BTC_RPC_URL=http://localhost:8332 ROOCH_BENCH_BTC_RPC_USERNAME=YourBTCUser ROOCH_BENCH_BTC_RPC_PASSWORD=YourBTCPass cargo bench -p rooch-benchmarks --bench bench_tx_exec
+```

crates/rooch-benchmarks/config/bench_tx.toml

@@ -2,4 +2,4 @@ tx_type = "empty"
 pprof_output = "flamegraph"
 # when tx_type is "btc-block", the following two options are required
 data_import_flag = true
-btc_block_dir = "/home/rooch/btc/blk"
+btc_block_dir = "target/btc_blocks"

crates/rooch-benchmarks/src/config.rs

@@ -50,6 +50,7 @@ pub enum TxType {
     Empty,
     Transfer,
     BtcBlock,
+    BtcTx,
 }
 
 impl FromStr for TxType {
@@ -58,6 +59,7 @@ impl FromStr for TxType {
         match s {
             "transfer" => Ok(TxType::Transfer),
             "btc_block" => Ok(TxType::BtcBlock),
+            "btc_tx" => Ok(TxType::BtcTx),
             "empty" => Ok(TxType::Empty),
             _ => Err(format!("invalid tx type: {}", s)),
         }
@@ -69,53 +71,69 @@ impl Display for TxType {
         let str = match self {
             TxType::Empty => "empty".to_string(),
             TxType::Transfer => "transfer".to_string(),
-            TxType::BtcBlock => "btc_blk".to_string(),
+            TxType::BtcBlock => "btc_block".to_string(),
+            TxType::BtcTx => "btc_tx".to_string(),
         };
         write!(f, "{}", str)
     }
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize, Parser, Eq)]
 pub struct BenchTxConfig {
-    pub tx_type: Option<TxType>, // empty(default)/transfer/btc-block
-    pub data_import_flag: bool,
+    pub tx_type: Option<TxType>,       // empty(default)/transfer/btc-block
     pub btc_block_dir: Option<String>, // btc block dir, file name: <height>.hex
+    pub btc_block_start_height: Option<u64>, // btc block start height
+    pub btc_rpc_url: Option<String>,
+    pub btc_rpc_username: Option<String>,
+    pub btc_rpc_password: Option<String>,
     pub pprof_output: Option<PProfOutput>, // flamegraph(default)/proto
 }
 
 impl Default for BenchTxConfig {
     fn default() -> Self {
         Self {
             tx_type: Some(TxType::Empty),
-            data_import_flag: false,
-            btc_block_dir: None,
+            btc_block_dir: Some("target/btc_blocks".to_string()),
+            btc_block_start_height: Some(820000),
+            btc_rpc_url: None,
+            btc_rpc_username: None,
+            btc_rpc_password: None,
             pprof_output: Some(PProfOutput::Flamegraph),
         }
     }
 }
 
 impl BenchTxConfig {
-    pub fn adjust(&mut self) {
-        self.tx_type.get_or_insert(TxType::Empty);
-        self.data_import_flag = false;
-        // if tx_type is btc_block, btc_block_dir must be existed, if not, panic
-        if self.tx_type == Some(TxType::BtcBlock) {
-            self.btc_block_dir
-                .as_ref()
-                .expect("btc_block_dir must be existed");
-        }
-        self.pprof_output.get_or_insert(PProfOutput::Flamegraph);
+    pub fn merge(&mut self, config: BenchTxConfig) {
+        if config.tx_type.is_some() {
+            self.tx_type = config.tx_type;
+        }
+        if config.btc_block_dir.is_some() {
+            self.btc_block_dir = config.btc_block_dir;
+        }
+        if config.btc_block_start_height.is_some() {
+            self.btc_block_start_height = config.btc_block_start_height;
+        }
+        if config.btc_rpc_url.is_some() {
+            self.btc_rpc_url = config.btc_rpc_url;
+        }
+        if config.btc_rpc_username.is_some() {
+            self.btc_rpc_username = config.btc_rpc_username;
+        }
+        if config.btc_rpc_password.is_some() {
+            self.btc_rpc_password = config.btc_rpc_password;
+        }
+        if config.pprof_output.is_some() {
+            self.pprof_output = config.pprof_output;
+        }
     }
 
     pub fn load() -> Self {
         let path = &*BENCH_TX_CONFIG_PATH;
         let mut config = BenchTxConfig::default();
         match std::fs::read_to_string(path) {
             Ok(config_data) => match toml::from_str::<BenchTxConfig>(&config_data) {
-                Ok(mut parsed_config) => {
-                    parsed_config.adjust();
-                    config = parsed_config;
-                }
+                Ok(parsed_config) => config.merge(parsed_config),
                 Err(e) => {
                     log::error!("Failed to parse config file: {}", e);
                 }
@@ -124,6 +142,28 @@ impl BenchTxConfig {
                 log::error!("Failed to read config file: {}", e);
             }
         };
+        // Override config with env variables
+        if let Ok(tx_type) = std::env::var("ROOCH_BENCH_TX_TYPE") {
+            config.tx_type = Some(tx_type.parse().unwrap());
+        }
+        if let Ok(btc_block_dir) = std::env::var("ROOCH_BENCH_BTC_BLOCK_DIR") {
+            config.btc_block_dir = Some(btc_block_dir);
+        }
+        if let Ok(btc_block_start_height) = std::env::var("ROOCH_BENCH_BTC_BLOCK_START_HEIGHT") {
+            config.btc_block_start_height = Some(btc_block_start_height.parse().unwrap());
+        }
+        if let Ok(btc_rpc_url) = std::env::var("ROOCH_BENCH_BTC_RPC_URL") {
+            config.btc_rpc_url = Some(btc_rpc_url);
+        }
+        if let Ok(btc_rpc_username) = std::env::var("ROOCH_BENCH_BTC_RPC_USERNAME") {
+            config.btc_rpc_username = Some(btc_rpc_username);
+        }
+        if let Ok(btc_rpc_password) = std::env::var("ROOCH_BENCH_BTC_RPC_PASSWORD") {
+            config.btc_rpc_password = Some(btc_rpc_password);
+        }
+        if let Ok(pprof_output) = std::env::var("ROOCH_BENCH_PPROF_OUTPUT") {
+            config.pprof_output = Some(pprof_output.parse().unwrap());
+        }
         config
     }
 }

crates/rooch-benchmarks/src/tx.rs

@@ -1,10 +1,13 @@
 // Copyright (c) RoochNetwork
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::config::TxType;
+use crate::tx::TxType::{Empty, Transfer};
 use anyhow::Result;
 use bitcoin::consensus::deserialize;
 use bitcoin::hashes::Hash;
 use bitcoin::hex::FromHex;
+use bitcoincore_rpc::RpcApi;
 use bitcoincore_rpc_json::bitcoin;
 use bitcoincore_rpc_json::bitcoin::Block;
 use rooch_sequencer::actor::sequencer::SequencerActor;
@@ -15,9 +18,8 @@ use rooch_types::multichain_id::RoochMultiChainID;
 use rooch_types::transaction::rooch::RoochTransaction;
 use rooch_types::transaction::L1BlockWithBody;
 use std::fs;
-
-use crate::config::TxType;
-use crate::tx::TxType::{Empty, Transfer};
+use std::path::Path;
+use tracing::info;
 
 pub const EXAMPLE_SIMPLE_BLOG_PACKAGE_NAME: &str = "simple_blog";
 pub const EXAMPLE_SIMPLE_BLOG_NAMED_ADDRESS: &str = "simple_blog";
@@ -52,7 +54,7 @@ pub fn create_l2_tx(
     test_transaction_builder.build_and_sign(action)
 }
 
-pub fn find_block_height(dir: String) -> Result<Vec<u64>> {
+pub fn find_block_height(dir: &Path) -> Result<Vec<u64>> {
     let mut block_heights = Vec::new();
 
     for entry in fs::read_dir(dir)? {
@@ -71,7 +73,7 @@ pub fn find_block_height(dir: String) -> Result<Vec<u64>> {
     Ok(block_heights)
 }
 
-pub fn create_btc_blk_tx(height: u64, block_file: String) -> Result<L1BlockWithBody> {
+pub fn create_btc_blk_tx(height: u64, block_file: &Path) -> Result<L1BlockWithBody> {
     let block_hex_str = fs::read_to_string(block_file).unwrap();
     let block_hex = Vec::<u8>::from_hex(&block_hex_str).unwrap();
     let origin_block: Block = deserialize(&block_hex).unwrap();
@@ -87,3 +89,38 @@ pub fn create_btc_blk_tx(height: u64, block_file: String) -> Result<L1BlockWithB
         block_body: move_block.encode(),
     })
 }
+
+// Download btc block data via bitcoin client
+pub fn prepare_btc_block(
+    btc_block_dir: &Path,
+    btc_rpc_url: String,
+    btc_rpc_username: String,
+    btc_rpc_password: String,
+    btc_block_start_height: u64,
+    btc_block_count: u64,
+) {
+    if !btc_block_dir.exists() {
+        fs::create_dir_all(btc_block_dir).unwrap();
+    }
+
+    let client = bitcoincore_rpc::Client::new(
+        btc_rpc_url.as_str(),
+        bitcoincore_rpc::Auth::UserPass(btc_rpc_username, btc_rpc_password),
+    )
+    .unwrap();
+
+    for i in 0..btc_block_count {
+        let height = btc_block_start_height + i;
+        let filename = format!("{}.hex", height);
+        let file_path = btc_block_dir.join(filename);
+
+        if file_path.exists() {
+            continue;
+        }
+
+        let block_hash = client.get_block_hash(height).unwrap();
+        let block_hex = client.get_block_hex(&block_hash).unwrap();
+        info!("Downloaded block {} to {}", height, file_path.display());
+        fs::write(file_path, block_hex).unwrap();
+    }
+}