forked from solana-labs/solana
-
Notifications
You must be signed in to change notification settings - Fork 259
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* use rate limit on connectings use rate limit on connectings; missing file * Change connection rate limit to 8/min instead of 4/s * Addressed some feedback from Trent * removed some comments * fix test failures which are opening connections more frequently * moved the flag up * turn off rate limiting to debug CI * Fix CI test failures * differentiate of the two throttling cases in stats: across connections or per ip addr * fmt issues * Addressed some feedback from Trent * Added unit tests Cleanup connection cache rate limiter if exceeding certain threshold missing files CONNECITON_RATE_LIMITER_CLEANUP_THRESHOLD to 100_000 clippy issue clippy issue sort crates * revert Cargo.lock changes * Addressed some feedback from Pankaj
- Loading branch information
1 parent
c0e51b4
commit f54c120
Showing
18 changed files
with
419 additions
and
8 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| use { | ||
| crate::nonblocking::{keyed_rate_limiter::KeyedRateLimiter, rate_limiter::RateLimiter}, | ||
| std::{net::IpAddr, time::Duration}, | ||
| }; | ||
|
|
||
| pub struct ConnectionRateLimiter { | ||
| limiter: KeyedRateLimiter<IpAddr>, | ||
| } | ||
|
|
||
| impl ConnectionRateLimiter { | ||
| /// Create a new rate limiter per IpAddr. The rate is specified as the count per minute to allow for | ||
| /// less frequent connections. | ||
| pub fn new(limit_per_minute: u64) -> Self { | ||
| Self { | ||
| limiter: KeyedRateLimiter::new(limit_per_minute, Duration::from_secs(60)), | ||
| } | ||
| } | ||
|
|
||
| /// Check if the connection from the said `ip` is allowed. | ||
| pub fn is_allowed(&self, ip: &IpAddr) -> bool { | ||
| // Acquire a permit from the rate limiter for the given IP address | ||
| if self.limiter.check_and_update(*ip) { | ||
| debug!("Request from IP {:?} allowed", ip); | ||
| true // Request allowed | ||
| } else { | ||
| debug!("Request from IP {:?} blocked", ip); | ||
| false // Request blocked | ||
| } | ||
| } | ||
|
|
||
| /// retain only keys whose throttle start date is within the throttle interval. | ||
| /// Otherwise drop them as inactive | ||
| pub fn retain_recent(&self) { | ||
| self.limiter.retain_recent() | ||
| } | ||
|
|
||
| /// Returns the number of "live" keys in the rate limiter. | ||
| pub fn len(&self) -> usize { | ||
| self.limiter.len() | ||
| } | ||
|
|
||
| /// Returns `true` if the rate limiter has no keys in it. | ||
| pub fn is_empty(&self) -> bool { | ||
| self.limiter.is_empty() | ||
| } | ||
| } | ||
|
|
||
| /// Connection rate limiter for enforcing connection rates from | ||
| /// all clients. | ||
| pub struct TotalConnectionRateLimiter { | ||
| limiter: RateLimiter, | ||
| } | ||
|
|
||
| impl TotalConnectionRateLimiter { | ||
| /// Create a new rate limiter. The rate is specified as the count per second. | ||
| pub fn new(limit_per_second: u64) -> Self { | ||
| Self { | ||
| limiter: RateLimiter::new(limit_per_second, Duration::from_secs(1)), | ||
| } | ||
| } | ||
|
|
||
| /// Check if a connection is allowed. | ||
| pub fn is_allowed(&mut self) -> bool { | ||
| self.limiter.check_and_update() | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| pub mod test { | ||
| use {super::*, std::net::Ipv4Addr}; | ||
|
|
||
| #[tokio::test] | ||
| async fn test_total_connection_rate_limiter() { | ||
| let mut limiter = TotalConnectionRateLimiter::new(2); | ||
| assert!(limiter.is_allowed()); | ||
| assert!(limiter.is_allowed()); | ||
| assert!(!limiter.is_allowed()); | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_connection_rate_limiter() { | ||
| let limiter = ConnectionRateLimiter::new(4); | ||
| let ip1 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)); | ||
| assert!(limiter.is_allowed(&ip1)); | ||
| assert!(limiter.is_allowed(&ip1)); | ||
| assert!(limiter.is_allowed(&ip1)); | ||
| assert!(limiter.is_allowed(&ip1)); | ||
| assert!(!limiter.is_allowed(&ip1)); | ||
|
|
||
| assert!(limiter.len() == 1); | ||
| let ip2 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)); | ||
| assert!(limiter.is_allowed(&ip2)); | ||
| assert!(limiter.len() == 2); | ||
| assert!(limiter.is_allowed(&ip2)); | ||
| assert!(limiter.is_allowed(&ip2)); | ||
| assert!(limiter.is_allowed(&ip2)); | ||
| assert!(!limiter.is_allowed(&ip2)); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,104 @@ | ||
| use { | ||
| crate::nonblocking::rate_limiter::RateLimiter, | ||
| dashmap::DashMap, | ||
| std::{hash::Hash, time::Duration}, | ||
| }; | ||
|
|
||
| pub struct KeyedRateLimiter<K> { | ||
| limiters: DashMap<K, RateLimiter>, | ||
| interval: Duration, | ||
| limit: u64, | ||
| } | ||
|
|
||
| impl<K> KeyedRateLimiter<K> | ||
| where | ||
| K: Eq + Hash, | ||
| { | ||
| /// Create a keyed rate limiter with `limit` count with a rate limit `interval` | ||
| pub fn new(limit: u64, interval: Duration) -> Self { | ||
| Self { | ||
| limiters: DashMap::default(), | ||
| interval, | ||
| limit, | ||
| } | ||
| } | ||
|
|
||
| /// Check if the connection from the said `key` is allowed to pass through the rate limiter. | ||
| /// When it is allowed, the rate limiter state is updated to reflect it has been | ||
| /// allowed. For a unique request, the caller should call it only once when it is allowed. | ||
| pub fn check_and_update(&self, key: K) -> bool { | ||
| let allowed = match self.limiters.entry(key) { | ||
| dashmap::mapref::entry::Entry::Occupied(mut entry) => { | ||
| let limiter = entry.get_mut(); | ||
| limiter.check_and_update() | ||
| } | ||
| dashmap::mapref::entry::Entry::Vacant(entry) => entry | ||
| .insert(RateLimiter::new(self.limit, self.interval)) | ||
| .value_mut() | ||
| .check_and_update(), | ||
| }; | ||
| allowed | ||
| } | ||
|
|
||
| /// retain only keys whose throttle start date is within the throttle interval. | ||
| /// Otherwise drop them as inactive | ||
| pub fn retain_recent(&self) { | ||
| let now = tokio::time::Instant::now(); | ||
| self.limiters.retain(|_key, limiter| { | ||
| now.duration_since(*limiter.throttle_start_instant()) <= self.interval | ||
| }); | ||
| } | ||
|
|
||
| /// Returns the number of "live" keys in the rate limiter. | ||
| pub fn len(&self) -> usize { | ||
| self.limiters.len() | ||
| } | ||
|
|
||
| /// Returns `true` if the rate limiter has no keys in it. | ||
| pub fn is_empty(&self) -> bool { | ||
| self.limiters.is_empty() | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| pub mod test { | ||
| use {super::*, tokio::time::sleep}; | ||
|
|
||
| #[allow(clippy::len_zero)] | ||
| #[tokio::test] | ||
| async fn test_rate_limiter() { | ||
| let limiter = KeyedRateLimiter::<u64>::new(2, Duration::from_millis(100)); | ||
| assert!(limiter.len() == 0); | ||
| assert!(limiter.is_empty()); | ||
| assert!(limiter.check_and_update(1)); | ||
| assert!(limiter.check_and_update(1)); | ||
| assert!(!limiter.check_and_update(1)); | ||
| assert!(limiter.len() == 1); | ||
| assert!(limiter.check_and_update(2)); | ||
| assert!(limiter.check_and_update(2)); | ||
| assert!(!limiter.check_and_update(2)); | ||
| assert!(limiter.len() == 2); | ||
|
|
||
| // sleep 150 ms, the throttle parameters should have been reset. | ||
| sleep(Duration::from_millis(150)).await; | ||
| assert!(limiter.len() == 2); | ||
|
|
||
| assert!(limiter.check_and_update(1)); | ||
| assert!(limiter.check_and_update(1)); | ||
| assert!(!limiter.check_and_update(1)); | ||
|
|
||
| assert!(limiter.check_and_update(2)); | ||
| assert!(limiter.check_and_update(2)); | ||
| assert!(!limiter.check_and_update(2)); | ||
| assert!(limiter.len() == 2); | ||
|
|
||
| // sleep another 150 and clean outdatated, key 2 will be removed | ||
| sleep(Duration::from_millis(150)).await; | ||
| assert!(limiter.check_and_update(1)); | ||
| assert!(limiter.check_and_update(1)); | ||
| assert!(!limiter.check_and_update(1)); | ||
|
|
||
| limiter.retain_recent(); | ||
| assert!(limiter.len() == 1); | ||
| } | ||
| } |
Oops, something went wrong.