[PATCH proxmox-backup-qemu 6/8] restore: implement `rate-limit` imposing chunk data stream limits

[Christian Ebner <c.ebner@proxmox.com>]

In addition to allowing only bandwidth limit downloaded data via the
limits imposed on the http client, also allow to explicitley limit
the chunk stream data when doing restores. By this not only zero
chunks are taken into account for the limit, but also the limit
is applied on the decompressed chunk data.

This further extends the new API function introduced by previous
changes.

Signed-off-by: Christian Ebner <c.ebner@proxmox.com>
---
 Cargo.toml     |  1 +
 current-api.h  |  1 +
 src/lib.rs     | 34 ++++++++++++++++++++++++----------
 src/restore.rs | 23 +++++++++++++++++++++++
 4 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index af7a84a..1b0af2e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -34,6 +34,7 @@ openssl = "0.10"
 proxmox-async = "0.5"
 proxmox-human-byte = "1"
 proxmox-lang = "1"
+proxmox-rate-limiter = "1"
 proxmox-schema = { version = "5", features = [ "api-macro" ] }
 proxmox-sortable-macro = "1"
 proxmox-sys = "1"
diff --git a/current-api.h b/current-api.h
index b9a3a44..25bf981 100644
--- a/current-api.h
+++ b/current-api.h
@@ -330,6 +330,7 @@ struct ProxmoxRestoreHandle *proxmox_restore_new_ns_rate_limited(const char *rep
                                                                  const char *key_password,
                                                                  const char *fingerprint,
                                                                  const char *bw_limit,
+                                                                 const char *rate_limit,
                                                                  char **error);
 
 /**
diff --git a/src/lib.rs b/src/lib.rs
index 4d4c28b..0a37e2f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -145,6 +145,7 @@ pub(crate) struct BackupSetup {
     pub master_keyfile: Option<std::path::PathBuf>,
     pub fingerprint: Option<String>,
     pub inbound_rate_limit: Option<ClientRateLimitConfig>,
+    pub outbound_rate_limit: Option<ClientRateLimitConfig>,
 }
 
 // helper class to implement synchronous interface
@@ -310,6 +311,7 @@ pub extern "C" fn proxmox_backup_new_ns(
             master_keyfile,
             fingerprint,
             inbound_rate_limit: None,
+            outbound_rate_limit: None,
         };
 
         BackupTask::new(setup, compress, crypt_mode)
@@ -778,6 +780,22 @@ fn restore_handle_to_task(handle: *mut ProxmoxRestoreHandle) -> Arc<RestoreTask>
     Arc::clone(restore_task)
 }
 
+fn client_rate_limit_config_from(
+    limit: *const c_char,
+) -> Result<Option<ClientRateLimitConfig>, Error> {
+    let rate_limit_config = match tools::utf8_c_string(limit)? {
+        Some(rate_limit) => {
+            let rate_limit: HumanByte = rate_limit.parse()?;
+            Some(ClientRateLimitConfig {
+                rate: Some(rate_limit),
+                burst: None,
+            })
+        }
+        None => None,
+    };
+    Ok(rate_limit_config)
+}
+
 /// DEPRECATED: Connect to the backup server for restore (sync)
 ///
 /// Deprecated in favor of `proxmox_restore_new_ns` which includes a namespace parameter.
@@ -823,6 +841,7 @@ pub extern "C" fn proxmox_restore_new(
             master_keyfile: None,
             fingerprint,
             inbound_rate_limit: None,
+            outbound_rate_limit: None,
         };
 
         RestoreTask::new(setup)
@@ -859,6 +878,7 @@ pub extern "C" fn proxmox_restore_new_ns(
         key_password,
         fingerprint,
         null(),
+        null(),
         error,
     )
 }
@@ -875,6 +895,7 @@ pub extern "C" fn proxmox_restore_new_ns_rate_limited(
     key_password: *const c_char,
     fingerprint: *const c_char,
     bw_limit: *const c_char,
+    rate_limit: *const c_char,
     error: *mut *mut c_char,
 ) -> *mut ProxmoxRestoreHandle {
     let result: Result<_, Error> = try_block!({
@@ -896,16 +917,8 @@ pub extern "C" fn proxmox_restore_new_ns_rate_limited(
         let keyfile = tools::utf8_c_string(keyfile)?.map(std::path::PathBuf::from);
         let key_password = tools::utf8_c_string(key_password)?;
         let fingerprint = tools::utf8_c_string(fingerprint)?;
-        let inbound_rate_limit = match tools::utf8_c_string(bw_limit)? {
-            Some(rate_limit) => {
-                let rate_limit: HumanByte = rate_limit.parse()?;
-                Some(ClientRateLimitConfig {
-                    rate: Some(rate_limit),
-                    burst: None,
-                })
-            }
-            None => None,
-        };
+        let inbound_rate_limit = client_rate_limit_config_from(bw_limit)?;
+        let outbound_rate_limit = client_rate_limit_config_from(rate_limit)?;
 
         let setup = BackupSetup {
             host: repo.host().to_owned(),
@@ -921,6 +934,7 @@ pub extern "C" fn proxmox_restore_new_ns_rate_limited(
             master_keyfile: None,
             fingerprint,
             inbound_rate_limit,
+            outbound_rate_limit,
         };
 
         RestoreTask::new(setup)
diff --git a/src/restore.rs b/src/restore.rs
index fff0465..6c1a0c6 100644
--- a/src/restore.rs
+++ b/src/restore.rs
@@ -1,5 +1,6 @@
 use std::convert::TryInto;
 use std::sync::{Arc, Mutex};
+use std::time::Instant;
 
 use anyhow::{bail, format_err, Error};
 use futures::StreamExt;
@@ -7,6 +8,7 @@ use once_cell::sync::OnceCell;
 use tokio::runtime::Runtime;
 
 use proxmox_async::runtime::get_runtime_with_builder;
+use proxmox_rate_limiter::{RateLimit, RateLimiter};
 
 use pbs_api_types::{BackupArchiveName, RateLimitConfig};
 use pbs_client::{BackupReader, HttpClient, HttpClientOptions, RemoteChunkReader};
@@ -43,6 +45,7 @@ pub(crate) struct RestoreTask {
     client: OnceCell<Arc<BackupReader>>,
     manifest: OnceCell<Arc<BackupManifest>>,
     image_registry: Arc<Mutex<Registry<ImageAccessInfo>>>,
+    rate_limiter: Option<Arc<Mutex<RateLimiter>>>,
 }
 
 impl RestoreTask {
@@ -61,6 +64,13 @@ impl RestoreTask {
                 Some(Arc::new(CryptConfig::new(key)?))
             }
         };
+        let rate_limiter = if let Some(limit) = &setup.outbound_rate_limit {
+            limit
+                .rate
+                .map(|rate| Arc::new(Mutex::new(RateLimiter::new(rate.as_u64(), rate.as_u64()))))
+        } else {
+            None
+        };
 
         Ok(Self {
             setup,
@@ -69,6 +79,7 @@ impl RestoreTask {
             client: OnceCell::new(),
             manifest: OnceCell::new(),
             image_registry: Arc::new(Mutex::new(Registry::<ImageAccessInfo>::new())),
+            rate_limiter,
         })
     }
 
@@ -214,6 +225,7 @@ impl RestoreTask {
             let res = res?;
             match res {
                 (None, offset) => {
+                    self.rate_limit_chunk_stream(index.chunk_size as u64).await;
                     let res = write_zero_callback(offset, index.chunk_size as u64);
                     if res < 0 {
                         bail!("write_zero_callback failed ({})", res);
@@ -222,6 +234,7 @@ impl RestoreTask {
                     zeroes += index.chunk_size;
                 }
                 (Some(raw_data), offset) => {
+                    self.rate_limit_chunk_stream(raw_data.len() as u64).await;
                     let res = write_data_callback(offset, &raw_data);
                     if res < 0 {
                         bail!("write_data_callback failed ({})", res);
@@ -260,6 +273,16 @@ impl RestoreTask {
         Ok(())
     }
 
+    async fn rate_limit_chunk_stream(&self, chunk_size: u64) {
+        if let Some(limiter) = &self.rate_limiter {
+            let delay = limiter
+                .lock()
+                .unwrap()
+                .register_traffic(Instant::now(), chunk_size);
+            tokio::time::sleep(delay).await;
+        }
+    }
+
     pub fn get_image_length(&self, aid: u8) -> Result<u64, Error> {
         let mut guard = self.image_registry.lock().unwrap();
         let info = guard.lookup(aid)?;
-- 
2.47.3