rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-10-27 22:16:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

share impl between proxy and ps

Browse Source
This commit is contained in:
Conrad Ludgate
2024-07-29 14:08:37 +01:00
parent 40d239560f
commit 28b85ca711
4 changed files with 191 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/src/rate_limiter/leaky_bucket.rs
+30 -82
View File
@@ -9,6 +9,7 @@ use dashmap::DashMap;
use rand::{thread_rng, Rng};
use tokio::time::Instant;
use tracing::info;
use utils::leaky_bucket::LeakyBucketState;
use crate::intern::EndpointIdInt;
@@ -17,7 +18,7 @@ pub type EndpointRateLimiter = LeakyBucketRateLimiter<EndpointIdInt>;
pub struct LeakyBucketRateLimiter<Key> {
map: DashMap<Key, LeakyBucketState, RandomState>,
config: LeakyBucketConfigInner,
config: utils::leaky_bucket::LeakyBucketConfig,
access_count: AtomicUsize,
}
@@ -46,9 +47,9 @@ impl<K: Hash + Eq> LeakyBucketRateLimiter<K> {
let mut entry = self
.map
.entry(key)
.or_insert_with(|| LeakyBucketState::new(now));
.or_insert_with(|| LeakyBucketState::new(now - self.config.epoch));
entry.check(&self.config, now, n as f64)
entry.add_tokens(&self.config, now, n as f64).is_ok()
}
fn do_gc(&self, now: Instant) {
@@ -60,7 +61,7 @@ impl<K: Hash + Eq> LeakyBucketRateLimiter<K> {
let shard = thread_rng().gen_range(0..n);
self.map.shards()[shard]
.write()
.retain(|_, value| value.get().should_retain(now));
.retain(|_, value| !value.get().bucket_is_empty(&self.config, now));
}
}
@@ -77,75 +78,17 @@ impl LeakyBucketConfig {
}
}
struct LeakyBucketConfigInner {
/// "time cost" of a single request unit.
/// loosely represents how long it takes to handle a request unit in active CPU time.
time_cost: Duration,
bucket_width: Duration,
}
impl From<LeakyBucketConfig> for LeakyBucketConfigInner {
impl From<LeakyBucketConfig> for utils::leaky_bucket::LeakyBucketConfig {
fn from(config: LeakyBucketConfig) -> Self {
// seconds_per_request = 1/(request_per_second)
let spr = config.rps.recip();
Self {
let bucket_width = Duration::from_secs_f64(config.max * spr);
utils::leaky_bucket::LeakyBucketConfig {
time_cost: Duration::from_secs_f64(spr),
bucket_width: Duration::from_secs_f64(config.max * spr),
}
}
}
struct LeakyBucketState {
/// Bucket is represented by `start..end` where `start = end - config.bucket_width`.
///
/// At any given time, `end - now` represents the number of tokens in the bucket, multiplied by the "time_cost".
/// Adding `n` tokens to the bucket is done by moving `end` forward by `n * config.time_cost`.
/// If `now < start`, the bucket is considered filled and cannot accept any more tokens.
/// Draining the bucket will happen naturally as `now` moves forward.
///
/// Let `n` be some "time cost" for the request,
/// If now is after end, the bucket is empty and the end is reset to now,
/// If now is within the `bucket window + n`, we are within time budget.
/// If now is before the `bucket window + n`, we have run out of budget.
///
/// This is inspired by the generic cell rate algorithm (GCRA) and works
/// exactly the same as a leaky-bucket.
end: Instant,
}
impl LeakyBucketState {
fn new(now: Instant) -> Self {
Self { end: now }
}
fn should_retain(&self, now: Instant) -> bool {
// if self.end is after now, the bucket is not empty
now < self.end
}
fn check(&mut self, config: &LeakyBucketConfigInner, now: Instant, n: f64) -> bool {
let start = self.end - config.bucket_width;
let n = config.time_cost.mul_f64(n);
// start end
// | start+n | end+n
// | / | /
// ------{o-[---------o-}--]----o----
// now1 ^ now2 ^ ^ now3
//
// at now1, the bucket would be completely filled if we add n tokens.
// at now2, the bucket would be partially filled if we add n tokens.
// at now3, the bucket would start completely empty before we add n tokens.
if self.end + n <= now {
self.end = now + n;
true
} else if start + n <= now {
self.end += n;
true
} else {
false
bucket_width,
// start a bit early to avoid certain underflow issues
epoch: Instant::now() - 2 * bucket_width,
refill_rate: Duration::ZERO,
}
}
}
@@ -155,51 +98,56 @@ mod tests {
use std::time::Duration;
use tokio::time::Instant;
use utils::leaky_bucket::LeakyBucketState;
use super::{LeakyBucketConfig, LeakyBucketConfigInner, LeakyBucketState};
use super::LeakyBucketConfig;
#[tokio::test(start_paused = true)]
async fn check() {
let config: LeakyBucketConfigInner = LeakyBucketConfig::new(500.0, 2000.0).into();
let config: utils::leaky_bucket::LeakyBucketConfig =
LeakyBucketConfig::new(500.0, 2000.0).into();
assert_eq!(config.time_cost, Duration::from_millis(2));
assert_eq!(config.bucket_width, Duration::from_secs(4));
let mut bucket = LeakyBucketState::new(Instant::now());
let mut bucket = LeakyBucketState::new(Instant::now() - config.epoch);
// should work for 2000 requests this second
for _ in 0..2000 {
assert!(bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_ok());
}
assert!(!bucket.check(&config, Instant::now(), 1.0));
assert_eq!(bucket.end - Instant::now(), config.bucket_width);
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_err());
assert_eq!(
bucket.end - (Instant::now() - config.epoch),
config.bucket_width
);
// in 1ms we should drain 0.5 tokens.
// make sure we don't lose any tokens
tokio::time::advance(Duration::from_millis(1)).await;
assert!(!bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_err());
tokio::time::advance(Duration::from_millis(1)).await;
assert!(bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_ok());
// in 10ms we should drain 5 tokens
tokio::time::advance(Duration::from_millis(10)).await;
for _ in 0..5 {
assert!(bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_ok());
}
assert!(!bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_err());
// in 10s we should drain 5000 tokens
// but cap is only 2000
tokio::time::advance(Duration::from_secs(10)).await;
for _ in 0..2000 {
assert!(bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_ok());
}
assert!(!bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_err());
// should sustain 500rps
for _ in 0..2000 {
tokio::time::advance(Duration::from_millis(10)).await;
for _ in 0..5 {
assert!(bucket.check(&config, Instant::now(), 1.0));
assert!(bucket.add_tokens(&config, Instant::now(), 1.0).is_ok());
}
}
}