[pbs-devel] [RFC proxmox-backup 1/2] pbs-client: add option to use the new RateLimiter

[Dietmar Maurer]

Signed-off-by: Dietmar Maurer <dietmar at proxmox.com>
---
 pbs-client/src/http_client.rs | 16 ++++++++++++++--
 pbs-client/src/tools/mod.rs   | 11 ++++++++---
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/pbs-client/src/http_client.rs b/pbs-client/src/http_client.rs
index 73c83f7a..4f395b48 100644
--- a/pbs-client/src/http_client.rs
+++ b/pbs-client/src/http_client.rs
@@ -20,7 +20,7 @@ use proxmox::{
 };
 use proxmox_router::HttpError;
 
-use proxmox_http::client::HttpsConnector;
+use proxmox_http::client::{HttpsConnector, RateLimiter};
 use proxmox_http::uri::build_authority;
 
 use pbs_api_types::{Authid, Userid};
@@ -51,6 +51,7 @@ pub struct HttpClientOptions {
     ticket_cache: bool,
     fingerprint_cache: bool,
     verify_cert: bool,
+    rate_limit: Option<u64>,
 }
 
 impl HttpClientOptions {
@@ -109,6 +110,11 @@ impl HttpClientOptions {
         self.verify_cert = verify_cert;
         self
     }
+
+    pub fn rate_limit(mut self, rate_limit:  Option<u64>) -> Self {
+        self.rate_limit = rate_limit;
+        self
+    }
 }
 
 impl Default for HttpClientOptions {
@@ -121,6 +127,7 @@ impl Default for HttpClientOptions {
             ticket_cache: false,
             fingerprint_cache: false,
             verify_cert: true,
+            rate_limit: None,
         }
     }
 }
@@ -343,7 +350,12 @@ impl HttpClient {
         httpc.enforce_http(false); // we want https...
 
         httpc.set_connect_timeout(Some(std::time::Duration::new(10, 0)));
-        let https = HttpsConnector::with_connector(httpc, ssl_connector_builder.build(), PROXMOX_BACKUP_TCP_KEEPALIVE_TIME);
+        let mut https = HttpsConnector::with_connector(httpc, ssl_connector_builder.build(), PROXMOX_BACKUP_TCP_KEEPALIVE_TIME);
+
+        if let Some(rate_limit) = options.rate_limit {
+            https.set_read_limiter(Some(Arc::new(Mutex::new(RateLimiter::new(rate_limit)))));
+            https.set_write_limiter(Some(Arc::new(Mutex::new(RateLimiter::new(rate_limit)))));
+        }
 
         let client = Client::builder()
         //.http2_initial_stream_window_size( (1 << 31) - 2)
diff --git a/pbs-client/src/tools/mod.rs b/pbs-client/src/tools/mod.rs
index a12635cf..b6135145 100644
--- a/pbs-client/src/tools/mod.rs
+++ b/pbs-client/src/tools/mod.rs
@@ -135,15 +135,20 @@ pub fn extract_repository_from_map(param: &HashMap<String, String>) -> Option<Ba
 }
 
 pub fn connect(repo: &BackupRepository) -> Result<HttpClient, Error> {
-    connect_do(repo.host(), repo.port(), repo.auth_id())
+    connect_do(repo.host(), repo.port(), repo.auth_id(), None)
         .map_err(|err| format_err!("error building client for repository {} - {}", repo, err))
 }
 
-fn connect_do(server: &str, port: u16, auth_id: &Authid) -> Result<HttpClient, Error> {
+pub fn connect_rate_limited(repo: &BackupRepository, rate: Option<u64>) -> Result<HttpClient, Error> {
+    connect_do(repo.host(), repo.port(), repo.auth_id(), rate)
+        .map_err(|err| format_err!("error building client for repository {} - {}", repo, err))
+}
+
+fn connect_do(server: &str, port: u16, auth_id: &Authid, rate_limit: Option<u64>) -> Result<HttpClient, Error> {
     let fingerprint = std::env::var(ENV_VAR_PBS_FINGERPRINT).ok();
 
     let password = get_secret_from_env(ENV_VAR_PBS_PASSWORD)?;
-    let options = HttpClientOptions::new_interactive(password, fingerprint);
+    let options = HttpClientOptions::new_interactive(password, fingerprint).rate_limit(rate_limit);
 
     HttpClient::new(server, port, auth_id, options)
 }
-- 
2.30.2