[pbs-devel] [PATCH proxmox 4/7] RateLimitedStream: allow periodic limiter updates

[Dietmar Maurer <dietmar@proxmox.com>]

---
 .../src/client/rate_limited_stream.rs         | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/proxmox-http/src/client/rate_limited_stream.rs b/proxmox-http/src/client/rate_limited_stream.rs
index 00ba066..ea99383 100644
--- a/proxmox-http/src/client/rate_limited_stream.rs
+++ b/proxmox-http/src/client/rate_limited_stream.rs
@@ -19,6 +19,8 @@ pub struct RateLimitedStream<S> {
     read_delay: Option<Pin<Box<Sleep>>>,
     write_limiter: Option<Arc<Mutex<RateLimiter>>>,
     write_delay: Option<Pin<Box<Sleep>>>,
+    update_limiter_cb: Option<Box<dyn Fn() -> (Option<Arc<Mutex<RateLimiter>>>, Option<Arc<Mutex<RateLimiter>>>) + Send>>,
+    last_limiter_update: Instant,
     stream: S,
 }
 
@@ -43,9 +45,44 @@ impl <S> RateLimitedStream<S> {
             read_delay: None,
             write_limiter,
             write_delay: None,
+            update_limiter_cb: None,
+            last_limiter_update: Instant::now(),
             stream,
         }
     }
+
+    /// Creates a new instance with limiter update callback.
+    ///
+    /// The fuction is called every minute to update/change the used limiters.
+    ///
+    /// Note: This function is called within an async context, so it
+    /// should be fast and must not block.
+    pub fn with_limiter_update_cb<F: Fn() -> (Option<Arc<Mutex<RateLimiter>>>, Option<Arc<Mutex<RateLimiter>>>) + Send + 'static>(
+        stream: S,
+        update_limiter_cb: F,
+    ) -> Self {
+        let (read_limiter, write_limiter) = update_limiter_cb();
+        Self {
+            read_limiter,
+            read_delay: None,
+            write_limiter,
+            write_delay: None,
+            update_limiter_cb: Some(Box::new(update_limiter_cb)),
+            last_limiter_update: Instant::now(),
+            stream,
+        }
+    }
+
+    fn update_limiters(&mut self) {
+        if let Some(ref update_limiter_cb) = self.update_limiter_cb {
+            if self.last_limiter_update.elapsed().as_secs() >= 5 {
+                self.last_limiter_update = Instant::now();
+                let (read_limiter, write_limiter) = update_limiter_cb();
+                self.read_limiter = read_limiter;
+                self.write_limiter = write_limiter;
+            }
+        }
+    }
 }
 
 fn register_traffic(
@@ -90,6 +127,8 @@ impl <S: AsyncWrite + Unpin> AsyncWrite for RateLimitedStream<S> {
 
         this.write_delay = None;
 
+        this.update_limiters();
+
         let result = Pin::new(&mut this.stream).poll_write(ctx, buf);
 
         if let Some(ref limiter) = this.write_limiter {
@@ -118,6 +157,8 @@ impl <S: AsyncWrite + Unpin> AsyncWrite for RateLimitedStream<S> {
 
         this.write_delay = None;
 
+        this.update_limiters();
+
         let result = Pin::new(&mut this.stream).poll_write_vectored(ctx, bufs);
 
         if let Some(ref limiter) = this.write_limiter {
@@ -161,6 +202,8 @@ impl <S: AsyncRead + Unpin> AsyncRead for RateLimitedStream<S> {
 
         this.read_delay = None;
 
+        this.update_limiters();
+
         let filled_len = buf.filled().len();
         let result = Pin::new(&mut this.stream).poll_read(ctx, buf);
 
-- 
2.30.2