性能优化

2026-08-06 10:33:16 +08:00
parent c7578a2e6b
commit e66c53e66d
105 changed files with 7273 additions and 5002 deletions
@@ -12,6 +12,12 @@ use tokio::task::JoinSet;
 use super::rate_limit::RateLimiter;
 use super::task::{ProbeResult, Segment};

+/// 请求超时：连接 + 响应头必须在 30s 内就绪（分段请求若无超时，
+/// 服务器挂死时任务将永久卡在 Active，pause→resume 会出现新旧任务并发写同一临时文件）
+const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
+/// 分块读取停滞超时：30s 内无任何数据视为连接挂死，主动中断（配合取消标志及时退出）
+const READ_STALL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
+
 /// HTTP/HTTPS 下载器
 #[derive(Clone)]
 pub struct HttpDownloader {
@@ -43,7 +49,10 @@ impl HttpDownloader {
            req = req.header(k, v);
        }

-        match req.send().await {
+        match tokio::time::timeout(REQUEST_TIMEOUT, req.send())
+            .await
+            .map_err(|_| "探测超时（30s 内未收到响应头）".to_string())?
+        {
            Ok(resp) => {
                let status = resp.status();
                let headers_map = resp.headers().clone();
@@ -92,9 +101,9 @@ impl HttpDownloader {
                for (k, v) in headers {
                    head_req = head_req.header(k, v);
                }
-                let resp = head_req
-                    .send()
+                let resp = tokio::time::timeout(REQUEST_TIMEOUT, head_req.send())
                    .await
+                    .map_err(|_| "探测超时（HEAD 30s 内未收到响应头）".to_string())?
                    .map_err(|e| format!("探测失败（GET 和 HEAD 均失败）: {}", e))?;
                let headers_map = resp.headers().clone();
                let total_size = headers_map
@@ -295,14 +304,26 @@ async fn download_segment_with_client(
        req = req.header(k, v);
    }

-    let resp = req
-        .send()
+    let resp = tokio::time::timeout(REQUEST_TIMEOUT, req.send())
        .await
+        .map_err(|_| "请求超时（30s 内未收到响应头）".to_string())?
        .map_err(|e| format!("请求失败: {}", e))?;

    let status = resp.status();
-    if !status.is_success() && status.as_u16() != 206 {
-        return Err(format!("服务器返回 HTTP {}", status));
+    if unknown_size {
+        // 未知大小：未发送 Range 头，接受任意 2xx
+        if !status.is_success() {
+            return Err(format!("服务器返回 HTTP {}", status));
+        }
+    } else {
+        // 已发送 Range 头：必须返回 206。若服务器忽略 Range 返回 200 全文，
+        // 按 range_start 偏移写入会错位 → 静默损坏文件；此处直接中断。
+        if status.as_u16() != 206 {
+            return Err(format!(
+                "服务器未按分段请求响应（期望 206，实际 {}），已中断以避免文件损坏",
+                status
+            ));
+        }
    }

    // 流式读取并写入文件
@@ -315,8 +336,9 @@ async fn download_segment_with_client(
            return Err("已取消".to_string());
        }

-        match stream.next().await {
-            Some(Ok(chunk)) => {
+        // 停滞超时：30s 无数据即中断，确保 cancel 标志能及时被感知（配合代际句柄防并发写）
+        match tokio::time::timeout(READ_STALL_TIMEOUT, stream.next()).await {
+            Ok(Some(Ok(chunk))) => {
                buf.extend_from_slice(&chunk);
                // 接收到数据立即更新进度（避免监控周期内进度无变化导致速度显示为 0）
                local_completed += chunk.len() as u64;
@@ -331,10 +353,10 @@ async fn download_segment_with_client(
                    buf.clear();
                }
            }
-            Some(Err(e)) => {
+            Ok(Some(Err(e))) => {
                return Err(format!("读取数据失败: {}", e));
            }
-            None => {
+            Ok(None) => {
                // 流结束，写入剩余数据
                if !buf.is_empty() {
                    file.write_all(&buf)
@@ -346,6 +368,9 @@ async fn download_segment_with_client(
                }
                break;
            }
+            Err(_) => {
+                return Err("读取超时（30s 无数据，已中断下载）".to_string());
+            }
        }
    }

@@ -436,10 +461,11 @@ fn percent_decode(s: &str) -> String {
 /// 将文件大小划分为 N 个分段
 pub fn split_segments(total_size: u64, num_connections: u32) -> Vec<Segment> {
    if total_size == 0 || num_connections == 0 {
+        // 空文件或未指定连接数：单段覆盖整个文件（end 按 total_size 推导，不能硬编码 0）
        return vec![Segment {
            index: 0,
            start: 0,
-            end: 0,
+            end: total_size.saturating_sub(1),
            completed: 0,
        }];
    }
@@ -473,3 +499,86 @@ pub fn split_segments(total_size: u64, num_connections: u32) -> Vec<Segment> {

    segments
 }
+
+#[cfg(test)]
+mod split_segments_tests {
+    use super::*;
+
+    /// 分段必须完整覆盖 [0, total_size)，且相互连续无重叠
+    fn assert_contiguous(segments: &[Segment], total_size: u64) {
+        assert!(!segments.is_empty());
+        let mut prev_end: i64 = -1;
+        for seg in segments {
+            assert_eq!(seg.start as i64, prev_end + 1, "分段不连续");
+            assert!(seg.end >= seg.start, "分段 start > end");
+            prev_end = seg.end as i64;
+        }
+        assert_eq!(
+            segments.last().unwrap().end,
+            total_size - 1,
+            "末段未覆盖文件末尾"
+        );
+    }
+
+    #[test]
+    fn zero_size_returns_single_segment() {
+        let segs = split_segments(0, 4);
+        assert_eq!(segs.len(), 1);
+        assert_eq!(segs[0].start, 0);
+        assert_eq!(segs[0].end, 0);
+    }
+
+    #[test]
+    fn zero_connections_covers_whole_file() {
+        let segs = split_segments(1024, 0);
+        assert_eq!(segs.len(), 1);
+        assert_eq!(segs[0].start, 0);
+        assert_eq!(segs[0].end, 1023);
+    }
+
+    #[test]
+    fn divides_evenly_with_contiguous_coverage() {
+        // 10MB / 4 连接 → 4 段完整覆盖
+        let total = 10 * 1024 * 1024;
+        let segs = split_segments(total, 4);
+        assert_eq!(segs.len(), 4);
+        assert_contiguous(&segs, total);
+    }
+
+    #[test]
+    fn clamps_connections_by_min_segment_size() {
+        // 2MB 文件请求 8 连接 → 受 1MB 最小分段限制，实际 ≤ 2 段
+        let total = 2 * 1024 * 1024;
+        let segs = split_segments(total, 8);
+        assert!(segs.len() <= 2, "连接数未按最小分段收敛: {}", segs.len());
+        assert_contiguous(&segs, total);
+    }
+
+    #[test]
+    fn respects_requested_connection_count() {
+        // 大文件按请求连接数切分
+        let total = 100 * 1024 * 1024;
+        let segs = split_segments(total, 3);
+        assert_eq!(segs.len(), 3);
+        assert_contiguous(&segs, total);
+        // 每段大小均匀
+        for seg in &segs {
+            let seg_len = seg.end - seg.start + 1;
+            assert!(
+                seg_len >= total / 3,
+                "分段大小不均: {} 段只有 {} 字节",
+                seg.index,
+                seg_len
+            );
+        }
+    }
+
+    #[test]
+    fn tiny_file_single_segment() {
+        // 小于 1MB 的文件始终单段
+        let total = 100;
+        let segs = split_segments(total, 4);
+        assert_eq!(segs.len(), 1);
+        assert_contiguous(&segs, total);
+    }
+}