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) } } -#[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;