Skip to main content

relay_quotas/
redis.rs

1use std::fmt::{self, Debug, Write};
2use std::sync::Arc;
3
4use itertools::Itertools;
5use relay_base_schema::metrics::MetricNamespace;
6use relay_base_schema::organization::OrganizationId;
7use relay_common::time::UnixTimestamp;
8use relay_log::protocol::value;
9use relay_redis::redis::{self, FromRedisValue, ParsingError, Script};
10use relay_redis::{AsyncRedisClient, RedisError, RedisScripts};
11use thiserror::Error;
12
13use crate::cache::OpportunisticQuotaCache;
14use crate::quota::{ItemScoping, Quota, QuotaScope};
15use crate::rate_limit::{RateLimit, RateLimits, RetryAfter};
16use crate::statsd::{QuotaCounters, QuotaTimers};
17use crate::{REJECT_ALL_SECS, cache};
18
19/// The `grace` period allows accommodating for clock drift in TTL
20/// calculation since the clock on the Redis instance used to store quota
21/// metrics may not be in sync with the computer running this code.
22const GRACE: u64 = 60;
23
24/// An error returned by [`RedisRateLimiter`].
25#[derive(Debug, Error)]
26#[error("failed to communicate with redis")]
27pub struct RateLimitingError(
28    #[from]
29    #[source]
30    pub RedisError,
31);
32
33/// Creates a refund key for a given counter key.
34///
35/// Refund keys are used to track credits that should be applied to a quota,
36/// allowing for more flexible quota management.
37fn get_refunded_quota_key(counter_key: &str) -> String {
38    format!("r:{counter_key}")
39}
40
41/// A transparent wrapper around an Option that only displays `Some`.
42struct OptionalDisplay<T>(Option<T>);
43
44impl<T> fmt::Display for OptionalDisplay<T>
45where
46    T: fmt::Display,
47{
48    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49        match self.0 {
50            Some(ref value) => write!(f, "{value}"),
51            None => Ok(()),
52        }
53    }
54}
55
56/// Owned version of [`RedisQuota`].
57#[derive(Debug, Clone)]
58pub struct OwnedRedisQuota {
59    /// The original quota.
60    quota: Quota,
61    /// Scopes of the item being tracked.
62    scoping: ItemScoping,
63    /// The Redis key prefix mapped from the quota id.
64    prefix: Arc<str>,
65    /// The Redis window in seconds mapped from the quota.
66    window: u64,
67    /// The quantity being checked.
68    quantity: u64,
69    /// The ingestion timestamp determining the rate limiting bucket.
70    timestamp: UnixTimestamp,
71}
72
73impl OwnedRedisQuota {
74    /// Returns an instance of [`RedisQuota`] which borrows from this [`OwnedRedisQuota`].
75    pub fn build_ref(&self) -> RedisQuota<'_> {
76        RedisQuota {
77            quota: &self.quota,
78            scoping: &self.scoping,
79            prefix: Arc::clone(&self.prefix),
80            window: self.window,
81            quantity: self.quantity,
82            timestamp: self.timestamp,
83        }
84    }
85}
86
87/// Reference to information required for tracking quotas in Redis.
88#[derive(Debug, Clone, Eq, PartialEq)]
89pub struct RedisQuota<'a> {
90    /// The original quota.
91    quota: &'a Quota,
92    /// Scopes of the item being tracked.
93    scoping: &'a ItemScoping,
94    /// The Redis key prefix mapped from the quota id.
95    prefix: Arc<str>,
96    /// The Redis window in seconds mapped from the quota.
97    window: u64,
98    /// The quantity being checked.
99    quantity: u64,
100    /// The ingestion timestamp determining the rate limiting bucket.
101    timestamp: UnixTimestamp,
102}
103
104impl<'a> RedisQuota<'a> {
105    /// Creates a new [`RedisQuota`] from a [`Quota`], item scoping, and timestamp.
106    ///
107    /// Returns `None` if the quota cannot be tracked in Redis because it's missing
108    /// required fields (ID or window). This allows forward compatibility with
109    /// future quota types.
110    pub fn new(
111        quota: &'a Quota,
112        quantity: u64,
113        scoping: &'a ItemScoping,
114        timestamp: UnixTimestamp,
115    ) -> Option<Self> {
116        // These fields indicate that we *can* track this quota.
117        let prefix = quota.id.clone()?;
118        let window = quota.window?;
119
120        Some(Self {
121            quota,
122            scoping,
123            prefix,
124            quantity,
125            window,
126            timestamp,
127        })
128    }
129
130    /// Converts this [`RedisQuota`] to an [`OwnedRedisQuota`] leaving the original
131    /// struct in place.
132    pub fn build_owned(&self) -> OwnedRedisQuota {
133        OwnedRedisQuota {
134            quota: self.quota.clone(),
135            scoping: self.scoping.clone(),
136            prefix: Arc::clone(&self.prefix),
137            window: self.window,
138            quantity: self.quantity,
139            timestamp: self.timestamp,
140        }
141    }
142
143    /// Returns the window size of the quota in seconds.
144    pub fn window(&self) -> u64 {
145        self.window
146    }
147
148    /// Returns the prefix of the quota used for Redis key generation.
149    pub fn prefix(&self) -> &str {
150        &self.prefix
151    }
152
153    /// Returns the quantity to rate limit.
154    pub fn quantity(&self) -> u64 {
155        self.quantity
156    }
157
158    /// Returns the limit value formatted for Redis.
159    ///
160    /// Returns `-1` for unlimited quotas or when the limit doesn't fit into an `i64`.
161    /// Otherwise, returns the limit value as an `i64`.
162    pub fn limit(&self) -> i64 {
163        self.limit
164            // If it does not fit into i64, treat as unlimited:
165            .and_then(|limit| limit.try_into().ok())
166            .unwrap_or(-1)
167    }
168
169    fn shift(&self) -> u64 {
170        self.scoping.organization_id.value() % self.window
171    }
172
173    /// Returns the current time slot of the quota based on the timestamp.
174    ///
175    /// Slots are used to determine the time bucket for rate limiting.
176    pub fn slot(&self) -> u64 {
177        (self.timestamp.as_secs() - self.shift()) / self.window
178    }
179
180    /// Returns the timestamp when the current quota window will expire.
181    pub fn expiry(&self) -> UnixTimestamp {
182        let next_slot = self.slot() + 1;
183        let next_start = next_slot * self.window + self.shift();
184        UnixTimestamp::from_secs(next_start)
185    }
186
187    /// Returns when the Redis key should expire.
188    ///
189    /// This is the expiry time plus a grace period.
190    pub fn key_expiry(&self) -> u64 {
191        self.expiry().as_secs() + GRACE
192    }
193
194    /// Returns the Redis key for this quota.
195    ///
196    /// The key includes the quota ID, organization ID, and other scoping information
197    /// based on the quota's scope type. Keys are structured to ensure proper isolation
198    /// between different organizations and scopes.
199    pub fn key(&self) -> QuotaCacheKey {
200        // The subscope id is only formatted into the key if the quota is not organization-scoped.
201        // The organization id is always included.
202        let subscope = match self.quota.scope {
203            QuotaScope::Organization => None,
204            scope => self.scoping.scope_id(scope),
205        };
206
207        QuotaCacheKey {
208            id: Arc::clone(&self.prefix),
209            org: self.scoping.organization_id,
210            subscope,
211            namespace: self.namespace,
212            slot: self.slot(),
213            dimensions: self.scoping.dimensions.to_string(self.quota),
214        }
215    }
216
217    /// Returns the maximum cardinality of the quota dimensions that we will support. This is
218    /// necessary to ensure we don't end up with huge numbers of buckets for some set of dimensions.
219    /// If we don't have any dimensions, than all values will hash to the same bucket, so we can
220    /// just say the cardinality is 1 in that case.
221    pub fn max_dimensions_cardinality(&self) -> u32 {
222        if let Some(group_by) = &self.group_by {
223            group_by.max_cardinality
224        } else {
225            1
226        }
227    }
228
229    /// Returns a [`cache::Quota`] built from this [`RedisQuota`].
230    fn for_cache(&self) -> cache::Quota<QuotaCacheKey> {
231        cache::Quota {
232            limit: self.limit(),
233            window: self.window,
234            key: self.key(),
235            expiry: UnixTimestamp::from_secs(self.key_expiry()),
236        }
237    }
238}
239
240impl std::ops::Deref for RedisQuota<'_> {
241    type Target = Quota;
242
243    fn deref(&self) -> &Self::Target {
244        self.quota
245    }
246}
247
248/// A key which uniquely identifies a quota.
249///
250/// Can be used as a Redis cache key by using the [`fmt::Display`] trait.
251///
252/// See also: [`RedisQuota::key`].
253#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
254pub struct QuotaCacheKey {
255    id: Arc<str>,
256    org: OrganizationId,
257    subscope: Option<u64>,
258    namespace: Option<MetricNamespace>,
259    slot: u64,
260    dimensions: Option<String>,
261}
262
263struct RedisQuotaKey {
264    redis_key: String,
265    hash_key: Option<String>,
266}
267
268impl QuotaCacheKey {
269    fn to_redis_quota_key(&self) -> RedisQuotaKey {
270        let mut redis_key = String::new();
271
272        write!(
273            &mut redis_key,
274            "quota:{id}{{{org}}}{subscope}{namespace}:{slot}",
275            id = self.id,
276            org = self.org,
277            subscope = OptionalDisplay(self.subscope),
278            namespace = OptionalDisplay(self.namespace),
279            slot = self.slot,
280        )
281        .expect("should be infallible");
282
283        RedisQuotaKey {
284            redis_key,
285            hash_key: self.dimensions.clone(),
286        }
287    }
288}
289
290/// A service that executes quotas and checks for rate limits in a shared cache.
291///
292/// Quotas handle tracking a project's usage and respond whether a project has been
293/// configured to throttle incoming data if they go beyond the specified quota.
294///
295/// Quotas can specify a window to be tracked in, such as per minute or per hour. Additionally,
296/// quotas allow to specify the data categories they apply to, for example error events or
297/// attachments. For more information on quota parameters, see [`Quota`].
298///
299/// Requires the `redis` feature.
300#[derive(Clone)]
301pub struct RedisRateLimiter {
302    client: AsyncRedisClient,
303    cache: Option<Arc<OpportunisticQuotaCache<QuotaCacheKey>>>,
304    script: &'static Script,
305    max_limit: Option<u64>,
306}
307
308impl RedisRateLimiter {
309    /// Creates a new [`RedisRateLimiter`] instance.
310    pub fn new(client: AsyncRedisClient) -> Self {
311        RedisRateLimiter {
312            client,
313            cache: None,
314            script: RedisScripts::load_is_rate_limited(),
315            max_limit: None,
316        }
317    }
318
319    /// Sets the maximum rate limit in seconds.
320    ///
321    /// By default, this rate limiter will return rate limits based on the quotas' `window` fields.
322    /// If a maximum rate limit is set, the returned rate limit will be bounded by this value.
323    pub fn max_limit(mut self, max_limit: Option<u64>) -> Self {
324        self.max_limit = max_limit;
325        self
326    }
327
328    /// Enables an opportunistic cache for quotas.
329    ///
330    /// The opportunistic cache, opportunistically serves quotas from a local cache, reducing the
331    /// load on Redis heavily.
332    ///
333    /// Caching considers a ratio of the remaining quota to be available and periodically
334    /// synchronizes with Redis.
335    pub fn cache(mut self, cache_ratio: Option<f32>, max: Option<f32>) -> Self {
336        self.cache = cache_ratio
337            .map(OpportunisticQuotaCache::new)
338            .map(|c| c.with_max(max))
339            .map(Arc::new);
340
341        self
342    }
343
344    /// Checks whether any of the quotas in effect have been exceeded and records consumption.
345    ///
346    /// By invoking this method, the caller signals that data is being ingested and needs to be
347    /// counted against the quota. This increment happens atomically if none of the quotas have been
348    /// exceeded. Otherwise, a rate limit is returned and data is not counted against the quotas.
349    ///
350    /// If no key is specified, then only organization-wide and project-wide quotas are checked. If
351    /// a key is specified, then key-quotas are also checked.
352    ///
353    /// When `over_accept_once` is set to `true` and the current quota would be exceeded by the
354    /// provided `quantity`, the data is accepted once and subsequent requests will be rejected
355    /// until the quota refreshes.
356    ///
357    /// A `quantity` of `0` can be used to check if the quota limit has been reached or exceeded
358    /// without incrementing it in the success case. This is useful for checking quotas in a different
359    /// data category.
360    pub async fn is_rate_limited<'a>(
361        &self,
362        quotas: impl IntoIterator<Item = &'a Quota>,
363        item_scoping: &ItemScoping,
364        quantity: usize,
365        over_accept_once: bool,
366    ) -> Result<RateLimits, RateLimitingError> {
367        let timestamp = UnixTimestamp::now();
368        self.is_rate_limited_at(quotas, item_scoping, quantity, over_accept_once, timestamp)
369            .await
370    }
371
372    /// Like [`Self::is_rate_limited`], but checks the quotas at the given `timestamp`.
373    async fn is_rate_limited_at<'a>(
374        &self,
375        quotas: impl IntoIterator<Item = &'a Quota>,
376        item_scoping: &ItemScoping,
377        quantity: usize,
378        over_accept_once: bool,
379        timestamp: UnixTimestamp,
380    ) -> Result<RateLimits, RateLimitingError> {
381        let mut invocation = self.script.prepare_invoke();
382        let mut tracked_quotas = Vec::new();
383        let mut rate_limits = RateLimits::new();
384
385        let quantity = u64::try_from(quantity).unwrap_or(u64::MAX);
386
387        for quota in quotas {
388            if !quota.matches(item_scoping) {
389                // Silently skip all quotas that do not apply to this item.
390            } else if quota.limit == Some(0) {
391                // A zero-sized quota is strongest. Do not call into Redis at all, and do not
392                // increment any keys, as one quota has reached capacity (this is how regular quotas
393                // behave as well).
394                let retry_after = self.retry_after(REJECT_ALL_SECS);
395                rate_limits.add(RateLimit::from_quota(quota, item_scoping, retry_after));
396            } else if let Some(mut quota) =
397                RedisQuota::new(quota, quantity, item_scoping, timestamp)
398            {
399                if let Some(cache) = &self.cache {
400                    quota.quantity = match cache.check_quota(quota.for_cache(), quantity) {
401                        cache::Action::Accept => continue,
402                        cache::Action::Check(quantity) => quantity,
403                    };
404                }
405
406                let redis_quota_key = quota.key().to_redis_quota_key();
407                let mut redis_key = redis_quota_key.redis_key;
408                // Remaining quotas are expected to be track-able in Redis.
409                let refund_key = get_refunded_quota_key(&redis_key);
410
411                let redis_dims_key = redis_quota_key.hash_key;
412
413                if redis_dims_key.is_some() {
414                    redis_key += ":hash";
415                }
416
417                let redis_dims_key = redis_dims_key.as_deref().unwrap_or(crate::EMPTY_DIMENSIONS);
418
419                let max_cardinality = quota.max_dimensions_cardinality();
420
421                invocation.key(redis_key);
422                invocation.key(refund_key);
423
424                invocation.arg(quota.limit());
425                invocation.arg(quota.key_expiry());
426                invocation.arg(quota.quantity);
427                invocation.arg(over_accept_once);
428                invocation.arg(redis_dims_key);
429                invocation.arg(max_cardinality);
430
431                tracked_quotas.push(quota);
432            } else {
433                // This quota is neither a static reject-all, nor can it be tracked in Redis due to
434                // missing fields. We're skipping this for forward-compatibility.
435                relay_log::with_scope(
436                    |scope| scope.set_extra("quota", value::to_value(quota).unwrap()),
437                    || relay_log::warn!("skipping unsupported quota"),
438                )
439            }
440        }
441
442        // Either there are no quotas to run against Redis, or we already have a rate limit from a
443        // zero-sized quota. In either cases, skip invoking the script and return early.
444        if tracked_quotas.is_empty() || rate_limits.is_limited() {
445            return Ok(rate_limits);
446        }
447
448        let mut connection = self.client.get_connection().await?;
449        let result: ScriptResult = invocation
450            .invoke_async(&mut connection)
451            .await
452            .map_err(RedisError::Redis)?;
453
454        for (quota, state) in tracked_quotas.iter().zip(result.0) {
455            if state.is_rejected {
456                // We can calculate the error by comparing how much the cache added to the
457                // quantity with remaining difference of the consumption and limit.
458                let cache_error = {
459                    let remaining = quota.limit().saturating_sub(state.consumed).max(0) as u64;
460                    let cache_quantity = quota.quantity.saturating_sub(quantity);
461
462                    cache_quantity.saturating_sub(remaining)
463                };
464                relay_statsd::metric!(
465                    counter(QuotaCounters::CacheError) += cache_error,
466                    category = item_scoping.category.name(),
467                );
468
469                let retry_after = self.retry_after((quota.expiry() - timestamp).as_secs());
470                rate_limits.add(RateLimit::from_quota(quota, item_scoping, retry_after));
471            } else if let Some(cache) = &self.cache {
472                // Only update the cache if it's really necessary. Quotas which are being rejected,
473                // will not be able to be handled from the cache anyways.
474                cache.set_quota(quota.for_cache(), state.consumed);
475            }
476        }
477        drop(connection);
478
479        if let Some(cache) = &self.cache {
480            let vacuum_start = std::time::Instant::now();
481            if cache.try_vacuum(timestamp) {
482                relay_statsd::metric!(
483                    timer(QuotaTimers::CacheVacuumDuration) = vacuum_start.elapsed()
484                );
485            }
486        }
487
488        Ok(rate_limits)
489    }
490
491    /// Creates a [`RetryAfter`] value that is bounded by the configured [`max_limit`](Self::max_limit).
492    ///
493    /// If a maximum rate limit has been set, the returned value will not exceed that limit.
494    fn retry_after(&self, mut seconds: u64) -> RetryAfter {
495        if let Some(max_limit) = self.max_limit {
496            seconds = std::cmp::min(seconds, max_limit);
497        }
498
499        RetryAfter::from_secs(seconds)
500    }
501}
502
503/// The result returned from the rate limiting Redis script.
504#[derive(Debug)]
505struct ScriptResult(Vec<QuotaState>);
506
507impl FromRedisValue for ScriptResult {
508    fn from_redis_value(v: redis::Value) -> Result<Self, ParsingError> {
509        let seq = v.into_sequence().map_err(|v| {
510            format!("Expected a sequence from the rate limiting script (value was: {v:?})")
511        })?;
512
513        if !seq.len().is_multiple_of(2) {
514            return Err(format!(
515                "Expected an even number of values from the rate limiting script (value was: {seq:?})"
516            ).into());
517        }
518
519        let mut quotas = Vec::with_capacity(seq.len() / 2);
520        for (is_rejected, consumed) in seq.into_iter().tuples() {
521            quotas.push(QuotaState {
522                is_rejected: bool::from_redis_value(is_rejected)?,
523                consumed: i64::from_redis_value(consumed)?,
524            });
525        }
526
527        Ok(Self(quotas))
528    }
529}
530
531/// The state returned from the rate limiting script for a single quota.
532#[derive(Debug)]
533struct QuotaState {
534    /// Whether the quota rejects the request.
535    is_rejected: bool,
536    /// How much of the quota has already been consumed.
537    consumed: i64,
538}
539
540#[cfg(test)]
541mod tests {
542    use std::collections::BTreeSet;
543    use std::time::{SystemTime, UNIX_EPOCH};
544
545    use super::*;
546    use crate::quota::{DataCategories, DataCategory, ReasonCode, Scoping};
547    use crate::rate_limit::RateLimitScope;
548    use crate::{Dimension, DimensionMap, EMPTY_DIMENSIONS, GroupBy, MetricNamespaceScoping};
549    use relay_base_schema::metrics::MetricNamespace;
550    use relay_base_schema::organization::OrganizationId;
551    use relay_base_schema::project::{ProjectId, ProjectKey};
552    use relay_redis::RedisConfigOptions;
553    use relay_redis::redis::AsyncCommands;
554    use smallvec::smallvec;
555
556    fn build_rate_limiter() -> RedisRateLimiter {
557        let url = std::env::var("RELAY_REDIS_URL")
558            .unwrap_or_else(|_| "redis://127.0.0.1:6379".to_owned());
559        let client =
560            AsyncRedisClient::single("test", &url, &RedisConfigOptions::default()).unwrap();
561
562        RedisRateLimiter {
563            client,
564            cache: None,
565            script: RedisScripts::load_is_rate_limited(),
566            max_limit: None,
567        }
568    }
569
570    #[tokio::test]
571    async fn test_zero_size_quotas() {
572        let quotas = &[
573            Quota {
574                id: None,
575                categories: DataCategories::new(),
576                scope: QuotaScope::Organization,
577                scope_id: None,
578                limit: Some(0),
579                window: None,
580                reason_code: Some(ReasonCode::new("get_lost")),
581                namespace: None,
582                group_by: None,
583            },
584            Quota {
585                id: Some("42".into()),
586                categories: DataCategories::new(),
587                scope: QuotaScope::Organization,
588                scope_id: None,
589                limit: None,
590                window: Some(42),
591                reason_code: Some(ReasonCode::new("unlimited")),
592                namespace: None,
593                group_by: None,
594            },
595        ];
596
597        let scoping = ItemScoping {
598            category: DataCategory::Error,
599            scoping: Scoping {
600                organization_id: OrganizationId::new(42),
601                project_id: ProjectId::new(43),
602                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
603                key_id: Some(44),
604            },
605            namespace: MetricNamespaceScoping::None,
606            dimensions: DimensionMap::default(),
607        };
608
609        let rate_limits: Vec<RateLimit> = build_rate_limiter()
610            .is_rate_limited(quotas, &scoping, 1, false)
611            .await
612            .expect("rate limiting failed")
613            .into_iter()
614            .collect();
615
616        assert_eq!(
617            rate_limits,
618            vec![RateLimit {
619                categories: DataCategories::new(),
620                scope: RateLimitScope::Organization(OrganizationId::new(42)),
621                reason_code: Some(ReasonCode::new("get_lost")),
622                retry_after: rate_limits[0].retry_after,
623                namespaces: smallvec![],
624                dimensions: DimensionMap::default(),
625            }]
626        );
627    }
628
629    /// Tests that a quota with and without namespace are counted separately.
630    #[tokio::test]
631    async fn test_namespace_quota() {
632        let quota_limit = 5;
633        let get_quota = |namespace: Option<MetricNamespace>| -> Quota {
634            Quota {
635                id: Some(format!("test_simple_quota_{}", uuid::Uuid::new_v4()).into()),
636                categories: DataCategories::new(),
637                scope: QuotaScope::Organization,
638                scope_id: None,
639                limit: Some(quota_limit),
640                window: Some(600),
641                reason_code: Some(ReasonCode::new(format!("ns: {namespace:?}"))),
642                namespace,
643                group_by: None,
644            }
645        };
646
647        let quotas = &[get_quota(None)];
648        let quota_with_namespace = &[get_quota(Some(MetricNamespace::Sessions))];
649
650        let scoping = ItemScoping {
651            category: DataCategory::Error,
652            scoping: Scoping {
653                organization_id: OrganizationId::new(42),
654                project_id: ProjectId::new(43),
655                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
656                key_id: Some(44),
657            },
658            namespace: MetricNamespaceScoping::Some(MetricNamespace::Sessions),
659            dimensions: DimensionMap::default(),
660        };
661
662        let rate_limiter = build_rate_limiter();
663
664        // First confirm normal behaviour without namespace.
665        for i in 0..10 {
666            let rate_limits: Vec<RateLimit> = rate_limiter
667                .is_rate_limited(quotas, &scoping, 1, false)
668                .await
669                .expect("rate limiting failed")
670                .into_iter()
671                .collect();
672
673            if i < quota_limit {
674                assert_eq!(rate_limits, vec![]);
675            } else {
676                assert_eq!(
677                    rate_limits[0].reason_code,
678                    Some(ReasonCode::new("ns: None"))
679                );
680            }
681        }
682
683        // Then, send identical quota with namespace and confirm it counts separately.
684        for i in 0..10 {
685            let rate_limits: Vec<RateLimit> = rate_limiter
686                .is_rate_limited(quota_with_namespace, &scoping, 1, false)
687                .await
688                .expect("rate limiting failed")
689                .into_iter()
690                .collect();
691
692            if i < quota_limit {
693                assert_eq!(rate_limits, vec![]);
694            } else {
695                assert_eq!(
696                    rate_limits[0].reason_code,
697                    Some(ReasonCode::new("ns: Some(Sessions)"))
698                );
699            }
700        }
701    }
702
703    #[tokio::test]
704    async fn test_simple_quota() {
705        let quotas = &[Quota {
706            id: Some(format!("test_simple_quota_{}", uuid::Uuid::new_v4()).into()),
707            categories: DataCategories::new(),
708            scope: QuotaScope::Organization,
709            scope_id: None,
710            limit: Some(5),
711            window: Some(60),
712            reason_code: Some(ReasonCode::new("get_lost")),
713            namespace: None,
714            group_by: None,
715        }];
716
717        let scoping = ItemScoping {
718            category: DataCategory::Error,
719            scoping: Scoping {
720                organization_id: OrganizationId::new(42),
721                project_id: ProjectId::new(43),
722                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
723                key_id: Some(44),
724            },
725            namespace: MetricNamespaceScoping::None,
726            dimensions: DimensionMap::default(),
727        };
728
729        let rate_limiter = build_rate_limiter();
730
731        for i in 0..10 {
732            let rate_limits: Vec<RateLimit> = rate_limiter
733                .is_rate_limited(quotas, &scoping, 1, false)
734                .await
735                .expect("rate limiting failed")
736                .into_iter()
737                .collect();
738
739            if i >= 5 {
740                assert_eq!(
741                    rate_limits,
742                    vec![RateLimit {
743                        categories: DataCategories::new(),
744                        scope: RateLimitScope::Organization(OrganizationId::new(42)),
745                        reason_code: Some(ReasonCode::new("get_lost")),
746                        retry_after: rate_limits[0].retry_after,
747                        namespaces: smallvec![],
748                        dimensions: DimensionMap::default(),
749                    }]
750                );
751            } else {
752                assert_eq!(rate_limits, vec![]);
753            }
754        }
755    }
756
757    #[tokio::test]
758    async fn test_quantity_0() {
759        let quotas = &[Quota {
760            id: Some(format!("test_quantity_0_{}", uuid::Uuid::new_v4()).into()),
761            categories: DataCategories::new(),
762            scope: QuotaScope::Organization,
763            scope_id: None,
764            limit: Some(1),
765            window: Some(60),
766            reason_code: Some(ReasonCode::new("get_lost")),
767            namespace: None,
768            group_by: None,
769        }];
770
771        let scoping = ItemScoping {
772            category: DataCategory::Error,
773            scoping: Scoping {
774                organization_id: OrganizationId::new(42),
775                project_id: ProjectId::new(43),
776                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
777                key_id: Some(44),
778            },
779            namespace: MetricNamespaceScoping::None,
780            dimensions: DimensionMap::default(),
781        };
782
783        let rate_limiter = build_rate_limiter();
784
785        // limit is 1, so first call not rate limited
786        assert!(
787            !rate_limiter
788                .is_rate_limited(quotas, &scoping, 1, false)
789                .await
790                .unwrap()
791                .is_limited()
792        );
793
794        // quota is now exhausted
795        assert!(
796            rate_limiter
797                .is_rate_limited(quotas, &scoping, 1, false)
798                .await
799                .unwrap()
800                .is_limited()
801        );
802
803        // quota is exhausted, regardless of the quantity
804        assert!(
805            rate_limiter
806                .is_rate_limited(quotas, &scoping, 0, false)
807                .await
808                .unwrap()
809                .is_limited()
810        );
811
812        // quota is exhausted, regardless of the quantity
813        assert!(
814            rate_limiter
815                .is_rate_limited(quotas, &scoping, 1, false)
816                .await
817                .unwrap()
818                .is_limited()
819        );
820    }
821
822    #[tokio::test]
823    async fn test_quota_go_over() {
824        let quotas = &[Quota {
825            id: Some(format!("test_quota_go_over{}", uuid::Uuid::new_v4()).into()),
826            categories: DataCategories::new(),
827            scope: QuotaScope::Organization,
828            scope_id: None,
829            limit: Some(2),
830            window: Some(60),
831            reason_code: Some(ReasonCode::new("get_lost")),
832            namespace: None,
833            group_by: None,
834        }];
835
836        let scoping = ItemScoping {
837            category: DataCategory::Error,
838            scoping: Scoping {
839                organization_id: OrganizationId::new(42),
840                project_id: ProjectId::new(43),
841                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
842                key_id: Some(44),
843            },
844            namespace: MetricNamespaceScoping::None,
845            dimensions: DimensionMap::default(),
846        };
847
848        let rate_limiter = build_rate_limiter();
849
850        // limit is 2, so first call not rate limited
851        let is_limited = rate_limiter
852            .is_rate_limited(quotas, &scoping, 1, true)
853            .await
854            .unwrap()
855            .is_limited();
856        assert!(!is_limited);
857
858        // go over limit, but first call is over-accepted
859        let is_limited = rate_limiter
860            .is_rate_limited(quotas, &scoping, 2, true)
861            .await
862            .unwrap()
863            .is_limited();
864        assert!(!is_limited);
865
866        // quota is exhausted, regardless of the quantity
867        let is_limited = rate_limiter
868            .is_rate_limited(quotas, &scoping, 0, true)
869            .await
870            .unwrap()
871            .is_limited();
872        assert!(is_limited);
873
874        // quota is exhausted, regardless of the quantity
875        let is_limited = rate_limiter
876            .is_rate_limited(quotas, &scoping, 1, true)
877            .await
878            .unwrap()
879            .is_limited();
880        assert!(is_limited);
881    }
882
883    #[tokio::test]
884    async fn test_bails_immediately_without_any_quota() {
885        let scoping = ItemScoping {
886            category: DataCategory::Error,
887            scoping: Scoping {
888                organization_id: OrganizationId::new(42),
889                project_id: ProjectId::new(43),
890                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
891                key_id: Some(44),
892            },
893            namespace: MetricNamespaceScoping::None,
894            dimensions: DimensionMap::default(),
895        };
896
897        let rate_limits: Vec<RateLimit> = build_rate_limiter()
898            .is_rate_limited(&[], &scoping, 1, false)
899            .await
900            .expect("rate limiting failed")
901            .into_iter()
902            .collect();
903
904        assert_eq!(rate_limits, vec![]);
905    }
906
907    #[tokio::test]
908    async fn test_limited_with_unlimited_quota() {
909        let quotas = &[
910            Quota {
911                id: Some("q0".into()),
912                categories: DataCategories::new(),
913                scope: QuotaScope::Organization,
914                scope_id: None,
915                limit: None,
916                window: Some(1),
917                reason_code: Some(ReasonCode::new("project_quota0")),
918                namespace: None,
919                group_by: None,
920            },
921            Quota {
922                id: Some("q1".into()),
923                categories: DataCategories::new(),
924                scope: QuotaScope::Organization,
925                scope_id: None,
926                limit: Some(1),
927                window: Some(1),
928                reason_code: Some(ReasonCode::new("project_quota1")),
929                namespace: None,
930                group_by: None,
931            },
932        ];
933
934        let scoping = ItemScoping {
935            category: DataCategory::Error,
936            scoping: Scoping {
937                organization_id: OrganizationId::new(42),
938                project_id: ProjectId::new(43),
939                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
940                key_id: Some(44),
941            },
942            namespace: MetricNamespaceScoping::None,
943            dimensions: DimensionMap::default(),
944        };
945
946        let rate_limiter = build_rate_limiter();
947
948        for i in 0..1 {
949            let rate_limits: Vec<RateLimit> = rate_limiter
950                .is_rate_limited(quotas, &scoping, 1, false)
951                .await
952                .expect("rate limiting failed")
953                .into_iter()
954                .collect();
955
956            if i == 0 {
957                assert_eq!(rate_limits, &[]);
958            } else {
959                assert_eq!(
960                    rate_limits,
961                    vec![RateLimit {
962                        categories: DataCategories::new(),
963                        scope: RateLimitScope::Organization(OrganizationId::new(42)),
964                        reason_code: Some(ReasonCode::new("project_quota1")),
965                        retry_after: rate_limits[0].retry_after,
966                        namespaces: smallvec![],
967                        dimensions: DimensionMap::default(),
968                    }]
969                );
970            }
971        }
972    }
973
974    #[tokio::test]
975    async fn test_quota_with_quantity() {
976        let quotas = &[Quota {
977            id: Some(format!("test_quantity_quota_{}", uuid::Uuid::new_v4()).into()),
978            categories: DataCategories::new(),
979            scope: QuotaScope::Organization,
980            scope_id: None,
981            limit: Some(500),
982            window: Some(60),
983            reason_code: Some(ReasonCode::new("get_lost")),
984            namespace: None,
985            group_by: None,
986        }];
987
988        let scoping = ItemScoping {
989            category: DataCategory::Error,
990            scoping: Scoping {
991                organization_id: OrganizationId::new(42),
992                project_id: ProjectId::new(43),
993                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
994                key_id: Some(44),
995            },
996            namespace: MetricNamespaceScoping::None,
997            dimensions: DimensionMap::default(),
998        };
999
1000        let rate_limiter = build_rate_limiter();
1001
1002        for i in 0..10 {
1003            let rate_limits: Vec<RateLimit> = rate_limiter
1004                .is_rate_limited(quotas, &scoping, 100, false)
1005                .await
1006                .expect("rate limiting failed")
1007                .into_iter()
1008                .collect();
1009
1010            if i >= 5 {
1011                assert_eq!(
1012                    rate_limits,
1013                    vec![RateLimit {
1014                        categories: DataCategories::new(),
1015                        scope: RateLimitScope::Organization(OrganizationId::new(42)),
1016                        reason_code: Some(ReasonCode::new("get_lost")),
1017                        retry_after: rate_limits[0].retry_after,
1018                        namespaces: smallvec![],
1019                        dimensions: DimensionMap::default(),
1020                    }]
1021                );
1022            } else {
1023                assert_eq!(rate_limits, vec![]);
1024            }
1025        }
1026    }
1027
1028    #[tokio::test]
1029    async fn test_get_redis_key_scoped() {
1030        let quota = Quota {
1031            id: Some("foo".into()),
1032            categories: DataCategories::new(),
1033            scope: QuotaScope::Project,
1034            scope_id: Some("42".into()),
1035            window: Some(2),
1036            limit: Some(0),
1037            reason_code: None,
1038            namespace: None,
1039            group_by: None,
1040        };
1041
1042        let scoping = ItemScoping {
1043            category: DataCategory::Error,
1044            scoping: Scoping {
1045                organization_id: OrganizationId::new(69420),
1046                project_id: ProjectId::new(42),
1047                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1048                key_id: Some(4711),
1049            },
1050            namespace: MetricNamespaceScoping::None,
1051            dimensions: DimensionMap::default(),
1052        };
1053
1054        let timestamp = UnixTimestamp::from_secs(123_123_123);
1055        let redis_quota = RedisQuota::new(&quota, 0, &scoping, timestamp).unwrap();
1056        assert_eq!(
1057            redis_quota.key().to_redis_quota_key().redis_key,
1058            "quota:foo{69420}42:61561561"
1059        );
1060    }
1061
1062    #[tokio::test]
1063    async fn test_get_redis_key_unscoped() {
1064        let quota = Quota {
1065            id: Some("foo".into()),
1066            categories: DataCategories::new(),
1067            scope: QuotaScope::Organization,
1068            scope_id: None,
1069            window: Some(10),
1070            limit: Some(0),
1071            reason_code: None,
1072            namespace: None,
1073            group_by: None,
1074        };
1075
1076        let scoping = ItemScoping {
1077            category: DataCategory::Error,
1078            scoping: Scoping {
1079                organization_id: OrganizationId::new(69420),
1080                project_id: ProjectId::new(42),
1081                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1082                key_id: Some(4711),
1083            },
1084            namespace: MetricNamespaceScoping::None,
1085            dimensions: DimensionMap::default(),
1086        };
1087
1088        let timestamp = UnixTimestamp::from_secs(234_531);
1089        let redis_quota = RedisQuota::new(&quota, 0, &scoping, timestamp).unwrap();
1090        assert_eq!(
1091            redis_quota.key().to_redis_quota_key().redis_key,
1092            "quota:foo{69420}:23453"
1093        );
1094    }
1095
1096    #[tokio::test]
1097    async fn test_large_redis_limit_large() {
1098        let quota = Quota {
1099            id: Some("foo".into()),
1100            categories: DataCategories::new(),
1101            scope: QuotaScope::Organization,
1102            scope_id: None,
1103            window: Some(10),
1104            limit: Some(9223372036854775808), // i64::MAX + 1
1105            reason_code: None,
1106            namespace: None,
1107            group_by: None,
1108        };
1109
1110        let scoping = ItemScoping {
1111            category: DataCategory::Error,
1112            scoping: Scoping {
1113                organization_id: OrganizationId::new(69420),
1114                project_id: ProjectId::new(42),
1115                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1116                key_id: Some(4711),
1117            },
1118            namespace: MetricNamespaceScoping::None,
1119            dimensions: DimensionMap::default(),
1120        };
1121
1122        let timestamp = UnixTimestamp::from_secs(234_531);
1123        let redis_quota = RedisQuota::new(&quota, 0, &scoping, timestamp).unwrap();
1124        assert_eq!(redis_quota.limit(), -1);
1125    }
1126
1127    #[tokio::test]
1128    async fn test_is_rate_limited_script() {
1129        let now = SystemTime::now()
1130            .duration_since(UNIX_EPOCH)
1131            .map(|duration| duration.as_secs())
1132            .unwrap();
1133
1134        let rate_limiter = build_rate_limiter();
1135        let mut conn = rate_limiter.client.get_connection().await.unwrap();
1136
1137        // define a few keys with random seed such that they do not collide with repeated test runs
1138        let foo = format!("foo___{now}");
1139        let r_foo = format!("r:foo___{now}");
1140        let bar = format!("bar___{now}");
1141        let r_bar = format!("r:bar___{now}");
1142        let apple = format!("apple___{now}");
1143        let orange = format!("orange___{now}");
1144        let baz = format!("baz___{now}");
1145
1146        let script = RedisScripts::load_is_rate_limited();
1147
1148        macro_rules! assert_invocation {
1149            ($invocation:expr, $($tt:tt)*) => {{
1150                let result = $invocation
1151                    .invoke_async::<ScriptResult>(&mut conn)
1152                    .await
1153                    .unwrap();
1154
1155                insta::assert_debug_snapshot!(result, $($tt)*);
1156            }};
1157        }
1158
1159        let mut invocation = script.prepare_invoke();
1160        invocation
1161            .key(&foo) // key
1162            .key(&r_foo) // refund key
1163            .key(&bar) // key
1164            .key(&r_bar) // refund key
1165            .arg(1) // limit
1166            .arg(now + 60) // expiry
1167            .arg(1) // quantity
1168            .arg(false) // over accept once
1169            .arg(EMPTY_DIMENSIONS) // dimensions
1170            .arg(1) // max cardinality
1171            .arg(2) // limit
1172            .arg(now + 120) // expiry
1173            .arg(1) // quantity
1174            .arg(false) // over accept once
1175            .arg(EMPTY_DIMENSIONS) // dimensions
1176            .arg(1); // max cardinality
1177
1178        // Craft a new invocation similar to the previous one, but it only applies to the quota
1179        // with a higher limit (2).
1180        let mut invocation2 = script.prepare_invoke();
1181        invocation2
1182            .key(&bar) // key
1183            .key(&r_bar) // refund key
1184            .arg(2) // limit
1185            .arg(now + 120) // expiry
1186            .arg(1) // quantity
1187            .arg(false) // over accept once
1188            .arg(EMPTY_DIMENSIONS) // dimensions
1189            .arg(1); // max cardinality
1190
1191        // 1 quantity used from both quotas.
1192        assert_invocation!(invocation, @r"
1193        ScriptResult(
1194            [
1195                QuotaState {
1196                    is_rejected: false,
1197                    consumed: 1,
1198                },
1199                QuotaState {
1200                    is_rejected: false,
1201                    consumed: 1,
1202                },
1203            ],
1204        )
1205        "
1206        );
1207
1208        // This invocation fails the rate limit on the first quota.
1209        // -> No changes are made to the counters.
1210        assert_invocation!(invocation, @r"
1211        ScriptResult(
1212            [
1213                QuotaState {
1214                    is_rejected: true,
1215                    consumed: 1,
1216                },
1217                QuotaState {
1218                    is_rejected: false,
1219                    consumed: 1,
1220                },
1221            ],
1222        )
1223        "
1224        );
1225
1226        // Another call, same result as before, just making sure there were no changes applied.
1227        assert_invocation!(invocation, @r"
1228        ScriptResult(
1229            [
1230                QuotaState {
1231                    is_rejected: true,
1232                    consumed: 1,
1233                },
1234                QuotaState {
1235                    is_rejected: false,
1236                    consumed: 1,
1237                },
1238            ],
1239        )
1240        "
1241        );
1242
1243        // Using the second invocation which only considers a quota with a higher limit, usage for
1244        // that quota is now 2.
1245        assert_invocation!(invocation2, @r"
1246        ScriptResult(
1247            [
1248                QuotaState {
1249                    is_rejected: false,
1250                    consumed: 2,
1251                },
1252            ],
1253        )
1254        "
1255        );
1256
1257        // Same invocation, but this time the limit is reached, quota should not increase.
1258        assert_invocation!(invocation2, @r"
1259        ScriptResult(
1260            [
1261                QuotaState {
1262                    is_rejected: true,
1263                    consumed: 2,
1264                },
1265            ],
1266        )
1267        "
1268        );
1269
1270        // Check again with the original invocation, this now yields `[1, 2]`.
1271        assert_invocation!(invocation, @r"
1272        ScriptResult(
1273            [
1274                QuotaState {
1275                    is_rejected: true,
1276                    consumed: 1,
1277                },
1278                QuotaState {
1279                    is_rejected: true,
1280                    consumed: 2,
1281                },
1282            ],
1283        )
1284        "
1285        );
1286
1287        assert_eq!(conn.get::<_, String>(&foo).await.unwrap(), "1");
1288
1289        let ttl: u64 = conn.ttl(&foo).await.unwrap();
1290        assert!(ttl >= 59);
1291        assert!(ttl <= 60);
1292
1293        assert_eq!(conn.get::<_, String>(&bar).await.unwrap(), "2");
1294
1295        let ttl: u64 = conn.ttl(&bar).await.unwrap();
1296        assert!(ttl >= 119);
1297        assert!(ttl <= 120);
1298
1299        // make sure "refund/negative" keys haven't been incremented
1300        let () = conn.get(r_foo).await.unwrap();
1301        let () = conn.get(r_bar).await.unwrap();
1302
1303        // Test that refunded quotas work
1304        let () = conn.set(&apple, 5).await.unwrap();
1305
1306        let mut invocation = script.prepare_invoke();
1307        invocation
1308            .key(&orange) // key
1309            .key(&baz) // refund key
1310            .arg(1) // limit
1311            .arg(now + 60) // expiry
1312            .arg(1) // quantity
1313            .arg(false) // over accept once
1314            .arg(EMPTY_DIMENSIONS) // dimensions
1315            .arg(1); // max cardinality
1316
1317        // increment, current quota usage is 1.
1318        assert_invocation!(invocation, @r"
1319        ScriptResult(
1320            [
1321                QuotaState {
1322                    is_rejected: false,
1323                    consumed: 1,
1324                },
1325            ],
1326        )
1327        "
1328        );
1329
1330        // test that it's rate limited without refund.
1331        assert_invocation!(invocation, @r"
1332        ScriptResult(
1333            [
1334                QuotaState {
1335                    is_rejected: true,
1336                    consumed: 1,
1337                },
1338            ],
1339        )
1340        "
1341        );
1342
1343        // Make sure, the counter wasn't incremented.
1344        assert_invocation!(invocation, @r"
1345        ScriptResult(
1346            [
1347                QuotaState {
1348                    is_rejected: true,
1349                    consumed: 1,
1350                },
1351            ],
1352        )
1353        "
1354        );
1355
1356        let mut invocation = script.prepare_invoke();
1357        invocation
1358            .key(&orange) // key
1359            .key(&apple) // refund key
1360            .arg(1) // limit
1361            .arg(now + 60) // expiry
1362            .arg(1) // quantity
1363            .arg(false) // over accept once
1364            .arg(EMPTY_DIMENSIONS) // dimensions
1365            .arg(1); // max cardinality
1366
1367        // test that refund key is used
1368        assert_invocation!(invocation, @r"
1369        ScriptResult(
1370            [
1371                QuotaState {
1372                    is_rejected: false,
1373                    consumed: -3,
1374                },
1375            ],
1376        )
1377        "
1378        );
1379    }
1380
1381    /// Each dimensions key is counted independently within the same Redis key.
1382    #[tokio::test]
1383    async fn test_is_rate_limited_script_dimensions() {
1384        let now = SystemTime::now()
1385            .duration_since(UNIX_EPOCH)
1386            .map(|duration| duration.as_secs())
1387            .unwrap();
1388
1389        let rate_limiter = build_rate_limiter();
1390        let mut conn = rate_limiter.client.get_connection().await.unwrap();
1391
1392        let key = format!("dims___{now}");
1393        let refund_key = format!("r:dims___{now}");
1394
1395        let script = RedisScripts::load_is_rate_limited();
1396
1397        // A single quota with a limit of 1, invoked for two different dimension values.
1398        let invoke = |dims: &str| {
1399            let mut invocation = script.prepare_invoke();
1400
1401            let mut key = key.clone();
1402            if dims != EMPTY_DIMENSIONS {
1403                key += ":hash";
1404            }
1405
1406            invocation
1407                .key(&key) // key
1408                .key(&refund_key) // refund key
1409                .arg(1) // limit
1410                .arg(now + 60) // expiry
1411                .arg(1) // quantity
1412                .arg(false) // over accept once
1413                .arg(dims) // dimensions
1414                .arg(999); // max cardinality
1415            invocation
1416        };
1417
1418        let invoke_async = async |dims: &str, conn: &mut _| {
1419            invoke(dims)
1420                .invoke_async::<ScriptResult>(conn)
1421                .await
1422                .unwrap()
1423                .0
1424        };
1425
1426        // The first check for `:2:a` consumes its bucket...
1427        let result = invoke_async(":2:a", &mut conn).await;
1428        assert!(!result[0].is_rejected);
1429        assert_eq!(result[0].consumed, 1);
1430
1431        // ...and the second one is rejected.
1432        let result = invoke_async(":2:a", &mut conn).await;
1433        assert!(result[0].is_rejected);
1434        assert_eq!(result[0].consumed, 1);
1435
1436        // A different dimensions key is a separate bucket, unaffected by the exhausted one.
1437        let result = invoke_async(":2:b", &mut conn).await;
1438        assert!(!result[0].is_rejected);
1439        assert_eq!(result[0].consumed, 1);
1440
1441        let result = invoke_async(":2:b", &mut conn).await;
1442        assert!(result[0].is_rejected);
1443        assert_eq!(result[0].consumed, 1);
1444
1445        // So is the undimensioned bucket.
1446        let result = invoke_async(EMPTY_DIMENSIONS, &mut conn).await;
1447        assert!(!result[0].is_rejected);
1448        assert_eq!(result[0].consumed, 1);
1449
1450        // All buckets live in the same hash, under the same expiry.
1451        let mut fields: Vec<(String, i64)> = conn.hgetall(&(key.clone() + ":hash")).await.unwrap();
1452        fields.sort();
1453        assert_eq!(
1454            fields,
1455            vec![(":2:a".to_owned(), 1), (":2:b".to_owned(), 1),]
1456        );
1457
1458        let ttl: u64 = conn.ttl(&key).await.unwrap();
1459        assert!(ttl >= 59, "{ttl}");
1460        assert!(ttl <= 60, "{ttl}");
1461    }
1462
1463    /// A quota with a maximum cardinality rejects items which would create a new bucket beyond
1464    /// that cardinality, but keeps serving the buckets which already exist.
1465    #[tokio::test]
1466    async fn test_is_rate_limited_script_dimensions_cardinality() {
1467        let now = SystemTime::now()
1468            .duration_since(UNIX_EPOCH)
1469            .map(|duration| duration.as_secs())
1470            .unwrap();
1471
1472        let rate_limiter = build_rate_limiter();
1473        let mut conn = rate_limiter.client.get_connection().await.unwrap();
1474
1475        let key = format!("cardinality___{now}");
1476        let refund_key = format!("r:cardinality___{now}");
1477
1478        let script = RedisScripts::load_is_rate_limited();
1479
1480        // A generous limit, so that only the cardinality can cause a rejection.
1481        let invoke_async = async |dims: &str, conn: &mut _| {
1482            let mut invocation = script.prepare_invoke();
1483            invocation
1484                .key(&key) // key
1485                .key(&refund_key) // refund key
1486                .arg(100) // limit
1487                .arg(now + 60) // expiry
1488                .arg(1) // quantity
1489                .arg(false) // over accept once
1490                .arg(dims) // dimensions
1491                .arg(2); // max cardinality
1492
1493            invocation
1494                .invoke_async::<ScriptResult>(conn)
1495                .await
1496                .unwrap()
1497                .0
1498        };
1499
1500        // The first two buckets fit within the cardinality of 2.
1501        assert!(!invoke_async(":2:a", &mut conn).await[0].is_rejected);
1502        assert!(!invoke_async(":2:b", &mut conn).await[0].is_rejected);
1503
1504        // A third, new bucket is rejected, even though its own limit is nowhere near reached.
1505        let result = invoke_async(":2:c", &mut conn).await;
1506        assert!(result[0].is_rejected);
1507        assert_eq!(result[0].consumed, 0);
1508
1509        // The existing buckets keep being counted, as they do not add cardinality.
1510        assert!(!invoke_async(":2:a", &mut conn).await[0].is_rejected);
1511        assert!(!invoke_async(":2:b", &mut conn).await[0].is_rejected);
1512
1513        // The rejected bucket was never created.
1514        let mut fields: Vec<(String, i64)> = conn.hgetall(&key).await.unwrap();
1515        fields.sort();
1516        assert_eq!(fields, vec![(":2:a".to_owned(), 2), (":2:b".to_owned(), 2)]);
1517    }
1518
1519    /// A cardinality of `0` rejects every item, and a cardinality of `-1` is unlimited.
1520    #[tokio::test]
1521    async fn test_is_rate_limited_script_dimensions_cardinality_edges() {
1522        let now = SystemTime::now()
1523            .duration_since(UNIX_EPOCH)
1524            .map(|duration| duration.as_secs())
1525            .unwrap();
1526
1527        let rate_limiter = build_rate_limiter();
1528        let mut conn = rate_limiter.client.get_connection().await.unwrap();
1529
1530        let script = RedisScripts::load_is_rate_limited();
1531
1532        let invoke_async = async |key: &str, dims: &str, cardinality: i32, conn: &mut _| {
1533            let mut invocation = script.prepare_invoke();
1534            invocation
1535                .key(key) // key
1536                .key(get_refunded_quota_key(key)) // refund key
1537                .arg(100) // limit
1538                .arg(now + 60) // expiry
1539                .arg(1) // quantity
1540                .arg(false) // over accept once
1541                .arg(dims) // dimensions
1542                .arg(cardinality); // max cardinality
1543
1544            invocation
1545                .invoke_async::<ScriptResult>(conn)
1546                .await
1547                .unwrap()
1548                .0
1549        };
1550
1551        let zero = format!("cardinality_zero___{now}");
1552        assert!(invoke_async(&zero, ":2:a", 0, &mut conn).await[0].is_rejected);
1553        assert_eq!(conn.hlen::<_, u64>(&zero).await.unwrap(), 0);
1554
1555        // `-1` means unlimited, so an arbitrary number of new buckets is accepted.
1556        let unlimited = format!("cardinality_unlimited___{now}");
1557        for i in 0..10 {
1558            let dims = format!(":2:{i}");
1559            assert!(!invoke_async(&unlimited, &dims, -1, &mut conn).await[0].is_rejected);
1560        }
1561        assert_eq!(conn.hlen::<_, u64>(&unlimited).await.unwrap(), 10);
1562    }
1563
1564    /// The refund counter is a plain key, and hence shared by all dimension buckets of a quota.
1565    #[tokio::test]
1566    async fn test_is_rate_limited_script_dimensions_share_refund_key() {
1567        let now = SystemTime::now()
1568            .duration_since(UNIX_EPOCH)
1569            .map(|duration| duration.as_secs())
1570            .unwrap();
1571
1572        let rate_limiter = build_rate_limiter();
1573        let mut conn = rate_limiter.client.get_connection().await.unwrap();
1574
1575        let key = format!("dims_refund___{now}");
1576        let refund_key = get_refunded_quota_key(&key);
1577        let () = conn.set(&refund_key, 5).await.unwrap();
1578
1579        let script = RedisScripts::load_is_rate_limited();
1580
1581        let invoke_async = async |dims: &str, conn: &mut _| {
1582            let mut invocation = script.prepare_invoke();
1583            invocation
1584                .key(&key) // key
1585                .key(&refund_key) // refund key
1586                .arg(1) // limit
1587                .arg(now + 60) // expiry
1588                .arg(1) // quantity
1589                .arg(false) // over accept once
1590                .arg(dims) // dimensions
1591                .arg(999); // max cardinality
1592
1593            invocation
1594                .invoke_async::<ScriptResult>(conn)
1595                .await
1596                .unwrap()
1597                .0
1598        };
1599
1600        // Both buckets are discounted by the full refund, since the refund is not dimensioned.
1601        let result = invoke_async(":2:a", &mut conn).await;
1602        assert!(!result[0].is_rejected);
1603        assert_eq!(result[0].consumed, -4);
1604
1605        let result = invoke_async(":2:b", &mut conn).await;
1606        assert!(!result[0].is_rejected);
1607        assert_eq!(result[0].consumed, -4);
1608    }
1609
1610    /// Usual rate limiting with a cache should just work as expected.
1611    #[tokio::test]
1612    async fn test_quota_with_cache() {
1613        let quotas = &[Quota {
1614            id: Some(format!("test_simple_quota_{}", uuid::Uuid::new_v4()).into()),
1615            categories: DataCategories::new(),
1616            scope: QuotaScope::Organization,
1617            scope_id: None,
1618            limit: Some(50),
1619            window: Some(60),
1620            reason_code: Some(ReasonCode::new("get_lost")),
1621            namespace: None,
1622            group_by: None,
1623        }];
1624
1625        let scoping = ItemScoping {
1626            category: DataCategory::Error,
1627            scoping: Scoping {
1628                organization_id: OrganizationId::new(42),
1629                project_id: ProjectId::new(43),
1630                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1631                key_id: Some(44),
1632            },
1633            namespace: MetricNamespaceScoping::None,
1634            dimensions: DimensionMap::default(),
1635        };
1636
1637        // For this test, with only a single rate limiter accessing Redis and always a quantity of
1638        // `1`, the parameters for the cache should not make any difference.
1639        let rate_limiter = build_rate_limiter().cache(Some(0.1), Some(0.9));
1640
1641        for _ in 0..50 {
1642            let rate_limits = rate_limiter
1643                .is_rate_limited(quotas, &scoping, 1, false)
1644                .await
1645                .unwrap();
1646
1647            assert!(rate_limits.is_empty());
1648        }
1649
1650        let rate_limits: Vec<RateLimit> = rate_limiter
1651            .is_rate_limited(quotas, &scoping, 1, false)
1652            .await
1653            .expect("rate limiting failed")
1654            .into_iter()
1655            .collect();
1656
1657        assert_eq!(
1658            rate_limits,
1659            vec![RateLimit {
1660                categories: DataCategories::new(),
1661                scope: RateLimitScope::Organization(OrganizationId::new(42)),
1662                reason_code: Some(ReasonCode::new("get_lost")),
1663                retry_after: rate_limits[0].retry_after,
1664                namespaces: smallvec![],
1665                dimensions: DimensionMap::default(),
1666            }]
1667        );
1668    }
1669
1670    #[tokio::test]
1671    async fn test_quota_with_cache_slightly_over_account() {
1672        let window = 60;
1673        let limit = 50 * window;
1674
1675        let quotas = &[Quota {
1676            id: Some(format!("test_simple_quota_{}", uuid::Uuid::new_v4()).into()),
1677            categories: DataCategories::new(),
1678            scope: QuotaScope::Organization,
1679            scope_id: None,
1680            limit: Some(limit),
1681            window: Some(window),
1682            reason_code: Some(ReasonCode::new("get_lost")),
1683            namespace: None,
1684            group_by: None,
1685        }];
1686
1687        let scoping = ItemScoping {
1688            category: DataCategory::Error,
1689            scoping: Scoping {
1690                organization_id: OrganizationId::new(42),
1691                project_id: ProjectId::new(43),
1692                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1693                key_id: Some(44),
1694            },
1695            namespace: MetricNamespaceScoping::None,
1696            dimensions: DimensionMap::default(),
1697        };
1698
1699        // 10% Quota cache.
1700        let rate_limiter1 = build_rate_limiter().cache(Some(0.1), None);
1701        let rate_limiter2 = build_rate_limiter().cache(Some(0.1), None);
1702
1703        // Prime the cache.
1704        let rate_limits = rate_limiter1
1705            .is_rate_limited(quotas, &scoping, 1, false)
1706            .await
1707            .unwrap();
1708        assert!(rate_limits.is_empty());
1709        // Reserve 3 out 5 in the cache.
1710        let rate_limits = rate_limiter1
1711            .is_rate_limited(quotas, &scoping, 3, false)
1712            .await
1713            .unwrap();
1714        assert!(rate_limits.is_empty());
1715
1716        // Consume right up to the limit on the other limiter
1717        let rate_limits = rate_limiter2
1718            .is_rate_limited(quotas, &scoping, limit as usize - 1, false)
1719            .await
1720            .unwrap();
1721        assert!(rate_limits.is_empty());
1722
1723        // There is still one more slot in the cache.
1724        let rate_limits = rate_limiter1
1725            .is_rate_limited(quotas, &scoping, 1, false)
1726            .await
1727            .unwrap();
1728        assert!(rate_limits.is_empty());
1729
1730        // This should now rate limit, as the cache is exhausted and Redis is checked.
1731        let rate_limits: Vec<RateLimit> = rate_limiter1
1732            .is_rate_limited(quotas, &scoping, 1, false)
1733            .await
1734            .unwrap()
1735            .into_iter()
1736            .collect();
1737
1738        assert_eq!(
1739            rate_limits,
1740            vec![RateLimit {
1741                categories: DataCategories::new(),
1742                scope: RateLimitScope::Organization(OrganizationId::new(42)),
1743                reason_code: Some(ReasonCode::new("get_lost")),
1744                retry_after: rate_limits[0].retry_after,
1745                namespaces: smallvec![],
1746                dimensions: DimensionMap::default(),
1747            }]
1748        );
1749    }
1750
1751    #[tokio::test]
1752    async fn test_quota_dimensions() {
1753        let quotas = &[
1754            Quota {
1755                id: Some(format!("test_quota_go_over{}", uuid::Uuid::new_v4()).into()),
1756                categories: DataCategories::new().add(DataCategory::Monitor).unwrap(),
1757                scope: QuotaScope::Project,
1758                scope_id: None,
1759                limit: Some(2),
1760                window: Some(60),
1761                reason_code: Some(ReasonCode::new("get_lost")),
1762                namespace: None,
1763                group_by: GroupBy {
1764                    max_cardinality: 999,
1765                    dimensions: BTreeSet::from([
1766                        Dimension::CheckInEnvironment,
1767                        Dimension::CheckInSlug,
1768                    ])
1769                    .into(),
1770                }
1771                .into(),
1772            },
1773            Quota {
1774                id: Some(format!("test_quota_wont_go_over{}", uuid::Uuid::new_v4()).into()),
1775                categories: DataCategories::new().add(DataCategory::Monitor).unwrap(),
1776                scope: QuotaScope::Project,
1777                scope_id: None,
1778                limit: Some(2),
1779                window: Some(60),
1780                reason_code: Some(ReasonCode::new("get_lost")),
1781                namespace: None,
1782                group_by: None,
1783            },
1784        ];
1785
1786        let scoping = ItemScoping {
1787            category: DataCategory::Monitor,
1788            scoping: Scoping {
1789                organization_id: OrganizationId::new(42),
1790                project_id: ProjectId::new(43),
1791                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1792                key_id: Some(44),
1793            },
1794            namespace: MetricNamespaceScoping::None,
1795            dimensions: DimensionMap::from([
1796                (Dimension::CheckInEnvironment, "us".to_owned()),
1797                (Dimension::CheckInSlug, "cron1".to_owned()),
1798            ]),
1799        };
1800
1801        let no_dims_scoping = ItemScoping {
1802            category: DataCategory::Monitor,
1803            scoping: Scoping {
1804                organization_id: OrganizationId::new(42),
1805                project_id: ProjectId::new(43),
1806                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1807                key_id: Some(44),
1808            },
1809            namespace: MetricNamespaceScoping::None,
1810            dimensions: DimensionMap::default(),
1811        };
1812
1813        let rate_limiter = build_rate_limiter();
1814
1815        // limit is 2, so first call not rate limited
1816        let is_limited = rate_limiter
1817            .is_rate_limited(quotas, &scoping, 1, true)
1818            .await
1819            .unwrap()
1820            .is_limited();
1821        assert!(!is_limited);
1822
1823        let is_limited = rate_limiter
1824            .is_rate_limited(quotas, &scoping, 2, false)
1825            .await
1826            .unwrap()
1827            .is_limited();
1828        assert!(is_limited);
1829
1830        // make sure the non-dimensioned monitor is not limited...
1831        let is_limited = rate_limiter
1832            .is_rate_limited(quotas, &no_dims_scoping, 1, true)
1833            .await
1834            .unwrap()
1835            .is_limited();
1836        assert!(!is_limited);
1837
1838        // ...until we invoke again
1839        let is_limited = rate_limiter
1840            .is_rate_limited(quotas, &no_dims_scoping, 2, false)
1841            .await
1842            .unwrap()
1843            .is_limited();
1844        assert!(is_limited);
1845    }
1846
1847    /// Builds a project-scoped monitor quota with the passed dimensions and limit.
1848    fn monitor_quota(limit: u64, group_by: Option<GroupBy>) -> Quota {
1849        Quota {
1850            id: Some(format!("test_dimensions_{}", uuid::Uuid::new_v4()).into()),
1851            categories: DataCategories::new().add(DataCategory::Monitor).unwrap(),
1852            scope: QuotaScope::Project,
1853            scope_id: None,
1854            limit: Some(limit),
1855            window: Some(60),
1856            reason_code: Some(ReasonCode::new("get_lost")),
1857            namespace: None,
1858            group_by,
1859        }
1860    }
1861
1862    /// Builds a monitor item scoping with the passed dimensions.
1863    fn monitor_scoping(dimensions: &[(Dimension, &str)]) -> ItemScoping {
1864        let mut owned_dims: Vec<_> = dimensions
1865            .iter()
1866            .map(|(d, v)| (*d, (*v).to_owned()))
1867            .collect();
1868        ItemScoping {
1869            category: DataCategory::Monitor,
1870            scoping: Scoping {
1871                organization_id: OrganizationId::new(42),
1872                project_id: ProjectId::new(43),
1873                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1874                key_id: Some(44),
1875            },
1876            namespace: MetricNamespaceScoping::None,
1877            dimensions: DimensionMap::from_iter(owned_dims.drain(..)),
1878        }
1879    }
1880
1881    /// Items which differ in any of the quota's dimensions consume separate buckets.
1882    #[tokio::test]
1883    async fn test_quota_dimensions_are_independent_buckets() {
1884        let quotas = &[monitor_quota(
1885            1,
1886            GroupBy {
1887                max_cardinality: 999,
1888                dimensions: BTreeSet::from([Dimension::CheckInEnvironment, Dimension::CheckInSlug])
1889                    .into(),
1890            }
1891            .into(),
1892        )];
1893
1894        let rate_limiter = build_rate_limiter();
1895
1896        // Every one of these differs from the others in at least one dimension, so each one
1897        // gets the full limit of 1 for itself.
1898        let distinct = [
1899            [
1900                (Dimension::CheckInEnvironment, "prod"),
1901                (Dimension::CheckInSlug, "cron1"),
1902            ],
1903            [
1904                (Dimension::CheckInEnvironment, "prod"),
1905                (Dimension::CheckInSlug, "cron2"),
1906            ],
1907            [
1908                (Dimension::CheckInEnvironment, "dev"),
1909                (Dimension::CheckInSlug, "cron1"),
1910            ],
1911        ];
1912
1913        for dims in &distinct {
1914            let scoping = monitor_scoping(dims);
1915
1916            let rate_limits = rate_limiter
1917                .is_rate_limited(quotas, &scoping, 1, false)
1918                .await
1919                .unwrap();
1920            assert!(!rate_limits.is_limited(), "first check for {dims:?}");
1921
1922            // The bucket for these dimensions is now exhausted.
1923            let rate_limits: Vec<RateLimit> = rate_limiter
1924                .is_rate_limited(quotas, &scoping, 1, false)
1925                .await
1926                .unwrap()
1927                .into_iter()
1928                .collect();
1929
1930            assert_eq!(
1931                rate_limits,
1932                vec![RateLimit {
1933                    categories: DataCategories::new().add(DataCategory::Monitor).unwrap(),
1934                    scope: RateLimitScope::Project(ProjectId::new(43)),
1935                    reason_code: Some(ReasonCode::new("get_lost")),
1936                    retry_after: rate_limits[0].retry_after,
1937                    namespaces: smallvec![],
1938                    dimensions: DimensionMap::from(*dims),
1939                }],
1940                "second check for {dims:?}"
1941            );
1942        }
1943    }
1944
1945    /// Dimensions the quota does not key on do not split the bucket.
1946    #[tokio::test]
1947    async fn test_quota_dimensions_ignores_unrelated_dimensions() {
1948        let quotas = &[monitor_quota(
1949            2,
1950            GroupBy {
1951                max_cardinality: 999,
1952                dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
1953            }
1954            .into(),
1955        )];
1956
1957        let rate_limiter = build_rate_limiter();
1958
1959        // The quota only keys on the slug, so these two items share a bucket despite their
1960        // differing environments.
1961        let prod = monitor_scoping(&[
1962            (Dimension::CheckInEnvironment, "prod"),
1963            (Dimension::CheckInSlug, "cron1"),
1964        ]);
1965        let dev = monitor_scoping(&[
1966            (Dimension::CheckInEnvironment, "dev"),
1967            (Dimension::CheckInSlug, "cron1"),
1968        ]);
1969
1970        assert!(
1971            !rate_limiter
1972                .is_rate_limited(quotas, &prod, 1, false)
1973                .await
1974                .unwrap()
1975                .is_limited()
1976        );
1977        assert!(
1978            !rate_limiter
1979                .is_rate_limited(quotas, &dev, 1, false)
1980                .await
1981                .unwrap()
1982                .is_limited()
1983        );
1984
1985        // Two of the limit of 2 are consumed, so both items are limited from here on.
1986        assert!(
1987            rate_limiter
1988                .is_rate_limited(quotas, &prod, 1, false)
1989                .await
1990                .unwrap()
1991                .is_limited()
1992        );
1993        assert!(
1994            rate_limiter
1995                .is_rate_limited(quotas, &dev, 1, false)
1996                .await
1997                .unwrap()
1998                .is_limited()
1999        );
2000    }
2001
2002    /// A quota with dimensions never applies to an item which does not carry them.
2003    #[tokio::test]
2004    async fn test_quota_dimensions_skipped_for_undimensioned_item() {
2005        let quotas = &[monitor_quota(
2006            1,
2007            GroupBy {
2008                max_cardinality: 999,
2009                dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
2010            }
2011            .into(),
2012        )];
2013
2014        let rate_limiter = build_rate_limiter();
2015
2016        // The quota does not match, so no amount of traffic can exhaust it.
2017        for _ in 0..5 {
2018            assert!(
2019                !rate_limiter
2020                    .is_rate_limited(quotas, &monitor_scoping(&[]), 1, false)
2021                    .await
2022                    .unwrap()
2023                    .is_limited()
2024            );
2025        }
2026
2027        // An item missing only one of two required dimensions is skipped the same way.
2028        let quotas = &[monitor_quota(
2029            1,
2030            GroupBy {
2031                max_cardinality: 999,
2032                dimensions: BTreeSet::from([Dimension::CheckInEnvironment, Dimension::CheckInSlug])
2033                    .into(),
2034            }
2035            .into(),
2036        )];
2037
2038        for _ in 0..5 {
2039            assert!(
2040                !rate_limiter
2041                    .is_rate_limited(
2042                        quotas,
2043                        &monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]),
2044                        1,
2045                        false
2046                    )
2047                    .await
2048                    .unwrap()
2049                    .is_limited()
2050            );
2051        }
2052    }
2053
2054    /// Once the quota's cardinality is used up, items with new dimension values are limited,
2055    /// while the existing ones keep their quota.
2056    #[tokio::test]
2057    async fn test_quota_dimensions_max_cardinality() {
2058        let quotas = &[monitor_quota(
2059            100,
2060            GroupBy {
2061                max_cardinality: 2,
2062                dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
2063            }
2064            .into(),
2065        )];
2066
2067        let rate_limiter = build_rate_limiter();
2068
2069        let cron1 = monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]);
2070        let cron2 = monitor_scoping(&[(Dimension::CheckInSlug, "cron2")]);
2071        let cron3 = monitor_scoping(&[(Dimension::CheckInSlug, "cron3")]);
2072
2073        // Two distinct slugs fit into the cardinality of 2.
2074        assert!(
2075            !rate_limiter
2076                .is_rate_limited(quotas, &cron1, 1, false)
2077                .await
2078                .unwrap()
2079                .is_limited()
2080        );
2081        assert!(
2082            !rate_limiter
2083                .is_rate_limited(quotas, &cron2, 1, false)
2084                .await
2085                .unwrap()
2086                .is_limited()
2087        );
2088
2089        // The third one is limited, despite the limit of 100 being nowhere near reached.
2090        let rate_limits: Vec<RateLimit> = rate_limiter
2091            .is_rate_limited(quotas, &cron3, 1, false)
2092            .await
2093            .unwrap()
2094            .into_iter()
2095            .collect();
2096
2097        assert_eq!(
2098            rate_limits,
2099            vec![RateLimit {
2100                categories: DataCategories::new().add(DataCategory::Monitor).unwrap(),
2101                scope: RateLimitScope::Project(ProjectId::new(43)),
2102                reason_code: Some(ReasonCode::new("get_lost")),
2103                retry_after: rate_limits[0].retry_after,
2104                namespaces: smallvec![],
2105                dimensions: DimensionMap::from([(Dimension::CheckInSlug, "cron3")]),
2106            }]
2107        );
2108
2109        // The slugs which made it in are unaffected.
2110        assert!(
2111            !rate_limiter
2112                .is_rate_limited(quotas, &cron1, 1, false)
2113                .await
2114                .unwrap()
2115                .is_limited()
2116        );
2117        assert!(
2118            !rate_limiter
2119                .is_rate_limited(quotas, &cron2, 1, false)
2120                .await
2121                .unwrap()
2122                .is_limited()
2123        );
2124    }
2125
2126    /// A quota with a cardinality of `0` rejects every dimensioned item.
2127    #[tokio::test]
2128    async fn test_quota_dimensions_zero_cardinality() {
2129        let quotas = &[monitor_quota(
2130            100,
2131            GroupBy {
2132                max_cardinality: 0,
2133                dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
2134            }
2135            .into(),
2136        )];
2137
2138        let rate_limiter = build_rate_limiter();
2139
2140        assert!(
2141            rate_limiter
2142                .is_rate_limited(
2143                    quotas,
2144                    &monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]),
2145                    1,
2146                    false
2147                )
2148                .await
2149                .unwrap()
2150                .is_limited()
2151        );
2152    }
2153
2154    /// Every dimension bucket gets its own limit, even when the opportunistic cache is enabled.
2155    ///
2156    /// The cache is keyed on the dimension values of an item, so every bucket of a quota has its
2157    /// own cache entry. This test pins down that accepting from the cache does not make the
2158    /// buckets share their limit.
2159    #[tokio::test]
2160    async fn test_quota_dimensions_with_cache() {
2161        let limit = 50 * 60;
2162
2163        let quotas = &[monitor_quota(
2164            limit,
2165            GroupBy {
2166                max_cardinality: 999,
2167                dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
2168            }
2169            .into(),
2170        )];
2171
2172        let rate_limiter = build_rate_limiter().cache(Some(0.1), Some(0.9));
2173
2174        let crons = [
2175            monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]),
2176            monitor_scoping(&[(Dimension::CheckInSlug, "cron2")]),
2177        ];
2178
2179        // The loop below takes long enough to cross into the next window, which would start
2180        // fresh buckets and over-accept. Check all items at the same timestamp instead.
2181        let timestamp = UnixTimestamp::now();
2182
2183        // Interleave the two slugs, so that both of them go through the cache, and send 50% more
2184        // than either of them is allowed.
2185        let mut accepted = [0u64; 2];
2186        for _ in 0..(limit + limit / 2) {
2187            for (scoping, accepted) in crons.iter().zip(accepted.iter_mut()) {
2188                if !rate_limiter
2189                    .is_rate_limited_at(quotas, scoping, 1, false, timestamp)
2190                    .await
2191                    .unwrap()
2192                    .is_limited()
2193                {
2194                    *accepted += 1;
2195                }
2196            }
2197        }
2198
2199        // The cache accepts opportunistically, so allow for the same 1% of over-acceptance the
2200        // undimensioned cache tests allow for.
2201        let tolerance = limit / 100;
2202        for accepted in accepted {
2203            assert!(
2204                accepted.abs_diff(limit) <= tolerance,
2205                "accepted {accepted} of a limit of {limit}"
2206            );
2207        }
2208    }
2209
2210    /// Dimensioned and undimensioned quotas on the same item are counted separately, and either
2211    /// of them can rate limit the item on its own.
2212    #[tokio::test]
2213    async fn test_quota_dimensions_alongside_undimensioned_quota() {
2214        // The undimensioned quota is the more generous of the two.
2215        let quotas = &[
2216            monitor_quota(
2217                1,
2218                GroupBy {
2219                    max_cardinality: 999,
2220                    dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
2221                }
2222                .into(),
2223            ),
2224            monitor_quota(3, None),
2225        ];
2226
2227        let rate_limiter = build_rate_limiter();
2228
2229        let cron1 = monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]);
2230        let cron2 = monitor_scoping(&[(Dimension::CheckInSlug, "cron2")]);
2231        let cron3 = monitor_scoping(&[(Dimension::CheckInSlug, "cron3")]);
2232        let cron4 = monitor_scoping(&[(Dimension::CheckInSlug, "cron4")]);
2233
2234        // Each new slug has its own per-dimension bucket, and all of them share the
2235        // undimensioned one.
2236        for scoping in [&cron1, &cron2, &cron3] {
2237            assert!(
2238                !rate_limiter
2239                    .is_rate_limited(quotas, scoping, 1, false)
2240                    .await
2241                    .unwrap()
2242                    .is_limited()
2243            );
2244        }
2245
2246        // The per-dimension bucket of a repeated slug is exhausted.
2247        assert!(
2248            rate_limiter
2249                .is_rate_limited(quotas, &cron1, 1, false)
2250                .await
2251                .unwrap()
2252                .is_limited()
2253        );
2254
2255        // A fresh slug has quota of its own, but the shared, undimensioned quota of 3 is now
2256        // used up, so it is limited all the same.
2257        assert!(
2258            rate_limiter
2259                .is_rate_limited(quotas, &cron4, 1, false)
2260                .await
2261                .unwrap()
2262                .is_limited()
2263        );
2264    }
2265
2266    /// Regression test for the project-wide over-application of a per-monitor quota.
2267    ///
2268    /// Exhausting one monitor's bucket must not produce anything which can be cached on the
2269    /// project, because a cached limit would then drop check-ins for every other monitor and be
2270    /// propagated to SDKs and downstream Relays as a scope-wide monitor limit.
2271    #[tokio::test]
2272    async fn test_quota_dimensions_are_not_cacheable() {
2273        let quotas = &[monitor_quota(
2274            1,
2275            GroupBy {
2276                max_cardinality: 999,
2277                dimensions: BTreeSet::from([Dimension::CheckInSlug]).into(),
2278            }
2279            .into(),
2280        )];
2281
2282        let rate_limiter = build_rate_limiter();
2283
2284        let cron1 = monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]);
2285        let cron2 = monitor_scoping(&[(Dimension::CheckInSlug, "cron2")]);
2286
2287        // Exhaust the bucket of a single monitor.
2288        rate_limiter
2289            .is_rate_limited(quotas, &cron1, 1, false)
2290            .await
2291            .unwrap();
2292        let limits = rate_limiter
2293            .is_rate_limited(quotas, &cron1, 1, false)
2294            .await
2295            .unwrap();
2296
2297        // The limit is enforced against the check-in which produced it.
2298        assert!(limits.is_limited());
2299
2300        assert!(!limits.check(&cron2).is_limited());
2301
2302        // The other monitor still has its own bucket in Redis.
2303        assert!(
2304            !rate_limiter
2305                .is_rate_limited(quotas, &cron2, 1, false)
2306                .await
2307                .unwrap()
2308                .is_limited()
2309        );
2310    }
2311
2312    /// An undimensioned quota on the same category stays propagatable.
2313    #[tokio::test]
2314    async fn test_quota_without_dimensions_is_propagatable() {
2315        let quotas = &[monitor_quota(1, None)];
2316
2317        let rate_limiter = build_rate_limiter();
2318
2319        let cron1 = monitor_scoping(&[(Dimension::CheckInSlug, "cron1")]);
2320
2321        rate_limiter
2322            .is_rate_limited(quotas, &cron1, 1, false)
2323            .await
2324            .unwrap();
2325        let limits = rate_limiter
2326            .is_rate_limited(quotas, &cron1, 1, false)
2327            .await
2328            .unwrap();
2329
2330        // This limit really does apply to the whole project, so it is kept for the cache even
2331        // though the item which produced it carried dimensions.
2332        assert!(limits.is_limited());
2333        assert!(!limits.propagatable().is_empty());
2334    }
2335}