From 39b71269207d463fae7be321bddeee093d2edb89 Mon Sep 17 00:00:00 2001 From: mrstreamer963 <698lo7epy@mozmail.com> Date: Thu, 9 Jul 2026 23:09:54 +0300 Subject: [PATCH] Add download speed limit feature - Introduced `-l ` option in CLI to limit total download speed. - Updated README to document the new speed limit feature. - Implemented `RateLimiter` for async download throttling in the downloader. - Enhanced download logic to incorporate rate limiting during file downloads. --- README.md | 1 + src/cli.rs | 7 ++++ src/downloader.rs | 22 +++++++++++-- src/main.rs | 14 +++++++- src/util.rs | 84 +++++++++++++++++++++++++++++++++++++++++++++-- 5 files changed, 122 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index bf9bb91..0fabb9e 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,7 @@ Run `dl --help` for a full list of commands and flags. > **Note:** You must provide only one of the following: `-f`, `-H`, `-m`, or direct URLs. * `-c `: (Optional) Number of concurrent downloads. Defaults to `3`. +* `-l `: (Optional) Limit total download speed across all concurrent downloads (e.g. `500K`, `5M`, `10MB/s`). * `-f `: Download from a text file of URLs. * `-H `: Download from a Hugging Face repo (`owner/repo_name` or full URL). * `-m `: Download a pre-defined model by alias. diff --git a/src/cli.rs b/src/cli.rs index 399feaa..5f6d031 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -16,6 +16,9 @@ It also includes utilities for searching Hugging Face models and self-updating." Download from a list in a file with concurrency 5: dl -f urls.txt -c 5 + Download with a total speed limit of 5 MB/s: + dl -f urls.txt -c 3 -l 5M + Download (and select files) from a Hugging Face repo using token: dl -H TheBloke/Llama-2-7B-GGUF -s --token @@ -61,6 +64,10 @@ pub struct Cli { /// Enable debug logging to log.log. #[arg(long)] pub debug: bool, + + /// Limit total download speed (e.g. 500K, 5M, 10MB/s). + #[arg(short = 'l', long, value_name = "SPEED")] + pub rate_limit: Option, } #[derive(Subcommand, Debug)] diff --git a/src/downloader.rs b/src/downloader.rs index c112b6e..91127cd 100644 --- a/src/downloader.rs +++ b/src/downloader.rs @@ -1,7 +1,10 @@ use crate::{ config::GGUF_SERIES_REGEX, hf::HFFile, - util::{format_bytes, format_duration_human, generate_actual_filename, get_client, shorten_error}, + util::{ + format_bytes, format_duration_human, generate_actual_filename, get_client, shorten_error, + RateLimiter, + }, }; use anyhow::{anyhow, Context, Result}; use futures_util::stream::StreamExt; @@ -17,6 +20,7 @@ use std::sync::{ }; use tokio::io::AsyncWriteExt; +use tokio::sync::Mutex as AsyncMutex; // A dedicated, higher concurrency level for fetching metadata. // This is much faster than the default download concurrency of 3. @@ -35,6 +39,7 @@ struct DownloadTask { overall_progress_bar: ProgressBar, multi_progress: Arc, client: reqwest::Client, + rate_limiter: Option>>, } pub async fn run_downloads( @@ -42,14 +47,21 @@ pub async fn run_downloads( base_dir: PathBuf, concurrency: usize, hf_token: String, + rate_limit: Option, ) -> Result<()> { + let rate_limit_msg = rate_limit + .map(|r| format!(", rate limit {}/s", format_bytes(r))) + .unwrap_or_default(); eprintln!( - "[INFO] Preparing to download {} file(s) to '{}' with concurrency {}.", + "[INFO] Preparing to download {} file(s) to '{}' with concurrency {}{}.", items.len(), base_dir.display(), - concurrency + concurrency, + rate_limit_msg ); + let rate_limiter = rate_limit.map(|r| Arc::new(AsyncMutex::new(RateLimiter::new(r)))); + let multi_progress = Arc::new(MultiProgress::new()); // --- Pre-scan for file sizes --- @@ -149,6 +161,7 @@ pub async fn run_downloads( overall_progress_bar: overall_pb.clone(), multi_progress: multi_progress.clone(), client: download_client.clone(), + rate_limiter: rate_limiter.clone(), }); } @@ -247,6 +260,9 @@ async fn download_file(task: DownloadTask) -> Result<()> { let chunk = chunk_result.context("Failed to read chunk from download stream")?; file.write_all(&chunk).await.context("Failed to write chunk to file")?; let chunk_len = chunk.len() as u64; + if let Some(limiter) = &task.rate_limiter { + limiter.lock().await.acquire(chunk_len).await; + } pb.inc(chunk_len); overall_pb.inc(chunk_len); } diff --git a/src/main.rs b/src/main.rs index ad5e1f3..ef47718 100644 --- a/src/main.rs +++ b/src/main.rs @@ -175,7 +175,19 @@ async fn run_downloader_flow(cli: Cli, hf_token: &str) -> Result<()> { tokio::fs::create_dir_all(&download_dir).await?; } - run_downloads(download_items, download_dir, cli.concurrency, hf_token.to_string()).await?; + let rate_limit = match cli.rate_limit { + Some(s) => Some(util::parse_speed_limit(&s)?), + None => None, + }; + + run_downloads( + download_items, + download_dir, + cli.concurrency, + hf_token.to_string(), + rate_limit, + ) + .await?; Ok(()) } \ No newline at end of file diff --git a/src/util.rs b/src/util.rs index 37d4523..6e79dbc 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,9 +1,9 @@ -use anyhow::Result; +use anyhow::{anyhow, Result}; use path_clean::PathClean; use std::backtrace::Backtrace; use std::panic::PanicHookInfo; use std::path::{Path, PathBuf}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; /// Formats a size in bytes into a human-readable string (KB, MB, GB, etc. - base 10). pub fn format_bytes(bytes: u64) -> String { @@ -183,6 +183,86 @@ pub fn log_panic(info: &PanicHookInfo<'_>) { } +/// Parses a human-readable speed limit string into bytes per second. +/// Supports suffixes K, KB, M, MB (base-10), optional `/s`, or raw bytes/sec. +pub fn parse_speed_limit(s: &str) -> Result { + let mut s = s.trim().to_uppercase(); + if s.is_empty() { + return Err(anyhow!("Speed limit cannot be empty")); + } + if s.ends_with("/S") { + s.truncate(s.len() - 2); + } + + let (num_str, multiplier) = if s.ends_with("MB") { + (&s[..s.len() - 2], 1_000_000u64) + } else if s.ends_with("KB") { + (&s[..s.len() - 2], 1_000u64) + } else if s.ends_with('M') { + (&s[..s.len() - 1], 1_000_000u64) + } else if s.ends_with('K') { + (&s[..s.len() - 1], 1_000u64) + } else if s.ends_with('B') { + (&s[..s.len() - 1], 1u64) + } else { + (s.as_str(), 1u64) + }; + + let value: f64 = num_str + .trim() + .parse() + .map_err(|_| anyhow!("Invalid speed limit: '{}'", s))?; + if value <= 0.0 { + return Err(anyhow!("Speed limit must be greater than zero")); + } + + let bytes = (value * multiplier as f64) as u64; + if bytes == 0 { + return Err(anyhow!("Speed limit must be greater than zero")); + } + Ok(bytes) +} + +/// Token-bucket rate limiter for async download throttling. +pub struct RateLimiter { + rate: u64, + tokens: f64, + last_update: Instant, +} + +impl RateLimiter { + pub fn new(rate: u64) -> Self { + Self { + rate, + tokens: rate as f64, + last_update: Instant::now(), + } + } + + pub async fn acquire(&mut self, bytes: u64) { + if self.rate == 0 || bytes == 0 { + return; + } + + loop { + let now = Instant::now(); + let elapsed = now.duration_since(self.last_update).as_secs_f64(); + self.tokens = (self.tokens + elapsed * self.rate as f64) + .min(self.rate as f64 * 2.0); + self.last_update = now; + + if self.tokens >= bytes as f64 { + self.tokens -= bytes as f64; + return; + } + + let deficit = bytes as f64 - self.tokens; + let wait_secs = deficit / self.rate as f64; + tokio::time::sleep(Duration::from_secs_f64(wait_secs)).await; + } + } +} + /// Creates a reqwest client with a default user agent and optional auth token. pub fn get_client(hf_token: &str) -> Result { let mut headers = reqwest::header::HeaderMap::new();