From b658ce7e75e85eb00ccb8109543d7346696f8907 Mon Sep 17 00:00:00 2001 From: huzibaca Date: Wed, 27 Nov 2024 07:34:34 +0800 Subject: [PATCH] refactor: backup implementation 1. timeouts can be set for different operations 2. performance optimization --- src-tauri/src/core/backup.rs | 204 +++++++++++++++++++++++------------ 1 file changed, 137 insertions(+), 67 deletions(-) diff --git a/src-tauri/src/core/backup.rs b/src-tauri/src/core/backup.rs index db36f494..4a2215ac 100644 --- a/src-tauri/src/core/backup.rs +++ b/src-tauri/src/core/backup.rs @@ -4,114 +4,184 @@ use anyhow::Error; use once_cell::sync::OnceCell; use parking_lot::Mutex; use reqwest_dav::list_cmd::{ListEntity, ListFile}; +use std::collections::HashMap; use std::env::{consts::OS, temp_dir}; use std::fs; use std::io::Write; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; +use tokio::time::timeout; use zip::write::SimpleFileOptions; +const TIMEOUT_UPLOAD: u64 = 300; // 上传超时 5 分钟 +const TIMEOUT_DOWNLOAD: u64 = 300; // 下载超时 5 分钟 +const TIMEOUT_LIST: u64 = 3; // 列表超时 30 秒 +const TIMEOUT_DELETE: u64 = 3; // 删除超时 30 秒 + +#[derive(Clone)] +struct WebDavConfig { + url: String, + username: String, + password: String, +} + +#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)] +enum Operation { + Upload, + Download, + List, + Delete, +} + +impl Operation { + fn timeout(&self) -> u64 { + match self { + Operation::Upload => TIMEOUT_UPLOAD, + Operation::Download => TIMEOUT_DOWNLOAD, + Operation::List => TIMEOUT_LIST, + Operation::Delete => TIMEOUT_DELETE, + } + } +} + pub struct WebDavClient { - client: Arc>>, + config: Arc>>, + clients: Arc>>, } impl WebDavClient { pub fn global() -> &'static WebDavClient { static WEBDAV_CLIENT: OnceCell = OnceCell::new(); WEBDAV_CLIENT.get_or_init(|| WebDavClient { - client: Arc::new(Mutex::new(None)), + config: Arc::new(Mutex::new(None)), + clients: Arc::new(Mutex::new(HashMap::new())), }) } - async fn get_client(&self) -> Result { - if self.client.lock().is_none() { - let verge = Config::verge().latest().clone(); - if verge.webdav_url.is_none() - || verge.webdav_username.is_none() - || verge.webdav_password.is_none() - { - let msg = - "Unable to create web dav client, please make sure the webdav config is correct" - .to_string(); - log::error!(target: "app","{}",msg); - return Err(anyhow::Error::msg(msg)); + async fn get_client(&self, op: Operation) -> Result { + // 先尝试从缓存获取 + { + let clients = self.clients.lock(); + if let Some(client) = clients.get(&op) { + return Ok(client.clone()); } - - let url = verge.webdav_url.unwrap_or_default(); - let username = verge.webdav_username.unwrap_or_default(); - let password = verge.webdav_password.unwrap_or_default(); - let url = url.trim_end_matches('/'); - let client = reqwest_dav::ClientBuilder::new() - .set_agent( - reqwest::Client::builder() - .danger_accept_invalid_certs(true) - .timeout(std::time::Duration::from_secs(15)) - .build() - .unwrap(), - ) - .set_host(url.to_owned()) - .set_auth(reqwest_dav::Auth::Basic( - username.to_owned(), - password.to_owned(), - )) - .build()?; - - if (client - .list(dirs::BACKUP_DIR, reqwest_dav::Depth::Number(0)) - .await) - .is_err() - { - client.mkcol(dirs::BACKUP_DIR).await?; - } - - *self.client.lock() = Some(client.clone()); } - Ok(self.client.lock().clone().unwrap()) + + // 获取或创建配置 + let config = { + let mut lock = self.config.lock(); + if let Some(cfg) = lock.as_ref() { + cfg.clone() + } else { + let verge = Config::verge().latest().clone(); + if verge.webdav_url.is_none() + || verge.webdav_username.is_none() + || verge.webdav_password.is_none() + { + let msg = "Unable to create web dav client, please make sure the webdav config is correct".to_string(); + return Err(anyhow::Error::msg(msg)); + } + + let config = WebDavConfig { + url: verge + .webdav_url + .unwrap_or_default() + .trim_end_matches('/') + .to_string(), + username: verge.webdav_username.unwrap_or_default(), + password: verge.webdav_password.unwrap_or_default(), + }; + + *lock = Some(config.clone()); + config + } + }; + + // 创建新的客户端 + let client = reqwest_dav::ClientBuilder::new() + .set_agent( + reqwest::Client::builder() + .danger_accept_invalid_certs(true) + .timeout(Duration::from_secs(op.timeout())) + .build() + .unwrap(), + ) + .set_host(config.url) + .set_auth(reqwest_dav::Auth::Basic(config.username, config.password)) + .build()?; + + // 确保备份目录存在 + let list_result = client + .list(dirs::BACKUP_DIR, reqwest_dav::Depth::Number(0)) + .await; + if list_result.is_err() { + client.mkcol(dirs::BACKUP_DIR).await?; + } + + // 缓存客户端 + { + let mut clients = self.clients.lock(); + clients.insert(op, client.clone()); + } + + Ok(client) } pub fn reset(&self) { - if !self.client.lock().is_none() { - self.client.lock().take(); - } + *self.config.lock() = None; + self.clients.lock().clear(); } pub async fn upload(&self, file_path: PathBuf, file_name: String) -> Result<(), Error> { - let client = self.get_client().await?; + let client = self.get_client(Operation::Upload).await?; let webdav_path: String = format!("{}/{}", dirs::BACKUP_DIR, file_name); - client - .put(webdav_path.as_ref(), fs::read(file_path)?) - .await?; + let fut = client.put(webdav_path.as_ref(), fs::read(file_path)?); + timeout(Duration::from_secs(TIMEOUT_UPLOAD), fut).await??; Ok(()) } pub async fn download(&self, filename: String, storage_path: PathBuf) -> Result<(), Error> { - let client = self.get_client().await?; + let client = self.get_client(Operation::Download).await?; let path = format!("{}/{}", dirs::BACKUP_DIR, filename); - let response = client.get(path.as_str()).await?; - let content = response.bytes().await?; - fs::write(&storage_path, &content)?; + + let fut = async { + let response = client.get(path.as_str()).await?; + let content = response.bytes().await?; + fs::write(&storage_path, &content)?; + Ok::<(), Error>(()) + }; + + timeout(Duration::from_secs(TIMEOUT_DOWNLOAD), fut).await??; Ok(()) } pub async fn list(&self) -> Result, Error> { - let client = self.get_client().await?; + let client = self.get_client(Operation::List).await?; let path = format!("{}/", dirs::BACKUP_DIR); - let files = client - .list(path.as_str(), reqwest_dav::Depth::Number(1)) - .await?; - let mut final_files = Vec::new(); - for file in files { - if let ListEntity::File(file) = file { - final_files.push(file); + + let fut = async { + let files = client + .list(path.as_str(), reqwest_dav::Depth::Number(1)) + .await?; + let mut final_files = Vec::new(); + for file in files { + if let ListEntity::File(file) = file { + final_files.push(file); + } } - } - Ok(final_files) + Ok::, Error>(final_files) + }; + + Ok(timeout(Duration::from_secs(TIMEOUT_LIST), fut).await??) } pub async fn delete(&self, file_name: String) -> Result<(), Error> { - let client = self.get_client().await?; + let client = self.get_client(Operation::Delete).await?; let path = format!("{}/{}", dirs::BACKUP_DIR, file_name); - client.delete(&path).await?; + + let fut = client.delete(&path); + timeout(Duration::from_secs(TIMEOUT_DELETE), fut).await??; Ok(()) } }