import: Add --num-threads to control parallel threads number

Signed-off-by: Eval EXEC <[email protected]>

Author: eval-exec

Cargo.lock

@@ -277,7 +277,6 @@ quote = "1.0"
 rand = "0.8.5"
 rand_distr = "0.4"
 rayon = "1.0"
-rayon-progress-bar = "1.9"
 regex = "1.1.6"
 reqwest = { version = "0.12", features = ["blocking", "json"] }
 rhai = "1.16.0"

Cargo.toml

@@ -35,6 +35,8 @@ pub const ARG_FORMAT: &str = "format";
 pub const ARG_TARGET: &str = "target";
 /// Command line argument `--source`.
 pub const ARG_SOURCE: &str = "source";
+/// Command line argument `--num-threads`.
+pub const ARG_NUM_THREADS: &str = "num-threads";
 /// Command line flag: `--skip-script-verify`.
 pub const ARG_SKIP_SCRIPT_VERIFY: &str = "skip-script-verify";
 /// Command line flag: `--skip-all-verify`.
@@ -410,6 +412,14 @@ fn export() -> Command {
 fn import() -> Command {
     Command::new(CMD_IMPORT)
         .about("Import CKB data")
+        .arg(
+            Arg::new(ARG_NUM_THREADS)
+                .long(ARG_NUM_THREADS)
+                .value_name(ARG_NUM_THREADS)
+                .value_parser(clap::value_parser!(usize))
+                .required(false)
+                .help("Specify the number of threads to use for parallel processing. If not specified, it will use the number of logical CPUs available on the system."),
+        )
         .arg(
             Arg::new(ARG_SOURCE)
                 .index(1)

ckb-bin/src/cli.rs

@@ -215,6 +215,10 @@ H256::from_str(&target[2..]).expect("default assume_valid_target for testnet mus
     pub fn import(self, matches: &ArgMatches) -> Result<ImportArgs, ExitCode> {
         let consensus = self.consensus()?;
         let config = self.config.into_ckb()?;
+        let num_threads = matches
+            .get_one::<usize>(cli::ARG_NUM_THREADS)
+            .cloned()
+            .unwrap_or(0);
         let source = {
             let source = matches
                 .get_one::<String>(cli::ARG_SOURCE)
@@ -252,6 +256,7 @@ H256::from_str(&target[2..]).expect("default assume_valid_target for testnet mus
             consensus,
             source,
             switch,
+            num_threads,
         })
     }
 

ckb-bin/src/setup.rs

@@ -20,10 +20,16 @@ pub fn import(args: ImportArgs, async_handle: Handle) -> Result<(), ExitCode> {
     pack.take_tx_pool_builder();
     pack.take_relay_tx_receiver();
 
-    Import::new(chain_controller, shared, args.source, args.switch)
-        .execute()
-        .map_err(|err| {
-            eprintln!("Import error: {err:?}");
-            ExitCode::Failure
-        })
+    Import::new(
+        chain_controller,
+        shared,
+        args.source,
+        args.switch,
+        args.num_threads,
+    )
+    .execute()
+    .map_err(|err| {
+        eprintln!("Import error: {err:?}");
+        ExitCode::Failure
+    })
 }

ckb-bin/src/subcommand/import.rs

@@ -8,8 +8,11 @@ use ckb_verification_traits::Switch;
 use std::path::PathBuf;
 use std::sync::Arc;
 
+/// The target directory to save the exported file or stdout.
 pub enum ExportTarget {
+    /// The path to the file to be exported.
     Path(PathBuf),
+    /// Export to stdout.
     Stdout,
 }
 
@@ -38,8 +41,11 @@ pub struct DaemonArgs {
     pub pid_file: PathBuf,
 }
 
+/// The source of the file to be imported.
 pub enum ImportSource {
+    /// The path to the file to be imported.
     Path(PathBuf),
+    /// Import from stdin, the file content must be encoded by base64.
     Stdin,
 }
 
@@ -53,6 +59,8 @@ pub struct ImportArgs {
     pub source: ImportSource,
     /// The switch to control the verification behavior.
     pub switch: Switch,
+    /// The number of threads to use for parallel processing.
+    pub num_threads: usize,
 }
 
 /// Parsed command line arguments for `ckb run`.

util/app-config/src/args.rs

@@ -12,15 +12,13 @@ repository = "https://github.com/nervosnetwork/ckb"
 ckb-types.workspace = true
 ckb-chain.workspace = true
 ckb-store.workspace = true
-ckb-chain-iter.workspace = true
 ckb-shared.workspace = true
 ckb-channel.workspace = true
 ckb-jsonrpc-types.workspace = true
 ckb-verification-traits.workspace = true
 serde_json.workspace = true
 indicatif = { workspace = true, optional = true }
 rayon.workspace = true
-rayon-progress-bar.workspace = true
 itertools.workspace = true
 ckb-app-config.workspace = true
 

util/instrument/Cargo.toml

@@ -21,6 +21,7 @@ pub struct Import {
     shared: Shared,
     chain: ChainController,
     switch: Switch,
+    num_threads: usize,
 }
 
 impl Import {
@@ -30,12 +31,14 @@ impl Import {
         shared: Shared,
         source: ImportSource,
         switch: Switch,
+        num_threads: usize,
     ) -> Self {
         Import {
             chain,
             shared,
             source,
             switch,
+            num_threads,
         }
     }
 
@@ -78,7 +81,7 @@ impl Import {
         }
 
         let f: Box<dyn Read + Send> = match &self.source {
-            ImportSource::Path(source) => Box::new(fs::File::open(&source)?),
+            ImportSource::Path(source) => Box::new(fs::File::open(source)?),
             ImportSource::Stdin => {
                 // read from stdin
                 Box::new(std::io::stdin())
@@ -131,10 +134,12 @@ impl Import {
             }
         }
 
-        #[cfg(feature = "progress_bar")]
         let progress_bar = {
             let bar = match &self.source {
-                ImportSource::Path(source) => ProgressBar::new(fs::metadata(&source)?.len()),
+                ImportSource::Path(source) => {
+                    let file_size = fs::metadata(source)?.len();
+                    ProgressBar::new(file_size)
+                }
                 ImportSource::Stdin => ProgressBar::new_spinner(),
             };
             let style = ProgressStyle::default_bar()
@@ -146,12 +151,13 @@ impl Import {
 
         let mut largest_block_number = 0;
         const BLOCKS_COUNT_PER_CHUNK: usize = 1024 * 6;
-        let (blocks_tx, blocks_rx) = ckb_channel::bounded::<Arc<BlockView>>(BLOCKS_COUNT_PER_CHUNK);
+        let (blocks_tx, blocks_rx) =
+            ckb_channel::bounded::<(Arc<BlockView>, usize)>(BLOCKS_COUNT_PER_CHUNK);
         std::thread::spawn({
-            #[cfg(feature = "progress_bar")]
-            let progress_bar = progress_bar.clone();
+            let num_threads = self.num_threads;
             move || {
                 let pool = rayon::ThreadPoolBuilder::new()
+                    .num_threads(num_threads)
                     .build()
                     .expect("rayon thread pool must build");
                 pool.install(|| {
@@ -168,33 +174,35 @@ impl Import {
                             let block: JsonBlock =
                                 serde_json::from_str(line).expect("parse block from json");
                             let block: Arc<core::BlockView> = Arc::new(block.into());
-                            blocks_tx.send(block).expect("send block to channel");
-
-                            #[cfg(feature = "progress_bar")]
-                            progress_bar.inc(line.len() as u64);
+                            blocks_tx
+                                .send((block, line.len()))
+                                .expect("send block to channel");
                         });
                     }
                     drop(blocks_tx);
                 });
             }
         });
 
-        let callback = |verify_result: VerifyResult| {
-            if let Err(err) = verify_result {
-                eprintln!("Error verifying block: {:?}", err);
-            }
-        };
-
-        for block in blocks_rx {
+        for (block, block_size) in blocks_rx {
             if !block.is_genesis() {
                 use ckb_chain::LonelyBlock;
 
                 largest_block_number = largest_block_number.max(block.number());
 
+                let progress_bar = progress_bar.clone();
+                let callback = Box::new(move |verify_result: VerifyResult| {
+                    if let Err(err) = verify_result {
+                        eprintln!("Error verifying block: {:?}", err);
+                    } else {
+                        progress_bar.inc(block_size as u64);
+                    }
+                });
+
                 let lonely_block = LonelyBlock {
                     block,
                     switch: Some(self.switch),
-                    verify_callback: Some(Box::new(callback)),
+                    verify_callback: Some(callback),
                 };
                 self.chain.asynchronous_process_lonely_block(lonely_block);
             }
@@ -209,7 +217,6 @@ impl Import {
             std::thread::sleep(std::time::Duration::from_secs(1));
         }
 
-        #[cfg(feature = "progress_bar")]
         progress_bar.finish_with_message("done!");
         Ok(())
     }