From 3675e8ec5de12bcf995a60842f05d3a325f1be5d Mon Sep 17 00:00:00 2001
From: shenjack <3695888@qq.com>
Date: Sun, 21 Jul 2024 00:01:02 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9C=80=E5=90=8E=E7=9A=84=E4=B8=80=E4=BA=9B?=
 =?UTF-8?q?=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

其实后面估计还会有一些空位检测优化
---
 sr_download/src/config.rs  |  4 ++++
 sr_download/src/db_part.rs |  4 ++--
 sr_download/src/main.rs    | 11 ++++++-----
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/sr_download/src/config.rs b/sr_download/src/config.rs
index 74639fc..dcbce42 100644
--- a/sr_download/src/config.rs
+++ b/sr_download/src/config.rs
@@ -8,6 +8,8 @@ pub struct ConfigFile {
     pub db_schema: String,
     pub max_connections: u32,
     pub sqlx_logging: bool,
+    pub worker_count: u32,
+    pub worker_size: u32,
 }
 
 impl Default for ConfigFile {
@@ -17,6 +19,8 @@ impl Default for ConfigFile {
             db_schema: "public".to_string(),
             max_connections: 10,
             sqlx_logging: true,
+            worker_count: 10,
+            worker_size: 10,
         }
     }
 }
diff --git a/sr_download/src/db_part.rs b/sr_download/src/db_part.rs
index ec3cb4c..7ad5d7a 100644
--- a/sr_download/src/db_part.rs
+++ b/sr_download/src/db_part.rs
@@ -29,8 +29,8 @@ pub async fn find_max_id(db: &DatabaseConnection) -> SaveId {
     // 我丢你老母, 有这时间写这个, 我都写完 sql 语句了
     match model::main_data::Entity::find()
         .order_by_desc(model::main_data::Column::SaveId)
-        .select_only()
-        .column(model::main_data::Column::SaveId)
+        // .select_only()
+        // .column(model::main_data::Column::SaveId)
         .one(db)
         .await
     {
diff --git a/sr_download/src/main.rs b/sr_download/src/main.rs
index 1817bf4..d5e99b4 100644
--- a/sr_download/src/main.rs
+++ b/sr_download/src/main.rs
@@ -43,7 +43,7 @@ async fn big_worker(db: sea_orm::DatabaseConnection, work_range: Range<SaveId>)
     }
 }
 
-#[tokio::main]
+#[tokio::main(flavor = "multi_thread", worker_threads = 10)]
 async fn main() -> anyhow::Result<()> {
     tracing_subscriber::fmt().with_max_level(Level::INFO).init();
     event!(Level::INFO, "Starting srdownload");
@@ -66,11 +66,11 @@ async fn main() -> anyhow::Result<()> {
 
     let mut current_id = start_id;
 
-    let batch_size = 100;
+    let batch_size = conf.worker_size;
     // 10 works
-    let mut works = Vec::with_capacity(10);
-    let max_works = 10;
-    for _ in 0..10 {
+    let mut works = Vec::with_capacity(conf.worker_count as usize);
+    let max_works = conf.worker_count as usize;
+    for _ in 0..works.len() {
         let end = current_id + batch_size;
         works.push(tokio::spawn(big_worker(
             db_connect.clone(),
@@ -78,6 +78,7 @@ async fn main() -> anyhow::Result<()> {
         )));
         current_id = end;
     }
+
     while current_id < end_id || !works.is_empty() {
         while current_id < end_id && works.len() < max_works {
             let end = current_id + batch_size;