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,
}
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(
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 {
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 {
this.read_delay = None;
+ this.update_limiters();
+
let filled_len = buf.filled().len();
let result = Pin::new(&mut this.stream).poll_read(ctx, buf);