From 0c0bbaae4fce4b885e27b1b853930a5d569f5ed0 Mon Sep 17 00:00:00 2001
From: 4o3F <4o3f@proton.me>
Date: Tue, 8 Oct 2024 08:43:22 +0800
Subject: [PATCH] feat: add thread pool limitation for rayon

---
 src/main.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index bf66df2..a3cd1db 100644
--- 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 {