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
19const GRACE: u64 = 60;
23
24#[derive(Debug, Error)]
26#[error("failed to communicate with redis")]
27pub struct RateLimitingError(
28 #[from]
29 #[source]
30 pub RedisError,
31);
32
33fn get_refunded_quota_key(counter_key: &str) -> String {
38 format!("r:{counter_key}")
39}
40
41struct 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#[derive(Debug, Clone)]
58pub struct OwnedRedisQuota {
59 quota: Quota,
61 scoping: ItemScoping,
63 prefix: Arc<str>,
65 window: u64,
67 quantity: u64,
69 timestamp: UnixTimestamp,
71}
72
73impl OwnedRedisQuota {
74 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#[derive(Debug, Clone, Eq, PartialEq)]
89pub struct RedisQuota<'a> {
90 quota: &'a Quota,
92 scoping: &'a ItemScoping,
94 prefix: Arc<str>,
96 window: u64,
98 quantity: u64,
100 timestamp: UnixTimestamp,
102}
103
104impl<'a> RedisQuota<'a> {
105 pub fn new(
111 quota: &'a Quota,
112 quantity: u64,
113 scoping: &'a ItemScoping,
114 timestamp: UnixTimestamp,
115 ) -> Option<Self> {
116 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 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 pub fn window(&self) -> u64 {
145 self.window
146 }
147
148 pub fn prefix(&self) -> &str {
150 &self.prefix
151 }
152
153 pub fn quantity(&self) -> u64 {
155 self.quantity
156 }
157
158 pub fn limit(&self) -> i64 {
163 self.limit
164 .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 pub fn slot(&self) -> u64 {
177 (self.timestamp.as_secs() - self.shift()) / self.window
178 }
179
180 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 pub fn key_expiry(&self) -> u64 {
191 self.expiry().as_secs() + GRACE
192 }
193
194 pub fn key(&self) -> QuotaCacheKey {
200 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 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 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#[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#[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 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 pub fn max_limit(mut self, max_limit: Option<u64>) -> Self {
324 self.max_limit = max_limit;
325 self
326 }
327
328 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 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 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 } else if quota.limit == Some(0) {
391 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 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 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 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 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 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 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#[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#[derive(Debug)]
533struct QuotaState {
534 is_rejected: bool,
536 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 #[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 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 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 assert!(
787 !rate_limiter
788 .is_rate_limited(quotas, &scoping, 1, false)
789 .await
790 .unwrap()
791 .is_limited()
792 );
793
794 assert!(
796 rate_limiter
797 .is_rate_limited(quotas, &scoping, 1, false)
798 .await
799 .unwrap()
800 .is_limited()
801 );
802
803 assert!(
805 rate_limiter
806 .is_rate_limited(quotas, &scoping, 0, false)
807 .await
808 .unwrap()
809 .is_limited()
810 );
811
812 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 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 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 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 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("a, 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("a, 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), 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("a, 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 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(&r_foo) .key(&bar) .key(&r_bar) .arg(1) .arg(now + 60) .arg(1) .arg(false) .arg(EMPTY_DIMENSIONS) .arg(1) .arg(2) .arg(now + 120) .arg(1) .arg(false) .arg(EMPTY_DIMENSIONS) .arg(1); let mut invocation2 = script.prepare_invoke();
1181 invocation2
1182 .key(&bar) .key(&r_bar) .arg(2) .arg(now + 120) .arg(1) .arg(false) .arg(EMPTY_DIMENSIONS) .arg(1); 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 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 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 assert_invocation!(invocation2, @r"
1246 ScriptResult(
1247 [
1248 QuotaState {
1249 is_rejected: false,
1250 consumed: 2,
1251 },
1252 ],
1253 )
1254 "
1255 );
1256
1257 assert_invocation!(invocation2, @r"
1259 ScriptResult(
1260 [
1261 QuotaState {
1262 is_rejected: true,
1263 consumed: 2,
1264 },
1265 ],
1266 )
1267 "
1268 );
1269
1270 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 let () = conn.get(r_foo).await.unwrap();
1301 let () = conn.get(r_bar).await.unwrap();
1302
1303 let () = conn.set(&apple, 5).await.unwrap();
1305
1306 let mut invocation = script.prepare_invoke();
1307 invocation
1308 .key(&orange) .key(&baz) .arg(1) .arg(now + 60) .arg(1) .arg(false) .arg(EMPTY_DIMENSIONS) .arg(1); assert_invocation!(invocation, @r"
1319 ScriptResult(
1320 [
1321 QuotaState {
1322 is_rejected: false,
1323 consumed: 1,
1324 },
1325 ],
1326 )
1327 "
1328 );
1329
1330 assert_invocation!(invocation, @r"
1332 ScriptResult(
1333 [
1334 QuotaState {
1335 is_rejected: true,
1336 consumed: 1,
1337 },
1338 ],
1339 )
1340 "
1341 );
1342
1343 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(&apple) .arg(1) .arg(now + 60) .arg(1) .arg(false) .arg(EMPTY_DIMENSIONS) .arg(1); assert_invocation!(invocation, @r"
1369 ScriptResult(
1370 [
1371 QuotaState {
1372 is_rejected: false,
1373 consumed: -3,
1374 },
1375 ],
1376 )
1377 "
1378 );
1379 }
1380
1381 #[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 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(&refund_key) .arg(1) .arg(now + 60) .arg(1) .arg(false) .arg(dims) .arg(999); 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 let result = invoke_async(":2:a", &mut conn).await;
1428 assert!(!result[0].is_rejected);
1429 assert_eq!(result[0].consumed, 1);
1430
1431 let result = invoke_async(":2:a", &mut conn).await;
1433 assert!(result[0].is_rejected);
1434 assert_eq!(result[0].consumed, 1);
1435
1436 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 let result = invoke_async(EMPTY_DIMENSIONS, &mut conn).await;
1447 assert!(!result[0].is_rejected);
1448 assert_eq!(result[0].consumed, 1);
1449
1450 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 #[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 let invoke_async = async |dims: &str, conn: &mut _| {
1482 let mut invocation = script.prepare_invoke();
1483 invocation
1484 .key(&key) .key(&refund_key) .arg(100) .arg(now + 60) .arg(1) .arg(false) .arg(dims) .arg(2); invocation
1494 .invoke_async::<ScriptResult>(conn)
1495 .await
1496 .unwrap()
1497 .0
1498 };
1499
1500 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 let result = invoke_async(":2:c", &mut conn).await;
1506 assert!(result[0].is_rejected);
1507 assert_eq!(result[0].consumed, 0);
1508
1509 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 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 #[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(get_refunded_quota_key(key)) .arg(100) .arg(now + 60) .arg(1) .arg(false) .arg(dims) .arg(cardinality); 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 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 #[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(&refund_key) .arg(1) .arg(now + 60) .arg(1) .arg(false) .arg(dims) .arg(999); invocation
1594 .invoke_async::<ScriptResult>(conn)
1595 .await
1596 .unwrap()
1597 .0
1598 };
1599
1600 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 #[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 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 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 let rate_limits = rate_limiter1
1705 .is_rate_limited(quotas, &scoping, 1, false)
1706 .await
1707 .unwrap();
1708 assert!(rate_limits.is_empty());
1709 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 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 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 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 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 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 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 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 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 #[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 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 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 #[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 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 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 #[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 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 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 #[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 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 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 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 #[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 #[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 let timestamp = UnixTimestamp::now();
2182
2183 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 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 #[tokio::test]
2213 async fn test_quota_dimensions_alongside_undimensioned_quota() {
2214 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 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 assert!(
2248 rate_limiter
2249 .is_rate_limited(quotas, &cron1, 1, false)
2250 .await
2251 .unwrap()
2252 .is_limited()
2253 );
2254
2255 assert!(
2258 rate_limiter
2259 .is_rate_limited(quotas, &cron4, 1, false)
2260 .await
2261 .unwrap()
2262 .is_limited()
2263 );
2264 }
2265
2266 #[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 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 assert!(limits.is_limited());
2299
2300 assert!(!limits.check(&cron2).is_limited());
2301
2302 assert!(
2304 !rate_limiter
2305 .is_rate_limited(quotas, &cron2, 1, false)
2306 .await
2307 .unwrap()
2308 .is_limited()
2309 );
2310 }
2311
2312 #[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 assert!(limits.is_limited());
2333 assert!(!limits.propagatable().is_empty());
2334 }
2335}