kata-containers/cloud-hypervisor
1
0
Fork 0
mirror of https://github.com/cloud-hypervisor/cloud-hypervisor.git synced 2026-08-05 02:19:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

rate-limiter: Allow RateLimiter to be shared between threads

Browse Source 
Signed-off-by: Thomas Barrett <tbarrett@crusoeenergy.com>
This commit is contained in:
Thomas Barrett
2023-09-04 14:10:45 -07:00
committed by Bo Chen
parent e880aeed8d
commit 8d2e590886
1 changed files with 106 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

195
rate_limiter/src/lib.rs
View File

@@ -46,9 +46,11 @@
#[macro_use] #[macro_use]
extern crate log; extern crate log;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd}; use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use std::{fmt, io};
use vmm_sys_util::timerfd::TimerFd; use vmm_sys_util::timerfd::TimerFd;
#[derive(Debug)] #[derive(Debug)]
@@ -271,27 +273,27 @@ pub enum BucketUpdate {
/// implementation. These events are meant to be consumed by the user of this struct. /// implementation. These events are meant to be consumed by the user of this struct.
/// On each such event, the user must call the `event_handler()` method. /// On each such event, the user must call the `event_handler()` method.
pub struct RateLimiter { pub struct RateLimiter {
inner: Mutex<RateLimiterInner>,
// Internal flag that quickly determines timer state.
timer_active: AtomicBool,
}
struct RateLimiterInner {
bandwidth: Option<TokenBucket>, bandwidth: Option<TokenBucket>,
ops: Option<TokenBucket>, ops: Option<TokenBucket>,
timer_fd: TimerFd, timer_fd: TimerFd,
// Internal flag that quickly determines timer state.
timer_active: bool,
} }
impl PartialEq for RateLimiter { impl RateLimiterInner {
fn eq(&self, other: &RateLimiter) -> bool { // Arm the timer of the rate limiter with the provided `Duration` (which will fire only once).
self.bandwidth == other.bandwidth && self.ops == other.ops fn activate_timer(&mut self, dur: Duration, flag: &AtomicBool) {
} // Panic when failing to arm the timer (same handling in crate TimerFd::set_state())
} self.timer_fd
.reset(dur, None)
impl fmt::Debug for RateLimiter { .expect("Can't arm the timer (unexpected 'timerfd_settime' failure).");
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { flag.store(true, Ordering::Relaxed)
write!(
f,
"RateLimiter {{ bandwidth: {:?}, ops: {:?} }}",
self.bandwidth, self.ops
)
} }
} }
@@ -355,35 +357,28 @@ impl RateLimiter {
} }
Ok(RateLimiter { Ok(RateLimiter {
bandwidth: bytes_token_bucket, inner: Mutex::new(RateLimiterInner {
ops: ops_token_bucket, bandwidth: bytes_token_bucket,
timer_fd, ops: ops_token_bucket,
timer_active: false, timer_fd,
}),
timer_active: AtomicBool::new(false),
}) })
} }
// Arm the timer of the rate limiter with the provided `Duration` (which will fire only once).
fn activate_timer(&mut self, dur: Duration) {
// Panic when failing to arm the timer (same handling in crate TimerFd::set_state())
self.timer_fd
.reset(dur, None)
.expect("Can't arm the timer (unexpected 'timerfd_settime' failure).");
self.timer_active = true;
}
/// Attempts to consume tokens and returns whether that is possible. /// Attempts to consume tokens and returns whether that is possible.
/// ///
/// If rate limiting is disabled on provided `token_type`, this function will always succeed. /// If rate limiting is disabled on provided `token_type`, this function will always succeed.
pub fn consume(&mut self, tokens: u64, token_type: TokenType) -> bool { pub fn consume(&self, tokens: u64, token_type: TokenType) -> bool {
// If the timer is active, we can't consume tokens from any bucket and the function fails. // If the timer is active, we can't consume tokens from any bucket and the function fails.
if self.timer_active { if self.is_blocked() {
return false; return false;
} }
let mut guard = self.inner.lock().unwrap();
// Identify the required token bucket. // Identify the required token bucket.
let token_bucket = match token_type { let token_bucket = match token_type {
TokenType::Bytes => self.bandwidth.as_mut(), TokenType::Bytes => guard.bandwidth.as_mut(),
TokenType::Ops => self.ops.as_mut(), TokenType::Ops => guard.ops.as_mut(),
}; };
// Try to consume from the token bucket. // Try to consume from the token bucket.
if let Some(bucket) = token_bucket { if let Some(bucket) = token_bucket {
@@ -393,8 +388,8 @@ impl RateLimiter {
// register a timer to replenish the bucket and resume processing; // register a timer to replenish the bucket and resume processing;
// make sure there is only one running timer for this limiter. // make sure there is only one running timer for this limiter.
BucketReduction::Failure => { BucketReduction::Failure => {
if !self.timer_active { if !self.is_blocked() {
self.activate_timer(TIMER_REFILL_DUR); guard.activate_timer(TIMER_REFILL_DUR, &self.timer_active);
} }
false false
} }
@@ -409,7 +404,10 @@ impl RateLimiter {
// order to enforce the bandwidth limit we need to prevent // order to enforce the bandwidth limit we need to prevent
// further calls to the rate limiter for // further calls to the rate limiter for
// `ratio * refill_time` milliseconds. // `ratio * refill_time` milliseconds.
self.activate_timer(Duration::from_millis((ratio * refill_time as f64) as u64)); guard.activate_timer(
Duration::from_millis((ratio * refill_time as f64) as u64),
&self.timer_active,
);
true true
} }
} }
@@ -424,11 +422,12 @@ impl RateLimiter {
/// ///
/// Can be used to *manually* add tokens to a bucket. Useful for reverting a /// Can be used to *manually* add tokens to a bucket. Useful for reverting a
/// `consume()` if needed. /// `consume()` if needed.
pub fn manual_replenish(&mut self, tokens: u64, token_type: TokenType) { pub fn manual_replenish(&self, tokens: u64, token_type: TokenType) {
let mut guard = self.inner.lock().unwrap();
// Identify the required token bucket. // Identify the required token bucket.
let token_bucket = match token_type { let token_bucket = match token_type {
TokenType::Bytes => self.bandwidth.as_mut(), TokenType::Bytes => guard.bandwidth.as_mut(),
TokenType::Ops => self.ops.as_mut(), TokenType::Ops => guard.ops.as_mut(),
}; };
// Add tokens to the token bucket. // Add tokens to the token bucket.
if let Some(bucket) = token_bucket { if let Some(bucket) = token_bucket {
@@ -442,7 +441,7 @@ impl RateLimiter {
/// budget for it. /// budget for it.
/// An event will be generated on the exported FD when the limiter 'unblocks'. /// An event will be generated on the exported FD when the limiter 'unblocks'.
pub fn is_blocked(&self) -> bool { pub fn is_blocked(&self) -> bool {
self.timer_active self.timer_active.load(Ordering::Relaxed)
} }
/// This function needs to be called every time there is an event on the /// This function needs to be called every time there is an event on the
@@ -451,11 +450,12 @@ impl RateLimiter {
/// # Errors /// # Errors
/// ///
/// If the rate limiter is disabled or is not blocked, an error is returned. /// If the rate limiter is disabled or is not blocked, an error is returned.
pub fn event_handler(&mut self) -> Result<(), Error> { pub fn event_handler(&self) -> Result<(), Error> {
let mut guard = self.inner.lock().unwrap();
loop { loop {
// Note: As we manually added the `O_NONBLOCK` flag to the FD, the following // Note: As we manually added the `O_NONBLOCK` flag to the FD, the following
// `timer_fd::wait()` won't block (which is different from its default behavior.) // `timer_fd::wait()` won't block (which is different from its default behavior.)
match self.timer_fd.wait() { match guard.timer_fd.wait() {
Err(e) => { Err(e) => {
let err: std::io::Error = e.into(); let err: std::io::Error = e.into();
match err.kind() { match err.kind() {
@@ -469,7 +469,7 @@ impl RateLimiter {
} }
} }
_ => { _ => {
self.timer_active = false; self.timer_active.store(false, Ordering::Relaxed);
return Ok(()); return Ok(());
} }
} }
@@ -479,27 +479,18 @@ impl RateLimiter {
/// Updates the parameters of the token buckets associated with this RateLimiter. /// Updates the parameters of the token buckets associated with this RateLimiter.
// TODO: Please note that, right now, the buckets become full after being updated. // TODO: Please note that, right now, the buckets become full after being updated.
pub fn update_buckets(&mut self, bytes: BucketUpdate, ops: BucketUpdate) { pub fn update_buckets(&mut self, bytes: BucketUpdate, ops: BucketUpdate) {
let mut guard = self.inner.lock().unwrap();
match bytes { match bytes {
BucketUpdate::Disabled => self.bandwidth = None, BucketUpdate::Disabled => guard.bandwidth = None,
BucketUpdate::Update(tb) => self.bandwidth = Some(tb), BucketUpdate::Update(tb) => guard.bandwidth = Some(tb),
BucketUpdate::None => (), BucketUpdate::None => (),
}; };
match ops { match ops {
BucketUpdate::Disabled => self.ops = None, BucketUpdate::Disabled => guard.ops = None,
BucketUpdate::Update(tb) => self.ops = Some(tb), BucketUpdate::Update(tb) => guard.ops = Some(tb),
BucketUpdate::None => (), BucketUpdate::None => (),
}; };
} }
/// Returns an immutable view of the inner bandwidth token bucket.
pub fn bandwidth(&self) -> Option<&TokenBucket> {
self.bandwidth.as_ref()
}
/// Returns an immutable view of the inner ops token bucket.
pub fn ops(&self) -> Option<&TokenBucket> {
self.ops.as_ref()
}
} }
impl AsRawFd for RateLimiter { impl AsRawFd for RateLimiter {
@@ -510,7 +501,8 @@ impl AsRawFd for RateLimiter {
/// Will return a negative value if rate limiting is disabled on both /// Will return a negative value if rate limiting is disabled on both
/// token types. /// token types.
fn as_raw_fd(&self) -> RawFd { fn as_raw_fd(&self) -> RawFd {
self.timer_fd.as_raw_fd() let guard = self.inner.lock().unwrap();
guard.timer_fd.as_raw_fd()
} }
} }
@@ -525,6 +517,7 @@ impl Default for RateLimiter {
#[cfg(test)] #[cfg(test)]
pub(crate) mod tests { pub(crate) mod tests {
use super::*; use super::*;
use std::fmt;
use std::thread; use std::thread;
use std::time::Duration; use std::time::Duration;
@@ -557,11 +550,33 @@ pub(crate) mod tests {
} }
impl RateLimiter { impl RateLimiter {
fn get_token_bucket(&self, token_type: TokenType) -> Option<&TokenBucket> { pub fn bandwidth(&self) -> Option<TokenBucket> {
match token_type { let guard = self.inner.lock().unwrap();
TokenType::Bytes => self.bandwidth.as_ref(), guard.bandwidth.clone()
TokenType::Ops => self.ops.as_ref(), }
}
pub fn ops(&self) -> Option<TokenBucket> {
let guard = self.inner.lock().unwrap();
guard.ops.clone()
}
}
impl PartialEq for RateLimiter {
fn eq(&self, other: &RateLimiter) -> bool {
let self_guard = self.inner.lock().unwrap();
let other_guard = other.inner.lock().unwrap();
self_guard.bandwidth == other_guard.bandwidth && self_guard.ops == other_guard.ops
}
}
impl fmt::Debug for RateLimiter {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let guard = self.inner.lock().unwrap();
write!(
f,
"RateLimiter {{ bandwidth: {:?}, ops: {:?} }}",
guard.bandwidth, guard.ops
)
} }
} }
@@ -640,7 +655,7 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_default() { fn test_rate_limiter_default() {
let mut l = RateLimiter::default(); let l = RateLimiter::default();
// limiter should not be blocked // limiter should not be blocked
assert!(!l.is_blocked()); assert!(!l.is_blocked());
@@ -659,14 +674,13 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_new() { fn test_rate_limiter_new() {
let l = RateLimiter::new(1000, 1001, 1002, 1003, 1004, 1005).unwrap(); let l = RateLimiter::new(1000, 1001, 1002, 1003, 1004, 1005).unwrap();
let bw = l.bandwidth().unwrap();
let bw = l.bandwidth.unwrap();
assert_eq!(bw.capacity(), 1000); assert_eq!(bw.capacity(), 1000);
assert_eq!(bw.one_time_burst(), 1001); assert_eq!(bw.one_time_burst(), 1001);
assert_eq!(bw.refill_time_ms(), 1002); assert_eq!(bw.refill_time_ms(), 1002);
assert_eq!(bw.budget(), 1000); assert_eq!(bw.budget(), 1000);
let ops = l.ops.unwrap(); let ops = l.ops().unwrap();
assert_eq!(ops.capacity(), 1003); assert_eq!(ops.capacity(), 1003);
assert_eq!(ops.one_time_burst(), 1004); assert_eq!(ops.one_time_burst(), 1004);
assert_eq!(ops.refill_time_ms(), 1005); assert_eq!(ops.refill_time_ms(), 1005);
@@ -676,20 +690,20 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_manual_replenish() { fn test_rate_limiter_manual_replenish() {
// rate limiter with limit of 1000 bytes/s and 1000 ops/s // rate limiter with limit of 1000 bytes/s and 1000 ops/s
let mut l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap(); let l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap();
// consume 123 bytes // consume 123 bytes
assert!(l.consume(123, TokenType::Bytes)); assert!(l.consume(123, TokenType::Bytes));
l.manual_replenish(23, TokenType::Bytes); l.manual_replenish(23, TokenType::Bytes);
{ {
let bytes_tb = l.get_token_bucket(TokenType::Bytes).unwrap(); let bytes_tb = l.bandwidth().unwrap();
assert_eq!(bytes_tb.budget(), 900); assert_eq!(bytes_tb.budget(), 900);
} }
// consume 123 ops // consume 123 ops
assert!(l.consume(123, TokenType::Ops)); assert!(l.consume(123, TokenType::Ops));
l.manual_replenish(23, TokenType::Ops); l.manual_replenish(23, TokenType::Ops);
{ {
let bytes_tb = l.get_token_bucket(TokenType::Ops).unwrap(); let bytes_tb = l.ops().unwrap();
assert_eq!(bytes_tb.budget(), 900); assert_eq!(bytes_tb.budget(), 900);
} }
} }
@@ -697,7 +711,7 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_bandwidth() { fn test_rate_limiter_bandwidth() {
// rate limiter with limit of 1000 bytes/s // rate limiter with limit of 1000 bytes/s
let mut l = RateLimiter::new(1000, 0, 1000, 0, 0, 0).unwrap(); let l = RateLimiter::new(1000, 0, 1000, 0, 0, 0).unwrap();
// limiter should not be blocked // limiter should not be blocked
assert!(!l.is_blocked()); assert!(!l.is_blocked());
@@ -730,7 +744,7 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_ops() { fn test_rate_limiter_ops() {
// rate limiter with limit of 1000 ops/s // rate limiter with limit of 1000 ops/s
let mut l = RateLimiter::new(0, 0, 0, 1000, 0, 1000).unwrap(); let l = RateLimiter::new(0, 0, 0, 1000, 0, 1000).unwrap();
// limiter should not be blocked // limiter should not be blocked
assert!(!l.is_blocked()); assert!(!l.is_blocked());
@@ -763,7 +777,7 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_full() { fn test_rate_limiter_full() {
// rate limiter with limit of 1000 bytes/s and 1000 ops/s // rate limiter with limit of 1000 bytes/s and 1000 ops/s
let mut l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap(); let l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap();
// limiter should not be blocked // limiter should not be blocked
assert!(!l.is_blocked()); assert!(!l.is_blocked());
@@ -799,7 +813,7 @@ pub(crate) mod tests {
#[test] #[test]
fn test_rate_limiter_overconsumption() { fn test_rate_limiter_overconsumption() {
// initialize the rate limiter // initialize the rate limiter
let mut l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap(); let l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap();
// try to consume 2.5x the bucket size // try to consume 2.5x the bucket size
// we are "borrowing" 1.5x the bucket size in tokens since // we are "borrowing" 1.5x the bucket size in tokens since
// the bucket is full // the bucket is full
@@ -818,7 +832,7 @@ pub(crate) mod tests {
assert!(!l.is_blocked()); assert!(!l.is_blocked());
// reset the rate limiter // reset the rate limiter
let mut l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap(); let l = RateLimiter::new(1000, 0, 1000, 1000, 0, 1000).unwrap();
// try to consume 1.5x the bucket size // try to consume 1.5x the bucket size
// we are "borrowing" 1.5x the bucket size in tokens since // we are "borrowing" 1.5x the bucket size in tokens since
// the bucket is full, should arm the timer to 0.5x replenish // the bucket is full, should arm the timer to 0.5x replenish
@@ -857,12 +871,12 @@ pub(crate) mod tests {
fn test_update_buckets() { fn test_update_buckets() {
let mut x = RateLimiter::new(1000, 2000, 1000, 10, 20, 1000).unwrap(); let mut x = RateLimiter::new(1000, 2000, 1000, 10, 20, 1000).unwrap();
let initial_bw = x.bandwidth.clone(); let initial_bw = x.bandwidth();
let initial_ops = x.ops.clone(); let initial_ops = x.ops();
x.update_buckets(BucketUpdate::None, BucketUpdate::None); x.update_buckets(BucketUpdate::None, BucketUpdate::None);
assert_eq!(x.bandwidth, initial_bw); assert_eq!(x.bandwidth(), initial_bw);
assert_eq!(x.ops, initial_ops); assert_eq!(x.ops(), initial_ops);
let new_bw = TokenBucket::new(123, 0, 57).unwrap(); let new_bw = TokenBucket::new(123, 0, 57).unwrap();
let new_ops = TokenBucket::new(321, 12346, 89).unwrap(); let new_ops = TokenBucket::new(321, 12346, 89).unwrap();
@@ -871,18 +885,21 @@ pub(crate) mod tests {
BucketUpdate::Update(new_ops.clone()), BucketUpdate::Update(new_ops.clone()),
); );
// We have manually adjust the last_update field, because it changes when update_buckets() {
// constructs new buckets (and thus gets a different value for last_update). We do this so let mut guard = x.inner.lock().unwrap();
// it makes sense to test the following assertions. // We have manually adjust the last_update field, because it changes when update_buckets()
x.bandwidth.as_mut().unwrap().last_update = new_bw.last_update; // constructs new buckets (and thus gets a different value for last_update). We do this so
x.ops.as_mut().unwrap().last_update = new_ops.last_update; // it makes sense to test the following assertions.
guard.bandwidth.as_mut().unwrap().last_update = new_bw.last_update;
guard.ops.as_mut().unwrap().last_update = new_ops.last_update;
}
assert_eq!(x.bandwidth, Some(new_bw)); assert_eq!(x.bandwidth(), Some(new_bw));
assert_eq!(x.ops, Some(new_ops)); assert_eq!(x.ops(), Some(new_ops));
x.update_buckets(BucketUpdate::Disabled, BucketUpdate::Disabled); x.update_buckets(BucketUpdate::Disabled, BucketUpdate::Disabled);
assert_eq!(x.bandwidth, None); assert_eq!(x.bandwidth(), None);
assert_eq!(x.ops, None); assert_eq!(x.ops(), None);
} }
#[test] #[test]