Line data Source code
1 : use std::{
2 : hash::Hash,
3 : sync::atomic::{AtomicUsize, Ordering},
4 : };
5 :
6 : use ahash::RandomState;
7 : use dashmap::DashMap;
8 : use rand::{thread_rng, Rng};
9 : use tokio::time::Instant;
10 : use tracing::info;
11 : use utils::leaky_bucket::LeakyBucketState;
12 :
13 : use crate::intern::EndpointIdInt;
14 :
15 : // Simple per-endpoint rate limiter.
16 : pub type EndpointRateLimiter = LeakyBucketRateLimiter<EndpointIdInt>;
17 :
18 : pub struct LeakyBucketRateLimiter<Key> {
19 : map: DashMap<Key, LeakyBucketState, RandomState>,
20 : config: utils::leaky_bucket::LeakyBucketConfig,
21 : access_count: AtomicUsize,
22 : }
23 :
24 : impl<K: Hash + Eq> LeakyBucketRateLimiter<K> {
25 : pub const DEFAULT: LeakyBucketConfig = LeakyBucketConfig {
26 : rps: 600.0,
27 : max: 1500.0,
28 : };
29 :
30 3 : pub fn new_with_shards(config: LeakyBucketConfig, shards: usize) -> Self {
31 3 : Self {
32 3 : map: DashMap::with_hasher_and_shard_amount(RandomState::new(), shards),
33 3 : config: config.into(),
34 3 : access_count: AtomicUsize::new(0),
35 3 : }
36 3 : }
37 :
38 : /// Check that number of connections to the endpoint is below `max_rps` rps.
39 3 : pub(crate) fn check(&self, key: K, n: u32) -> bool {
40 3 : let now = Instant::now();
41 3 :
42 3 : if self.access_count.fetch_add(1, Ordering::AcqRel) % 2048 == 0 {
43 3 : self.do_gc(now);
44 3 : }
45 :
46 3 : let mut entry = self
47 3 : .map
48 3 : .entry(key)
49 3 : .or_insert_with(|| LeakyBucketState { empty_at: now });
50 3 :
51 3 : entry.add_tokens(&self.config, now, n as f64).is_ok()
52 3 : }
53 :
54 3 : fn do_gc(&self, now: Instant) {
55 3 : info!(
56 0 : "cleaning up bucket rate limiter, current size = {}",
57 0 : self.map.len()
58 : );
59 3 : let n = self.map.shards().len();
60 3 : let shard = thread_rng().gen_range(0..n);
61 3 : self.map.shards()[shard]
62 3 : .write()
63 3 : .retain(|_, value| !value.get().bucket_is_empty(now));
64 3 : }
65 : }
66 :
67 : pub struct LeakyBucketConfig {
68 : pub rps: f64,
69 : pub max: f64,
70 : }
71 :
72 : #[cfg(test)]
73 : impl LeakyBucketConfig {
74 1 : pub(crate) fn new(rps: f64, max: f64) -> Self {
75 1 : assert!(rps > 0.0, "rps must be positive");
76 1 : assert!(max > 0.0, "max must be positive");
77 1 : Self { rps, max }
78 1 : }
79 : }
80 :
81 : impl From<LeakyBucketConfig> for utils::leaky_bucket::LeakyBucketConfig {
82 4 : fn from(config: LeakyBucketConfig) -> Self {
83 4 : utils::leaky_bucket::LeakyBucketConfig::new(config.rps, config.max)
84 4 : }
85 : }
86 :
87 : #[cfg(test)]
88 : #[allow(clippy::float_cmp)]
89 : mod tests {
90 : use std::time::Duration;
91 :
92 : use tokio::time::Instant;
93 : use utils::leaky_bucket::LeakyBucketState;
94 :
95 : use super::LeakyBucketConfig;
96 :
97 : #[tokio::test(start_paused = true)]
98 1 : async fn check() {
99 1 : let config: utils::leaky_bucket::LeakyBucketConfig =
100 1 : LeakyBucketConfig::new(500.0, 2000.0).into();
101 1 : assert_eq!(config.cost, Duration::from_millis(2));
102 1 : assert_eq!(config.bucket_width, Duration::from_secs(4));
103 1 :
104 1 : let mut bucket = LeakyBucketState {
105 1 : empty_at: Instant::now(),
106 1 : };
107 1 :
108 1 : // should work for 2000 requests this second
109 2001 : for _ in 0..2000 {
110 2000 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap();
111 2000 : }
112 1 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap_err();
113 1 : assert_eq!(bucket.empty_at - Instant::now(), config.bucket_width);
114 1 :
115 1 : // in 1ms we should drain 0.5 tokens.
116 1 : // make sure we don't lose any tokens
117 1 : tokio::time::advance(Duration::from_millis(1)).await;
118 1 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap_err();
119 1 : tokio::time::advance(Duration::from_millis(1)).await;
120 1 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap();
121 1 :
122 1 : // in 10ms we should drain 5 tokens
123 1 : tokio::time::advance(Duration::from_millis(10)).await;
124 6 : for _ in 0..5 {
125 5 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap();
126 5 : }
127 1 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap_err();
128 1 :
129 1 : // in 10s we should drain 5000 tokens
130 1 : // but cap is only 2000
131 1 : tokio::time::advance(Duration::from_secs(10)).await;
132 2001 : for _ in 0..2000 {
133 2000 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap();
134 2000 : }
135 1 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap_err();
136 1 :
137 1 : // should sustain 500rps
138 2001 : for _ in 0..2000 {
139 2000 : tokio::time::advance(Duration::from_millis(10)).await;
140 12000 : for _ in 0..5 {
141 10000 : bucket.add_tokens(&config, Instant::now(), 1.0).unwrap();
142 10000 : }
143 1 : }
144 1 : }
145 : }
|
