下载非内核

src-tauri/src/download_engine/rate_limit.rs

@@ -0,0 +1,74 @@
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+/// 令牌桶限速器（字节/秒，0=不限）
+///
+/// 滑动窗口算法：1 秒窗口内累计消耗字节数，超出限制时 sleep 到下一窗口。
+pub struct RateLimiter {
+    /// 限速 bytes/s（0=不限），用 Arc<AtomicU64> 以便运行时动态调整
+    limit: Arc<AtomicU64>,
+    /// 当前窗口已消耗字节
+    consumed: AtomicU64,
+    /// 窗口起始时间
+    window_start: Mutex<std::time::Instant>,
+}
+
+impl RateLimiter {
+    pub fn new(limit: u64) -> Self {
+        Self {
+            limit: Arc::new(AtomicU64::new(limit)),
+            consumed: AtomicU64::new(0),
+            window_start: Mutex::new(std::time::Instant::now()),
+        }
+    }
+
+    /// 动态设置限速（bytes/s，0=不限）
+    pub fn set_limit(&self, limit: u64) {
+        self.limit.store(limit, Ordering::Relaxed);
+    }
+
+    /// 消费指定字节数，若超出限速则 sleep 等待
+    pub async fn consume(&self, bytes: u64) {
+        let limit = self.limit.load(Ordering::Relaxed);
+        if limit == 0 || bytes == 0 {
+            return;
+        }
+
+        // 尝试在当前窗口消费（用作用域确保 MutexGuard 在 await 前释放）
+        let over_limit = {
+            let now = std::time::Instant::now();
+            let mut start = self.window_start.lock().unwrap();
+            let elapsed = now.duration_since(*start);
+
+            // 窗口过期，重置
+            if elapsed >= Duration::from_secs(1) {
+                self.consumed.store(0, Ordering::Relaxed);
+                *start = now;
+            }
+
+            let current = self.consumed.fetch_add(bytes, Ordering::Relaxed) + bytes;
+            if current > limit {
+                Some(current - limit)
+            } else {
+                None
+            }
+        }; // MutexGuard 在此释放
+
+        // 超出限速，计算需要等待的时间
+        if let Some(over) = over_limit {
+            let wait_ms = (over as f64 / limit as f64 * 1000.0).ceil() as u64;
+            let wait = Duration::from_millis(wait_ms.min(1000));
+            tokio::time::sleep(wait).await;
+        }
+    }
+}
+
+/// 克隆限速器的共享句柄（共享同一限速状态）
+impl Clone for RateLimiter {
+    fn clone(&self) -> Self {
+        // 注意：clone 的是限速配置，不是状态。新实例独立计数。
+        // 全局限速器应通过引用共享，而非 clone。
+        Self::new(self.limit.load(Ordering::Relaxed))
+    }
+}