use anyhow::Result;
use bytes::Bytes;
use kanal::{AsyncReceiver, AsyncSender};
use reqwest::{IntoUrl, Url};
use std::collections::VecDeque;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::Mutex;
// 为了在不同平台（Unix/Windows）上使用 pwrite (positional write)
#[cfg(unix)]
use std::os::unix::fs::FileExt;
#[cfg(windows)]
use std::os::windows::fs::FileExt;
// --- 自定义错误类型，方便统一处理 ---
#[derive(Debug, thiserror::Error)]
pub enum DownloadError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("HTTP request error: {0}")]
Reqwest(#[from] reqwest::Error),
#[error("URL parse error: {0}")]
UrlParse(#[from] url::ParseError),
#[error("Concurrency error (task join): {0}")]
Join(#[from] tokio::task::JoinError),
#[error("Channel communication error: {0}")]
KanalSend(#[from] kanal::SendError<(u64, Bytes)>),
#[error("Channel communication error: {0}")]
KanalRecv(#[from] kanal::ReceiveError),
#[error("All download workers failed to start")]
NoWorkers,
}
// 定义一个通用的 Result 类型
pub type Void = Result<(), DownloadError>;
// 定义块大小常量
const CHUNK_SIZE: u64 = 256 * 1024; // 256KB
/// 多线程分块下载器
///
/// # 参数
/// * `filesize` - 目标文件的总大小.
/// * `url_li` - 一个URL列表的迭代器，所有URL都指向同一个文件.
/// * `to_path` - 下载后保存的文件路径.
pub async fn down<U: IntoUrl + Send + 'static>(
filesize: u64,
url_li: impl IntoIterator<Item = U>,
to_path: impl AsRef<Path>,
) -> Void {
// 1. 创建并预分配稀疏文件
let file = Arc::new(std::fs::File::create(to_path)?);
file.set_len(filesize)?;
// 2. 初始化任务队列和数据通道
let total_chunks = (filesize + CHUNK_SIZE - 1) / CHUNK_SIZE;
// 使用 Mutex<VecDeque> 作为任务队列，以支持失败后 "push_first"
let tasks = Arc::new(Mutex::new(VecDeque::from_iter(0..total_chunks)));
// 使用 kanal 作为下载数据块的通信通道
// 容量设为32，可以缓冲32个下载好的块
let (data_tx, data_rx): (AsyncSender<(u64, Bytes)>, AsyncReceiver<(u64, Bytes)>) =
kanal::bounded_async(32);
// 3. 启动文件写入任务
// 这个任务独立运行，负责从通道接收数据并写入文件
let writer_handle = {
let file = file.clone();
let writer_task = async move {
// 我们期望接收 total_chunks 个数据块
for i in 0..total_chunks {
match data_rx.recv().await {
Ok((chunk_id, data)) => {
let offset = chunk_id * CHUNK_SIZE;
let file_clone = file.clone();
// spawn_blocking 用于执行阻塞的IO操作，防止阻塞tokio运行时
tokio::task::spawn_blocking(move || {
#[cfg(unix)]
file_clone.write_all_at(&data, offset)?;
#[cfg(windows)]
file_clone.write_all_at(&data, offset)?;
Ok::<(), std::io::Error>(())
})
.await??; // 第一个 `?` 处理 JoinError, 第二个处理 IO Error
println!("[Writer] Wrote chunk {}/{} (ID: {})", i + 1, total_chunks, chunk_id);
}
Err(e) => {
// 如果通道关闭且仍有块未写入，说明出错了
eprintln!("[Writer] Error receiving chunk data: {}", e);
return Err(DownloadError::from(e));
}
}
}
println!("[Writer] All chunks written. Finishing.");
Ok(())
};
tokio::spawn(writer_task)
};
// 4. 启动下载工人（downloader workers）
let client = reqwest::Client::new();
let mut worker_handles = Vec::new();
for url_input in url_li.into_iter() {
let url = url_input.into_url()?;
let tasks_clone = tasks.clone();
let data_tx_clone = data_tx.clone();
let client_clone = client.clone();
let worker_task = async move {
loop {
// --- 从队列认领任务 ---
let task_id = {
let mut queue = tasks_clone.lock().await;
queue.pop_front() // 从头部取任务
};
if let Some(id) = task_id {
let start = id * CHUNK_SIZE;
let end = std::cmp::min(start + CHUNK_SIZE - 1, filesize - 1);
let range = format!("bytes={}-{}", start, end);
println!("[Worker {:?}] Downloading chunk {} (range: {})", std::thread::current().id(), id, range);
// --- 执行下载 ---
match client_clone.get(url.clone()).header("Range", range).send().await {
Ok(response) if response.status().is_success() => {
match response.bytes().await {
Ok(data) => {
// 发送数据到写入任务
if data_tx_clone.send((id, data)).await.is_err() {
// 如果发送失败，说明写入任务已退出，无需继续
eprintln!("[Worker] Failed to send chunk {} to writer. Aborting.", id);
break;
}
}
Err(e) => {
eprintln!("[Worker] Failed to get bytes for chunk {}: {}. Re-queueing.", id, e);
// 下载失败，将任务ID放回队列头部，以便快速重试
tasks_clone.lock().await.push_front(id);
}
}
}
Ok(response) => { // HTTP 错误 (e.g., 404, 500)
eprintln!("[Worker] HTTP error for chunk {}: Status {}. Re-queueing.", id, response.status());
tasks_clone.lock().await.push_front(id);
}
Err(e) => { // 网络错误
eprintln!("[Worker] Network error for chunk {}: {}. Re-queueing.", id, e);
tasks_clone.lock().await.push_front(id);
}
}
} else {
// 队列为空，工人任务结束
println!("[Worker {:?}] No more tasks. Shutting down.", std::thread::current().id());
break;
}
}
};
worker_handles.push(tokio::spawn(worker_task));
}
if worker_handles.is_empty() {
return Err(DownloadError::NoWorkers);
}
// 5. 等待所有任务完成
// 首先丢弃此处的 data_tx，这样当所有工人都结束后，data_rx 才会感知到通道关闭
drop(data_tx);
// 等待所有工人完成
for handle in worker_handles {
handle.await?; // 这里只处理 JoinError，工人内部的逻辑错误已自行处理
}
println!("All workers finished.");
// 最后等待写入任务完成
writer_handle.await??;
println!("Download completed successfully!");
Ok(())
}
// --- main 函数，用于演示和测试 ---
#[tokio::main]
async fn main() -> Result<()> {
// 设置日志，方便观察
tracing_subscriber::fmt::init();
// 1. 启动一个本地文件服务器
let server_handle = tokio::spawn(start_test_server());
// 等待服务器启动
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
// 2. 准备下载参数
// 我们假设知道文件大小。在真实场景中，可能需要先发一个 HEAD 请求获取。
let filesize = 10 * 1024 * 1024; // 10 MB
let urls = vec![
"http://127.0.0.1:3000/testfile",
"http://127.0.0.1:3000/testfile", // 多个URL指向同一个文件
"http://127.0.0.1:3000/testfile",
];
let output_path = "downloaded_file.bin";
// 3. 执行下载
println!("Starting download...");
match down(filesize, urls, output_path).await {
Ok(_) => println!("\nSUCCESS: File '{}' downloaded.", output_path),
Err(e) => eprintln!("\nERROR: Download failed: {}", e),
}
// 4. 验证文件大小
if let Ok(metadata) = std::fs::metadata(output_path) {
println!("Verification: Output file size is {} bytes.", metadata.len());
assert_eq!(metadata.len(), filesize, "File size mismatch!");
// 清理
std::fs::remove_file(output_path)?;
}
// 停止服务器
server_handle.abort();
Ok(())
}
/// 启动一个 axum 测试服务器，提供一个大文件供下载
async fn start_test_server() {
use axum::{routing::get_service, Router};
use tower_http::services::ServeFile;
// 创建一个 10MB 的临时文件
let temp_file_path = "testfile.tmp";
let mut file = std::fs::File::create(temp_file_path).unwrap();
// 写入一些数据以确保它不是全零
file.write_all(&vec![0u8; 10 * 1024 * 1024]).unwrap();
let app = Router::new().route(
"/testfile",
get_service(ServeFile::new(temp_file_path)),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
.await
.unwrap();
println!("Test server listening on {}", listener.local_addr().unwrap());
axum::serve(listener, app).await.unwrap();
}