feat(quotas): Add global throughput limit

CHANGELOG.md

@@ -1,9 +1,10 @@
 # Changelog
 
-## Unreleased
+## 23.12.1
 
 **Features**:
 
+- Add a global throughput rate limiter for metric buckets. ([#2854](https://github.com/getsentry/relay/pull/2854))
 - Group db spans with repeating logical conditions together. ([#2929](https://github.com/getsentry/relay/pull/2929))
 
 **Bug Fixes**:

relay-quotas/Cargo.toml

@@ -14,6 +14,7 @@ default = []
 redis = ["dep:thiserror", "dep:relay-log", "relay-redis/impl"]
 
 [dependencies]
+hashbrown = { workspace = true }
 relay-base-schema = { path = "../relay-base-schema" }
 relay-common = { path = "../relay-common" }
 relay-log = { path = "../relay-log", optional = true }

relay-quotas/src/global.rs

@@ -0,0 +1,256 @@
+use std::cmp::Ordering;
+use std::sync::{Arc, Mutex, OnceLock, RwLock};
+
+use relay_common::time::UnixTimestamp;
+use relay_redis::redis::Script;
+use relay_redis::{PooledClient, RedisError};
+
+use crate::RedisQuota;
+
+fn load_global_lua_script() -> &'static Script {
+    static SCRIPT: OnceLock<Script> = OnceLock::new();
+    SCRIPT.get_or_init(|| Script::new(include_str!("global_quota.lua")))
+}
+
+fn current_slot(window: u64) -> usize {
+    UnixTimestamp::now()
+        .as_secs()
+        .checked_div(window)
+        .unwrap_or_default() as usize
+}
+
+/// Counters used as a cache for global quotas.
+///
+/// When we want to ratelimit across all relay-instances, we need to use redis to synchronize.
+/// Calling Redis every time we want to check if an item should be ratelimited would be very expensive,
+/// which is why we have this cache. It works by 'taking' a certain budget from redis, by pre-incrementing
+/// a global counter. We Put the amount we pre-incremented into this local cache and count down until
+/// we have no more budget, then we ask for more from redis. If we find the global counter is above
+/// the quota limit, we will ratelimit the item.
+#[derive(Default)]
+pub struct GlobalRateLimits {
+    counters: RwLock<hashbrown::HashMap<BudgetKey, Arc<Mutex<SlottedCounter>>>>,
+}
+
+impl GlobalRateLimits {
+    /// Returns `true` if the global quota should be ratelimited.
+    ///
+    /// Certain errors can be resolved by syncing to redis, so in those cases
+    /// we try again to decrement the budget after syncing.
+    pub fn is_rate_limited(
+        &self,
+        client: &mut PooledClient,
+        quota: &RedisQuota,
+        quantity: usize,
+    ) -> Result<bool, RedisError> {
+        let Some(limit) = quota.limit else {
+            return Ok(false);
+        };
+
+        let key = BudgetKeyRef::new(quota);
+
+        let opt_val = self.counters.read().unwrap().get(&key).cloned();
+        let val = match opt_val {
+            Some(val) => val,
+            None => Arc::clone(self.counters.write().unwrap().entry_ref(&key).or_default()),
+        };
+
+        let mut lock = val.lock().unwrap();
+        lock.is_rate_limited(client, quota, quantity, limit)
+    }
+}
+
+/// Key for storing global quota-budgets locally.
+///
+/// Note: must not be used in redis. For that, use RedisQuota.key().
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct BudgetKey {
+    prefix: String,
+    window: u64,
+}
+
+impl From<&BudgetKeyRef<'_>> for BudgetKey {
+    fn from(value: &BudgetKeyRef<'_>) -> Self {
+        BudgetKey {
+            prefix: value.prefix.to_owned(),
+            window: value.window,
+        }
+    }
+}
+
+/// Used to look up a hashmap of [`BudgetKey`]-keys without a string allocation.
+///
+/// This works due to the 'Equivalent' trait in the hashbrown crate.
+#[derive(Clone, Copy, Hash)]
+struct BudgetKeyRef<'a> {
+    prefix: &'a str,
+    window: u64,
+}
+
+impl<'a> BudgetKeyRef<'a> {
+    fn new(quota: &'a RedisQuota) -> Self {
+        Self {
+            prefix: quota.prefix(),
+            window: quota.window(),
+        }
+    }
+}
+
+impl hashbrown::Equivalent<BudgetKey> for BudgetKeyRef<'_> {
+    fn equivalent(&self, key: &BudgetKey) -> bool {
+        self.prefix == key.prefix && self.window == key.window
+    }
+}
+
+struct SlottedCounter {
+    slot: usize,
+    counter: Counter,
+}
+
+impl SlottedCounter {
+    pub fn new() -> Self {
+        Self {
+            slot: 0,
+            counter: Counter::new(),
+        }
+    }
+
+    pub fn is_rate_limited(
+        &mut self,
+        client: &mut PooledClient,
+        quota: &RedisQuota,
+        quantity: usize,
+        limit: u64,
+    ) -> Result<bool, RedisError> {
+        let current_slot = current_slot(quota.window());
+
+        match self.slot.cmp(&current_slot) {
+            Ordering::Greater => {
+                // TODO double check logic here
+                // in theory this should never happen only if someone messes with the system time
+                // be safe and dont rate limit
+                relay_log::error!("budget slot ahead of current slot");
+                Ok(false)
+            }
+            Ordering::Less => {
+                self.counter = Counter::new();
+                self.slot = current_slot;
+
+                self.counter.is_rate_limited(client, limit, quantity, quota)
+            }
+            Ordering::Equal => self.counter.is_rate_limited(client, limit, quantity, quota),
+        }
+    }
+}
+
+impl Default for SlottedCounter {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Represents the local budget taken from a global quota.
+struct Counter {
+    local_counter: LocalCounter,
+    redis_counter: RedisCounter,
+}
+
+impl Counter {
+    pub fn new() -> Self {
+        Self {
+            local_counter: LocalCounter::new(),
+            redis_counter: RedisCounter::new(),
+        }
+    }
+
+    pub fn is_rate_limited(
+        &mut self,
+        client: &mut PooledClient,
+        limit: u64,
+        quantity: usize,
+        quota: &RedisQuota,
+    ) -> Result<bool, RedisError> {
+        if self.local_counter.try_consume(quantity) {
+            return Ok(false);
+        }
+
+        if !self.redis_counter.can_satisfy(quantity, limit) {
+            return Ok(true);
+        }
+
+        let budget_to_reserve = quantity.max(self.default_request_size_based_on_limit());
+        let reserved = self
+            .redis_counter
+            .try_reserve(client, budget_to_reserve, limit, quota)? as usize;
+
+        self.local_counter.increase_budget(reserved);
+        Ok(!self.local_counter.try_consume(quantity))
+    }
+
+    fn default_request_size_based_on_limit(&self) -> usize {
+        100
+    }
+}
+
+struct LocalCounter {
+    budget: usize,
+}
+
+impl LocalCounter {
+    fn new() -> Self {
+        Self { budget: 0 }
+    }
+
+    fn try_consume(&mut self, quantity: usize) -> bool {
+        if self.budget >= quantity {
+            self.budget -= quantity;
+            true
+        } else {
+            false
+        }
+    }
+
+    fn increase_budget(&mut self, quantity: usize) {
+        self.budget = self.budget.saturating_add(quantity);
+    }
+}
+
+struct RedisCounter {
+    last_seen: u64,
+}
+
+impl RedisCounter {
+    fn new() -> Self {
+        Self { last_seen: 0 }
+    }
+
+    fn can_satisfy(&self, quantity: usize, limit: u64) -> bool {
+        self.last_seen + quantity as u64 <= limit
+    }
+
+    fn try_reserve(
+        &mut self,
+        client: &mut PooledClient,
+        quantity: usize,
+        limit: u64,
+        quota: &RedisQuota,
+    ) -> Result<u64, RedisError> {
+        let script = load_global_lua_script();
+
+        let redis_key = quota.key();
+        let expiry = UnixTimestamp::now().as_secs() + quota.window();
+
+        let (budget, redis_count): (u64, u64) = script
+            .prepare_invoke()
+            .key(redis_key.as_str())
+            .arg(limit)
+            .arg(expiry)
+            .arg(quantity)
+            .invoke(&mut client.connection()?)
+            .map_err(RedisError::Redis)?;
+
+        self.last_seen = redis_count;
+
+        Ok(budget)
+    }
+}

relay-quotas/src/global_quota.lua

@@ -0,0 +1,66 @@
+-- Global quota system.
+--
+--
+-- Input:
+--
+-- ``KEY``:
+--  * [string] Key of the counter.
+--
+-- ``ARGV``:
+--  * [number]  Quota limit. will not go over this limit while taking budget.
+--  * [number]  Absolute Expiration time as Unix timestamp (secs since 1.1.1970 ) for the key.
+--  * [number]  Quantity we want to take. Limited by the Quota limit.
+--
+-- Output:
+--
+-- A two-element table containing the following:
+--
+-- 1. Reserved Budget (number):
+--    - The amount of budget that has been successfully allocated based on the request.
+--    - This value is 0 if the current quota limit has already been reached or exceeded.
+--    - Otherwise, it reflects the actual budget allocated, which may be less than or equal to the requested budget,
+--      depending on the available quota (headroom).
+--
+-- 2. Redis count (number):
+--    - Represents the value of the counter after we have allocated our budget.
+--
+--
+-- Made to work with a local cache of quota limit. The caller will "take" a certain budget
+-- from the global redis counter, given that the limit have not been exceeded. This is to dramatically
+-- reduce the amount of redis calls while at the same time not overshooting the quota.
+-- We return both the size of the taken budget and the size of the redis count.
+-- The reason we return the size of the redis count is to avoid asking for more budget when
+-- the previously seen redis count is still bigger than the limit.
+--
+--
+-- The redis keys are unique to their timeslot, which is why we let them expire in order to not
+-- fill up redis with dead keys.
+
+
+-- The key to the global quota.
+local key = KEYS[1]
+-- The max amount that we want to take within the given slot. We won't take a budget if
+-- the count is higher than the limit.
+local limit = tonumber(ARGV[1])
+-- When the redis key/val should be deleted.
+local expiry = tonumber(ARGV[2])
+-- The budget that the caller intends to take. Will be capped if too close to the limit.
+local requested_budget = tonumber(ARGV[3])
+
+local redis_count = tonumber(redis.call('GET', key) or 0)
+
+if redis_count >= limit then
+    return {0, redis_count}
+else
+    -- Ensures the budget is not more than the quantity needed to hit the limit.
+    local headroom = limit - redis_count
+    local budget = math.min(headroom, requested_budget)
+    redis.call('INCRBY', key, budget)
+    if redis_count == 0 then
+        redis.call('EXPIREAT', key, expiry)
+    end
+
+    return {budget, redis_count + budget}
+end
+
+

relay-quotas/src/lib.rs

@@ -11,9 +11,14 @@
 /// typically happens for disabled keys, projects, or organizations.
 const REJECT_ALL_SECS: u64 = 60;
 
+#[cfg(feature = "redis")]
+mod global;
 mod quota;
 mod rate_limit;
 
+#[cfg(feature = "redis")]
+pub use self::global::*;
+
 pub use self::quota::*;
 pub use self::rate_limit::*;