From 2a7c632b2bbcde1c804402108f10ed2ef76886c8 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Mon, 18 Nov 2024 17:40:04 +0000
Subject: [PATCH] Expose ratelimits in CLI

---
 crates/librqbit/src/lib.rs |  2 +-
 crates/rqbit/src/main.rs   | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/crates/librqbit/src/lib.rs b/crates/librqbit/src/lib.rs
index 0eadf6f..e13b9f1 100644
--- a/crates/librqbit/src/lib.rs
+++ b/crates/librqbit/src/lib.rs
@@ -51,7 +51,7 @@ mod file_ops;
 pub mod http_api;
 #[cfg(feature = "http-api")]
 pub mod http_api_client;
-mod limits;
+pub mod limits;
 mod merge_streams;
 mod peer_connection;
 mod peer_info_reader;
diff --git a/crates/rqbit/src/main.rs b/crates/rqbit/src/main.rs
index 9a0eb5b..726cfd7 100644
--- a/crates/rqbit/src/main.rs
+++ b/crates/rqbit/src/main.rs
@@ -1,6 +1,7 @@
 use std::{
     io,
     net::SocketAddr,
+    num::NonZeroU32,
     path::{Path, PathBuf},
     sync::Arc,
     thread,
@@ -14,6 +15,7 @@ use librqbit::{
     api::ApiAddTorrentResponse,
     http_api::{HttpApi, HttpApiOptions},
     http_api_client, librqbit_spawn,
+    limits::LimitsConfig,
     storage::{
         filesystem::{FilesystemStorageFactory, MmapFilesystemStorageFactory},
         StorageFactory, StorageFactoryExt,
@@ -218,6 +220,14 @@ struct Opts {
     #[cfg(feature = "disable-upload")]
     #[arg(long, env = "RQBIT_DISABLE_UPLOAD")]
     disable_upload: bool,
+
+    /// Limit download to bytes-per-second.
+    #[arg(long = "ratelimit-download", env = "RQBIT_RATELIMIT_DOWNLOAD")]
+    ratelimit_download_bps: Option<NonZeroU32>,
+
+    /// Limit upload to bytes-per-second.
+    #[arg(long = "ratelimit-upload", env = "RQBIT_RATELIMIT_UPLOAD")]
+    ratelimit_upload_bps: Option<NonZeroU32>,
 }
 
 #[derive(Parser)]
@@ -480,8 +490,10 @@ async fn async_main(opts: Opts, cancel: CancellationToken) -> anyhow::Result<()>
         cancellation_token: Some(cancel.clone()),
         #[cfg(feature = "disable-upload")]
         disable_upload: opts.disable_upload,
-        // TODO: expose
-        ratelimits: Default::default(),
+        ratelimits: LimitsConfig {
+            upload_bps: opts.ratelimit_upload_bps,
+            download_bps: opts.ratelimit_download_bps,
+        },
     };
 
     let stats_printer = |session: Arc<Session>| async move {