优化gofile上传，防止一次性读取大文件至内存

Chikage0o0 · Oct 16, 2023 · 3f46948 · 3f46948
1 parent 5c548da
commit 3f46948
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 10 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/matrix_bot/src/main.rs b/matrix_bot/src/main.rs
@@ -61,7 +61,10 @@ async fn main() {
     let _ = load_plugins(&matrix_client, args.data.join("plugins"), &args.plugins);
 
     let ctrlc = tokio::signal::ctrl_c();
-    let mut term= tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).unwrap();
+
+    #[cfg(unix)]
+    let mut term =
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).unwrap();
 
     let client = matrix_client.clone();
 
@@ -79,7 +82,7 @@ async fn main() {
         _=handle => {
             log::error!("Syncing stopped");
         }
-    
+
         _=term.recv() => {
             log::info!("SIGTERM received, stopping");
         }

diff --git a/plugins/qbittorrent/Cargo.toml b/plugins/qbittorrent/Cargo.toml
@@ -18,10 +18,12 @@ reqwest = { version = "0.11.6", default-features = false, features = [
     "json",
     "trust-dns",
     "rustls-tls-native-roots",
+    "stream",
 ] }
 regex = "1.10.0"
 once_cell = "1.8.0"
 url = "2.4.1"
+tokio-util = { version = "0.7.9", default-features = false, features = ["io"] }
 serde_json = "1.0.68"
 walkdir = "2.4.0"
 

diff --git a/plugins/qbittorrent/src/upload/gofile.rs b/plugins/qbittorrent/src/upload/gofile.rs
@@ -1,6 +1,9 @@
 use anyhow::Result;
+use reqwest::Body;
 use serde::{Deserialize, Serialize};
-use std::{borrow::Cow, path::Path};
+use std::path::Path;
+use tokio::fs::File;
+use tokio_util::io::ReaderStream;
 use walkdir::WalkDir;
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -38,12 +41,13 @@ pub async fn upload<P: AsRef<Path>>(file_path: P) -> Result<String> {
             })
         {
             let file_name = entry.file_name().to_string_lossy().to_string();
-            let file = std::fs::read(entry.path())?;
+            let file = File::open(entry.path()).await?;
+            let file_stream = ReaderStream::new(file);
             let upload = UploadParams {
                 server: server.clone(),
                 token: token.clone(),
                 folder_id: folder_id.clone(),
-                file,
+                file_stream,
                 file_name,
             };
             let upload = upload_file(upload).await?;
@@ -69,12 +73,13 @@ pub async fn upload<P: AsRef<Path>>(file_path: P) -> Result<String> {
             .unwrap()
             .to_string_lossy()
             .to_string();
-        let file = std::fs::read(file_path)?;
+        let file = File::open(file_path).await?;
+        let file_stream = ReaderStream::new(file);
         let upload = UploadParams {
             server,
             token: None,
             folder_id: None,
-            file,
+            file_stream,
             file_name,
         };
         let upload = upload_file(upload).await?;
@@ -87,7 +92,7 @@ struct UploadParams {
     server: String,
     token: Option<String>,
     folder_id: Option<String>,
-    file: Vec<u8>,
+    file_stream: ReaderStream<File>,
     file_name: String,
 }
 
@@ -113,8 +118,8 @@ async fn upload_file(parameter: UploadParams) -> Result<UploadInfo> {
     let url = format!("https://{}.gofile.io/uploadFile", parameter.server);
 
     // 一定要有file_name方法，且参数不能为空，否则数据上传失败
-    let part =
-        reqwest::multipart::Part::bytes(Cow::from(parameter.file)).file_name(parameter.file_name);
+    let part = reqwest::multipart::Part::stream(Body::wrap_stream(parameter.file_stream))
+        .file_name(parameter.file_name);
     let mut form = reqwest::multipart::Form::new().part("file", part);
     if let Some(token) = parameter.token {
         form = form.text("token", token);