feat: add thread pool limitation for rayon

4o3F · Oct 8, 2024 · 0c0bbaa · 0c0bbaa
1 parent 35082a8
commit 0c0bbaa
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -13,6 +13,9 @@ struct Cli {
     /// Commands
     #[command(subcommand)]
     command: Option<Commands>,
+
+    #[arg(long, default_value = "100", help = "Thread pool size")]
+    thread: usize,
 }
 
 #[derive(Subcommand)]
@@ -330,8 +333,6 @@ enum YoloCommands {
 
 #[tokio::main]
 async fn main() {
-    // env_logger::init();
-
     // Do tracing init
     let subscriber = tracing_subscriber::fmt()
         .with_max_level(Level::TRACE)
@@ -342,6 +343,11 @@ async fn main() {
 
     let cli = Cli::parse();
 
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(cli.thread)
+        .build_global()
+        .unwrap();
+
     match &cli.command {
         Some(Commands::Common { command }) => match command {
             CommonCommands::CropRectangle {