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())
@@ -134,7 +137,7 @@ 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) => ProgressBar::new(fs::metadata(source)?.len()),
                 ImportSource::Stdin => ProgressBar::new_spinner(),
             };
             let style = ProgressStyle::default_bar()
@@ -148,10 +151,12 @@ impl Import {
         const BLOCKS_COUNT_PER_CHUNK: usize = 1024 * 6;
         let (blocks_tx, blocks_rx) = ckb_channel::bounded::<Arc<BlockView>>(BLOCKS_COUNT_PER_CHUNK);
         std::thread::spawn({
+            let num_threads = self.num_threads;
             #[cfg(feature = "progress_bar")]
             let progress_bar = progress_bar.clone();
             move || {
                 let pool = rayon::ThreadPoolBuilder::new()
+                    .num_threads(num_threads)
                     .build()
                     .expect("rayon thread pool must build");
                 pool.install(|| {