Skip to content

Commit

Permalink
build: Update itertools to 0.13.0 (#3993)
Browse files Browse the repository at this point in the history
Updates itertools to the latest version. `group_by` has been replaced by
`chunk_by` to clarify that this method only collects _consecutive_
elements with the same key.

To improve on this, this PR uses `GroupingMap`, which offers more
efficient folding operations such as `min_by_key` directly. This
implements the desired outcome of collecing the shortest quota for each
unique `KeyRef`.
  • Loading branch information
jan-auer authored Sep 5, 2024
1 parent 61d5fb9 commit 4131a20
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 20 deletions.
31 changes: 20 additions & 11 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ hyper-util = { version = "0.1.7", features = ["tokio"] }
indexmap = "2.2.5"
insta = { version = "1.31.0", features = ["json", "redactions", "ron"] }
ipnetwork = "0.20.0"
itertools = "0.10.5"
itertools = "0.13.0"
json-forensics = "0.1.1"
lru = "0.9.0"
maxminddb = "0.23.0"
Expand Down
18 changes: 10 additions & 8 deletions relay-quotas/src/global.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
use std::sync::{Mutex, OnceLock, PoisonError};

use crate::RedisQuota;
use itertools::Itertools;
use relay_base_schema::metrics::MetricNamespace;
use relay_redis::redis::Script;
use relay_redis::{PooledClient, RedisError};

use crate::RedisQuota;

/// Default percentage of the quota limit to reserve from Redis as a local cache.
const DEFAULT_BUDGET_RATIO: f32 = 0.001;

Expand Down Expand Up @@ -37,11 +38,12 @@ impl GlobalRateLimits {
let mut ratelimited = vec![];
let mut not_ratelimited = vec![];

for (key, quotas) in &quotas.iter().group_by(|quota| KeyRef::new(quota)) {
let Some(quota) = quotas.min_by_key(|quota| quota.limit()) else {
continue;
};
let min_by_keyref = quotas
.iter()
.into_grouping_map_by(|q| KeyRef::new(q))
.min_by_key(|_, q| q.limit());

for (key, quota) in min_by_keyref {
let val = guard.entry_ref(&key).or_default();

if val.is_rate_limited(client, quota, key, quantity as u64)? {
Expand Down Expand Up @@ -243,6 +245,7 @@ impl Default for GlobalRateLimit {

#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use std::time::Duration;

use super::*;
Expand Down Expand Up @@ -312,13 +315,12 @@ mod tests {
.unwrap();

// Only the quotas that are less than the quantity gets ratelimited.
assert_eq!(rate_limited_quotas.len(), 2);
assert_eq!(
vec![100, 150],
BTreeSet::from([100, 150]),
rate_limited_quotas
.iter()
.map(|quota| quota.limit())
.collect_vec(),
.collect()
);
}

Expand Down

0 comments on commit 4131a20

Please sign in to comment.