添加下载功能

This commit is contained in:
柚子
2025-03-04 04:10:19 +08:00
parent 2e963b65e2
commit aadc137055
10 changed files with 682 additions and 241 deletions

View File

@@ -1,61 +1,203 @@
use std::collections::HashMap;
use std::sync::Mutex;
use crate::models::download::{DownloadStatus, DownloadTask};
use crate::utils::metalink::parse_metalink;
use std::process::Command;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::net::TcpListener;
use crate::models::download::{DownloadStatus, DownloadTask, DownloadTaskResponse};
use crate::handlers::server::get_json_server_url;
use crate::utils::{format_icon_url, UA};
use crate::utils::{format_icon_url, UA, aria2::Aria2Client};
use super::format::{format_size, format_speed};
pub struct DownloadManager {
queue: Mutex<HashMap<String, DownloadTask>>,
aria2_started: Arc<AtomicBool>,
aria2_port: Arc<Mutex<u16>>,
}
impl DownloadManager {
pub fn new() -> Self {
DownloadManager {
queue: Mutex::new(HashMap::new()),
aria2_started: Arc::new(AtomicBool::new(false)),
aria2_port: Arc::new(Mutex::new(5144)),
}
}
// 获取所有下载任务
pub fn get_downloads(&self) -> Result<Vec<DownloadTask>, String> {
let downloads = self.queue.lock().map_err(|e| e.to_string())?;
Ok(downloads.values().cloned().collect())
}
// 检查是否有正在下载的任务
fn has_downloading_task(&self) -> Result<bool, String> {
let downloads = self.queue.lock().map_err(|e| e.to_string())?;
Ok(downloads.values().any(|task| matches!(task.status, DownloadStatus::Downloading)))
}
// 获取下一个等待下载的任务
fn get_next_queued_task(&self) -> Result<Option<String>, String> {
let downloads = self.queue.lock().map_err(|e| e.to_string())?;
Ok(downloads
.iter()
.find(|(_, task)| matches!(task.status, DownloadStatus::Queued))
.map(|(task_id, _)| task_id.clone()))
}
// 开始下载下一个任务
fn start_next_download(&self) -> Result<(), String> {
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
if let Some(task_id) = self.get_next_queued_task()? {
if let Some(task) = downloads.get_mut(&task_id) {
task.status = DownloadStatus::Downloading;
pub async fn get_downloads(&self) -> Result<Vec<DownloadTaskResponse>, String> {
// 获取队列中的任务信息并立即克隆所需数据,然后释放锁
let tasks_clone: Vec<DownloadTask>;
let aria2_started;
let port;
{
// 使用作用域限制锁的生命周期
let downloads = self.queue.lock().map_err(|e| e.to_string())?;
tasks_clone = downloads.values().cloned().collect();
aria2_started = self.aria2_started.load(Ordering::SeqCst);
}
// 如果 aria2 未启动,直接返回队列中的任务
if !aria2_started {
return Ok(tasks_clone.into_iter().map(|task| DownloadTaskResponse {
category: task.category,
pkgname: task.pkgname,
filename: task.filename,
status: DownloadStatus::Error, // 如果 aria2 未启动,标记为错误状态
icon: task.icon,
name: task.name,
progress: 0.0,
speed: None,
size: None,
}).collect());
}
// 获取端口(在单独的作用域中获取锁)
{
port = *self.aria2_port.lock().map_err(|e| e.to_string())?;
}
// 创建 Aria2Client 实例
let aria2_client = Aria2Client::new("127.0.0.1", port, None);
// 获取所有活动中的下载任务
let active_downloads = aria2_client.tell_active().await
.map_err(|e| format!("获取活动下载任务失败: {}", e))?;
// 获取所有等待中的下载任务最多100个
let waiting_downloads = aria2_client.tell_waiting(0, 100).await
.map_err(|e| format!("获取等待中下载任务失败: {}", e))?;
// 获取所有已停止的下载任务最多100个
let stopped_downloads = aria2_client.tell_stopped(0, 100).await
.map_err(|e| format!("获取已停止下载任务失败: {}", e))?;
// 创建一个映射,用于存储 GID 到 aria2 任务状态的映射
let mut aria2_tasks = HashMap::new();
// 处理活动中的下载任务
for task in active_downloads {
if let Some(gid) = task["gid"].as_str() {
aria2_tasks.insert(gid.to_string(), (task, DownloadStatus::Downloading));
}
}
Ok(())
// 处理等待中的下载任务
for task in waiting_downloads {
if let Some(gid) = task["gid"].as_str() {
let status = if task["status"].as_str() == Some("paused") {
DownloadStatus::Paused
} else {
DownloadStatus::Queued
};
aria2_tasks.insert(gid.to_string(), (task, status));
}
}
// 处理已停止的下载任务
for task in stopped_downloads {
if let Some(gid) = task["gid"].as_str() {
let status = if task["status"].as_str() == Some("complete") {
DownloadStatus::Completed
} else {
DownloadStatus::Error
};
aria2_tasks.insert(gid.to_string(), (task, status));
}
}
// 将队列中的任务与 aria2 任务状态结合,生成响应
let mut result = Vec::new();
for task in tasks_clone {
let mut response = DownloadTaskResponse {
category: task.category.clone(),
pkgname: task.pkgname.clone(),
filename: task.filename.clone(),
status: DownloadStatus::Error, // 默认为错误状态
icon: task.icon.clone(),
name: task.name.clone(),
progress: 0.0,
speed: None,
size: None,
};
// 如果在 aria2 任务中找到对应的 GID更新状态信息
if let Some((aria2_task, status)) = aria2_tasks.get(&task.gid) {
response.status = status.clone();
// 计算进度(百分比)
if let (Some(completed_length), Some(total_length)) = (
aria2_task["completedLength"].as_str().and_then(|s| s.parse::<f64>().ok()),
aria2_task["totalLength"].as_str().and_then(|s| s.parse::<f64>().ok())
) {
if total_length > 0.0 {
// 保留两位小数
let progress = ((completed_length / total_length) * 100.0) as f32;
response.progress = (progress * 100.0).round() / 100.0;
}
}
// 获取下载速度
if let Some(download_speed) = aria2_task["downloadSpeed"].as_str() {
if let Ok(speed) = download_speed.parse::<u64>() {
response.speed = Some(format_speed(speed));
}
}
// 获取文件大小
if let Some(total_length) = aria2_task["totalLength"].as_str() {
if let Ok(size) = total_length.parse::<u64>() {
response.size = Some(format_size(size));
}
}
}
result.push(response);
}
Ok(result)
}
fn start_aria2(&self) {
// 寻找可用端口
let mut port = 5144;
while TcpListener::bind(format!("127.0.0.1:{}", port)).is_err() {
port += 1;
}
if let Ok(mut port_guard) = self.aria2_port.lock() {
*port_guard = port;
}
// 启动 aria2c
Command::new("aria2c")
.args([
"--enable-rpc",
"--rpc-listen-all=false",
&format!("--rpc-listen-port={}", port),
&format!("--user-agent={}", UA),
])
.spawn()
.map_err(|e| format!("启动 aria2 失败: {}", e)).unwrap();
self.aria2_started.store(true, Ordering::SeqCst);
}
// 添加下载任务
pub async fn add_download(&self, category: String, pkgname: String, filename: String, name: String) -> Result<(), String> {
// 检查并启动 aria2如果还没启动
if !self.aria2_started.load(Ordering::SeqCst) {
self.start_aria2();
}
let task_id = format!("{}/{}", category, pkgname);
// 获取metalink文件URL和内容在获取锁之前完成
let json_server_url = get_json_server_url();
let metalink_url = format!("{}{}/{}/{}.metalink", json_server_url, category, pkgname, filename);
// 发送请求获取metalink文件
let client = reqwest::Client::new();
let response = client
@@ -70,14 +212,24 @@ impl DownloadManager {
return Err(format!("获取metalink文件失败: HTTP {}", response.status()));
}
// 解析 metalink 内容
let metalink_content = response.text().await
.map_err(|e| format!("读取metalink内容失败: {}", e))?;
// 获取metalink文件内容
let metalink_content = response.bytes()
.await
.map_err(|e| format!("读取metalink内容失败: {}", e))?;
let metalink_data = parse_metalink(&metalink_content)?;
// 创建Aria2Client并添加下载任务
let port = *self.aria2_port.lock().map_err(|e| e.to_string())?;
let aria2_client = Aria2Client::new("127.0.0.1", port, None);
// 使用Aria2Client添加metalink下载
let gids = aria2_client.add_metalink(
&metalink_content,
None,
None
)
.await
.map_err(|e| format!("添加下载任务失败: {}", e))?;
println!("{:?}", metalink_data);
// 获取一次锁,完成所有状态更新操作
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
@@ -85,15 +237,7 @@ impl DownloadManager {
if downloads.contains_key(&task_id) {
return Ok(());
}
// 检查是否有正在下载的任务
let has_downloading = downloads.values().any(|task| matches!(task.status, DownloadStatus::Downloading));
let initial_status = if has_downloading {
DownloadStatus::Queued
} else {
DownloadStatus::Downloading
};
// 创建下载任务
let task = DownloadTask {
category: category.clone(),
@@ -101,11 +245,7 @@ impl DownloadManager {
filename,
name,
icon: format_icon_url(&category, &pkgname),
status: initial_status,
progress: 0.0,
speed: None,
size: None,
metalink: metalink_data,
gid: gids[0].clone()
};
// 添加任务到队列
@@ -115,64 +255,95 @@ impl DownloadManager {
}
// 暂停下载任务
pub fn pause_download(&self, category: String, pkgname: String) -> Result<(), String> {
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
pub async fn pause_download(&self, category: String, pkgname: String) -> Result<(), String> {
let task_id = format!("{}/{}", category, pkgname);
if let Some(task) = downloads.get_mut(&task_id) {
let old_status = task.status.clone();
task.status = DownloadStatus::Paused;
if matches!(old_status, DownloadStatus::Downloading) {
self.start_next_download()?;
// 获取任务信息,并在作用域结束时释放锁
let task_gid = {
let downloads = self.queue.lock().map_err(|e| e.to_string())?;
match downloads.get(&task_id) {
Some(task) => task.gid.clone(),
None => return Err(format!("找不到下载任务: {}", task_id)),
}
};
// 如果 aria2 未启动,返回错误
if !self.aria2_started.load(Ordering::SeqCst) {
return Err("aria2 未启动".to_string());
}
// 创建 Aria2Client 实例
let port = *self.aria2_port.lock().map_err(|e| e.to_string())?;
let aria2_client = Aria2Client::new("127.0.0.1", port, None);
// 调用 aria2 的 pause 方法
aria2_client.pause(task_gid.as_str()).await
.map_err(|e| format!("暂停下载任务失败: {}", e))?;
Ok(())
}
// 恢复下载任务
pub fn resume_download(&self, category: String, pkgname: String) -> Result<(), String> {
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
pub async fn resume_download(&self, category: String, pkgname: String) -> Result<(), String> {
let task_id = format!("{}/{}", category, pkgname);
// 先检查任务是否存在且处于暂停状态
let should_resume = downloads
.get(&task_id)
.map(|task| matches!(task.status, DownloadStatus::Paused))
.unwrap_or(false);
if should_resume {
// 检查是否有其他正在下载的任务
let has_downloading = self.has_downloading_task()?;
// 更新任务状态
if let Some(task) = downloads.get_mut(&task_id) {
task.status = if has_downloading {
DownloadStatus::Queued
} else {
DownloadStatus::Downloading
};
// 获取任务信息,并在作用域结束时释放锁
let task_gid = {
let downloads = self.queue.lock().map_err(|e| e.to_string())?;
match downloads.get(&task_id) {
Some(task) => task.gid.clone(),
None => return Err(format!("找不到下载任务: {}", task_id)),
}
};
// 如果 aria2 未启动,返回错误
if !self.aria2_started.load(Ordering::SeqCst) {
return Err("aria2 未启动".to_string());
}
// 创建 Aria2Client 实例
let port = *self.aria2_port.lock().map_err(|e| e.to_string())?;
let aria2_client = Aria2Client::new("127.0.0.1", port, None);
// 调用 aria2 的 unpause 方法
aria2_client.unpause(task_gid.as_str()).await
.map_err(|e| format!("恢复下载任务失败: {}", e))?;
Ok(())
}
// 取消下载任务
pub fn cancel_download(&self, category: String, pkgname: String) -> Result<(), String> {
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
pub async fn cancel_download(&self, category: String, pkgname: String) -> Result<(), String> {
let task_id = format!("{}/{}", category, pkgname);
let was_downloading = downloads
.get(&task_id)
.map(|task| matches!(task.status, DownloadStatus::Downloading))
.unwrap_or(false);
// 获取任务信息,并在作用域结束时释放锁
let task_gid = {
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
match downloads.get(&task_id) {
Some(task) => {
// 如果 aria2 未启动,只从队列中移除任务
if !self.aria2_started.load(Ordering::SeqCst) {
downloads.remove(&task_id);
return Ok(());
}
task.gid.clone()
},
None => return Err(format!("找不到下载任务: {}", task_id)),
}
};
// 创建 Aria2Client 实例
let port = *self.aria2_port.lock().map_err(|e| e.to_string())?;
let aria2_client = Aria2Client::new("127.0.0.1", port, None);
// 调用 aria2 的 remove 方法
aria2_client.remove(task_gid.as_str()).await
.map_err(|e| format!("取消下载任务失败: {}", e))?;
// 从队列中移除任务
let mut downloads = self.queue.lock().map_err(|e| e.to_string())?;
downloads.remove(&task_id);
if was_downloading {
self.start_next_download()?;
}
Ok(())
}
}