relay_quotas/
quota.rs

1use std::fmt;
2use std::str::FromStr;
3use std::sync::Arc;
4
5use relay_base_schema::metrics::MetricNamespace;
6use relay_base_schema::organization::OrganizationId;
7use relay_base_schema::project::{ProjectId, ProjectKey};
8use serde::ser::SerializeSeq;
9use serde::{Deserialize, Serialize};
10use smallvec::SmallVec;
11
12#[doc(inline)]
13pub use relay_base_schema::data_category::{CategoryUnit, DataCategory};
14
15/// Data scoping information for rate limiting and quota enforcement.
16///
17/// [`Scoping`] holds all the identifiers needed to attribute data to specific
18/// organizations, projects, and keys. This allows the rate limiting and quota
19/// systems to enforce limits at the appropriate scope levels.
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
21pub struct Scoping {
22    /// The organization id.
23    pub organization_id: OrganizationId,
24
25    /// The project id.
26    pub project_id: ProjectId,
27
28    /// The DSN public key.
29    pub project_key: ProjectKey,
30
31    /// The public key's internal id.
32    pub key_id: Option<u64>,
33}
34
35impl Scoping {
36    /// Creates an [`ItemScoping`] for a specific data category in this scope.
37    ///
38    /// The returned item scoping contains a reference to this scope and the provided
39    /// data category. This is a cheap operation that allows for efficient rate limiting
40    /// of individual items.
41    pub fn item(&self, category: DataCategory) -> ItemScoping {
42        ItemScoping {
43            category,
44            scoping: *self,
45            namespace: MetricNamespaceScoping::None,
46        }
47    }
48
49    /// Creates an [`ItemScoping`] specifically for metric buckets in this scope.
50    ///
51    /// The returned item scoping contains a reference to this scope, the
52    /// [`DataCategory::MetricBucket`] category, and the provided metric namespace.
53    /// This is specialized for handling metrics with namespaces.
54    pub fn metric_bucket(&self, namespace: MetricNamespace) -> ItemScoping {
55        ItemScoping {
56            category: DataCategory::MetricBucket,
57            scoping: *self,
58            namespace: MetricNamespaceScoping::Some(namespace),
59        }
60    }
61}
62
63/// Describes the metric namespace scoping of an item.
64///
65/// This enum is used within [`ItemScoping`] to represent the metric namespace of an item.
66/// It handles the different cases: no namespace, a specific namespace, or any namespace.
67#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash, PartialOrd)]
68pub enum MetricNamespaceScoping {
69    /// The item does not contain metrics of any namespace.
70    ///
71    /// This should only be used for non-metric items.
72    #[default]
73    None,
74
75    /// The item contains metrics of a specific namespace.
76    Some(MetricNamespace),
77
78    /// The item contains metrics of any namespace.
79    ///
80    /// The specific namespace is not known or relevant. This can be used to check rate
81    /// limits or quotas that should apply to any namespace.
82    Any,
83}
84
85impl MetricNamespaceScoping {
86    /// Checks if the given namespace matches this namespace scoping.
87    ///
88    /// Returns `true` in the following cases:
89    /// - If `self` is [`MetricNamespaceScoping::Some`] with the same namespace
90    /// - If `self` is [`MetricNamespaceScoping::Any`], matching any namespace
91    pub fn matches(&self, namespace: MetricNamespace) -> bool {
92        match self {
93            Self::None => false,
94            Self::Some(ns) => *ns == namespace,
95            Self::Any => true,
96        }
97    }
98}
99
100impl From<MetricNamespace> for MetricNamespaceScoping {
101    fn from(namespace: MetricNamespace) -> Self {
102        Self::Some(namespace)
103    }
104}
105
106/// Data categorization and scoping information for a single item.
107///
108/// [`ItemScoping`] combines a data category, scoping information, and optional
109/// metric namespace to fully define an item for rate limiting purposes.
110#[derive(Debug, Copy, Clone, Eq, PartialEq)]
111pub struct ItemScoping {
112    /// The data category of the item.
113    pub category: DataCategory,
114
115    /// Scoping of the data.
116    pub scoping: Scoping,
117
118    /// Namespace for metric items, requiring [`DataCategory::MetricBucket`].
119    pub namespace: MetricNamespaceScoping,
120}
121
122impl std::ops::Deref for ItemScoping {
123    type Target = Scoping;
124
125    fn deref(&self) -> &Self::Target {
126        &self.scoping
127    }
128}
129
130impl ItemScoping {
131    /// Returns the identifier for the given quota scope.
132    ///
133    /// Maps the quota scope type to the corresponding identifier from this scoping,
134    /// or `None` if the scope type doesn't have an applicable identifier.
135    pub fn scope_id(&self, scope: QuotaScope) -> Option<u64> {
136        match scope {
137            QuotaScope::Organization => Some(self.organization_id.value()),
138            QuotaScope::Project => Some(self.project_id.value()),
139            QuotaScope::Key => self.key_id,
140            QuotaScope::Unknown => None,
141        }
142    }
143
144    /// Checks whether the category matches any of the quota's categories.
145    pub(crate) fn matches_categories(&self, categories: DataCategories) -> bool {
146        // An empty list of categories means that this quota matches all categories. Note that we
147        // skip `Unknown` categories silently. If the list of categories only contains `Unknown`s,
148        // we do **not** match, since apparently the quota is meant for some data this Relay does
149        // not support yet.
150        categories.is_empty() || categories.contains(&self.category)
151    }
152
153    /// Returns `true` if the rate limit namespace matches the namespace of the item.
154    ///
155    /// Matching behavior depends on the passed namespaces and the namespace of the scoping:
156    ///  - If the list of namespaces is empty, this check always returns `true`.
157    ///  - If the list of namespaces contains at least one namespace, a namespace on the scoping is
158    ///    required. [`MetricNamespaceScoping::None`] will not match.
159    ///  - If the namespace of this scoping is [`MetricNamespaceScoping::Any`], this check will
160    ///    always return true.
161    ///  - Otherwise, an exact match of the scoping's namespace must be found in the list.
162    ///
163    /// `namespace` can be either a slice, an iterator, or a reference to an
164    /// `Option<MetricNamespace>`. In case of `None`, this method behaves like an empty list and
165    /// permits any namespace.
166    pub(crate) fn matches_namespaces<'a, I>(&self, namespaces: I) -> bool
167    where
168        I: IntoIterator<Item = &'a MetricNamespace>,
169    {
170        let mut iter = namespaces.into_iter().peekable();
171        iter.peek().is_none() || iter.any(|ns| self.namespace.matches(*ns))
172    }
173}
174
175/// An efficient container for data categories that avoids allocations.
176///
177/// It is a read only and has set like properties, allowing for fast comparisons.
178#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, Hash)]
179pub struct DataCategories(u64);
180
181impl Serialize for DataCategories {
182    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
183    where
184        S: serde::Serializer,
185    {
186        let mut ser = serializer.serialize_seq(Some(self.len()))?;
187        for category in self.iter() {
188            ser.serialize_element(&category)?;
189        }
190        ser.end()
191    }
192}
193
194impl DataCategories {
195    fn category_to_mask(category: DataCategory) -> u64 {
196        if matches!(category, DataCategory::Unknown) {
197            // Handle Unknown by throwing it into the last bit.
198            1u64 << 63
199        } else {
200            1u64 << (category as u8)
201        }
202    }
203
204    fn bit_number_to_category(bit_number: u8) -> DataCategory {
205        bit_number.try_into().unwrap_or(DataCategory::Unknown)
206    }
207
208    /// Creates new and empty [`DataCategories`].
209    pub fn new() -> Self {
210        Default::default()
211    }
212
213    /// Creates a new [`DataCategories`] from the supplied slice of [`DataCategory`] values.
214    pub fn from_slice(slice: &[DataCategory]) -> Self {
215        let mut categories = 0;
216        for category in slice {
217            categories |= DataCategories::category_to_mask(*category);
218        }
219
220        Self(categories)
221    }
222
223    /// Adds a data category to [`Self`].
224    ///
225    /// Returns `None` if the category was already contained, otherwise creates a new [`Self`] with
226    /// the `category` added.
227    pub fn add(&self, category: DataCategory) -> Option<Self> {
228        let category_mask = DataCategories::category_to_mask(category);
229
230        if (self.0 & category_mask) != 0 {
231            return None;
232        }
233
234        Some(Self(self.0 | category_mask))
235    }
236
237    /// Returns true iff the category is contained.
238    pub fn contains(&self, category: &DataCategory) -> bool {
239        let category_mask = DataCategories::category_to_mask(*category);
240        (self.0 & category_mask) != 0
241    }
242
243    /// Returns an iterator over this [`DataCategories`] container.
244    pub fn iter(&self) -> DataCategoryIterator {
245        // We start our iteration from the lsb, so the number of trailings zeroes is our 0-based
246        // starting index.
247        let bit_start = self.0.trailing_zeros();
248
249        // Shift our bitfield so that the first bit is in the 0th place.
250        let (categories, _) = self.0.overflowing_shr(bit_start);
251        DataCategoryIterator {
252            categories,
253            current_bit: bit_start as u8,
254        }
255    }
256
257    /// Returns the number of categories in this container.
258    pub fn len(&self) -> usize {
259        self.0.count_ones() as usize
260    }
261
262    /// Returns true iff this container contains no categories.
263    pub fn is_empty(&self) -> bool {
264        self.0 == 0
265    }
266}
267
268/// An iterator over a [`DataCategories`] container.
269pub struct DataCategoryIterator {
270    categories: u64,
271    current_bit: u8,
272}
273
274impl Iterator for DataCategoryIterator {
275    type Item = DataCategory;
276
277    fn next(&mut self) -> Option<Self::Item> {
278        if self.categories == 0 {
279            return None;
280        }
281        let category = DataCategories::bit_number_to_category(self.current_bit);
282
283        // Consume the last bit
284        self.categories &= u64::MAX - 1;
285
286        // Find the next set bit
287        let trailing_zeroes = self.categories.trailing_zeros() as u8;
288
289        // Put that bit in the 0th (to-consume) position
290        (self.categories, _) = self.categories.overflowing_shr(trailing_zeroes as u32);
291
292        // Update the current bit state
293        self.current_bit += trailing_zeroes;
294
295        Some(category)
296    }
297}
298
299impl<'de> Deserialize<'de> for DataCategories {
300    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
301    where
302        D: serde::Deserializer<'de>,
303    {
304        SmallVec::<[DataCategory; 12]>::deserialize(deserializer).map(Self::from)
305    }
306}
307
308impl<const N: usize> From<SmallVec<[DataCategory; N]>> for DataCategories {
309    fn from(categories: SmallVec<[DataCategory; N]>) -> Self {
310        Self::from_slice(&categories)
311    }
312}
313
314impl<const N: usize> From<[DataCategory; N]> for DataCategories {
315    fn from(categories: [DataCategory; N]) -> Self {
316        Self::from_slice(&categories)
317    }
318}
319
320impl FromIterator<DataCategory> for DataCategories {
321    fn from_iter<T: IntoIterator<Item = DataCategory>>(iter: T) -> Self {
322        let v: SmallVec<[DataCategory; 12]> = iter.into_iter().collect();
323        Self::from_slice(&v)
324    }
325}
326
327/// The scope at which a quota is applied.
328///
329/// Defines the granularity at which quotas are enforced, from organizations
330/// down to individual project keys. This enum only defines the type of scope,
331/// not the specific instance.
332///
333/// This type is directly related to [`crate::rate_limit::RateLimitScope`], which
334/// includes the specific scope identifiers.
335#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)]
336#[serde(rename_all = "lowercase")]
337pub enum QuotaScope {
338    /// The organization level.
339    ///
340    /// This is the top-level scope.
341    Organization,
342
343    /// The project level.
344    ///
345    /// Projects are contained within organizations.
346    Project,
347
348    /// The project key level (corresponds to a DSN).
349    ///
350    /// This is the most specific scope level and is contained within projects.
351    Key,
352
353    /// Any scope type not recognized by this Relay.
354    #[serde(other)]
355    Unknown,
356}
357
358impl QuotaScope {
359    /// Returns the quota scope corresponding to the given name string.
360    ///
361    /// If the string doesn't match any known scope, returns [`QuotaScope::Unknown`].
362    pub fn from_name(string: &str) -> Self {
363        match string {
364            "organization" => Self::Organization,
365            "project" => Self::Project,
366            "key" => Self::Key,
367            _ => Self::Unknown,
368        }
369    }
370
371    /// Returns the canonical string name of this scope.
372    ///
373    /// This is the lowercase string representation used in serialization.
374    pub fn name(self) -> &'static str {
375        match self {
376            Self::Key => "key",
377            Self::Project => "project",
378            Self::Organization => "organization",
379            Self::Unknown => "unknown",
380        }
381    }
382}
383
384impl fmt::Display for QuotaScope {
385    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
386        write!(f, "{}", self.name())
387    }
388}
389
390impl FromStr for QuotaScope {
391    type Err = ();
392
393    fn from_str(string: &str) -> Result<Self, Self::Err> {
394        Ok(Self::from_name(string))
395    }
396}
397
398fn default_scope() -> QuotaScope {
399    QuotaScope::Organization
400}
401
402/// A machine-readable reason code for rate limits.
403///
404/// Reason codes provide a standardized way to communicate why a particular
405/// item was rate limited.
406#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, Hash)]
407pub struct ReasonCode(Arc<str>);
408
409impl ReasonCode {
410    /// Creates a new reason code from a string.
411    ///
412    /// This method is primarily intended for testing. In production, reason codes
413    /// should typically be deserialized from quota configurations rather than
414    /// constructed manually.
415    pub fn new<S: Into<Arc<str>>>(code: S) -> Self {
416        Self(code.into())
417    }
418
419    /// Returns the string representation of this reason code.
420    pub fn as_str(&self) -> &str {
421        &self.0
422    }
423}
424
425impl fmt::Display for ReasonCode {
426    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
427        self.0.fmt(f)
428    }
429}
430
431/// Configuration for a data ingestion quota.
432///
433/// A quota defines restrictions on data ingestion based on data categories, scopes,
434/// and time windows. The system applies multiple quotas to incoming data, and items
435/// are counted against all matching quotas based on their categories.
436///
437/// Quotas can either:
438/// - Reject all data (`limit` = 0)
439/// - Limit data to a specific quantity per time window (`limit` > 0)
440/// - Count data without limiting it (`limit` = None)
441///
442/// Different quotas may apply at different scope levels (organization, project, key).
443#[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
444#[serde(rename_all = "camelCase")]
445pub struct Quota {
446    /// The unique identifier for counting this quota.
447    ///
448    /// Required for all quotas except those with `limit` = 0, which are statically enforced.
449    #[serde(default)]
450    pub id: Option<Arc<str>>,
451
452    /// Data categories this quota applies to.
453    ///
454    /// If missing or empty, this quota applies to all data categories.
455    #[serde(default)]
456    pub categories: DataCategories,
457
458    /// The scope level at which this quota is enforced.
459    ///
460    /// The quota is enforced separately within each instance of this scope
461    /// (e.g., for each project key separately). Defaults to [`QuotaScope::Organization`].
462    #[serde(default = "default_scope")]
463    pub scope: QuotaScope,
464
465    /// Specific scope instance identifier this quota applies to.
466    ///
467    /// If set, this quota only applies to the specified scope instance
468    /// (e.g., a specific project key). Requires `scope` to be set explicitly.
469    #[serde(default, skip_serializing_if = "Option::is_none")]
470    pub scope_id: Option<Arc<str>>,
471
472    /// Maximum number of events allowed within the time window.
473    ///
474    /// Possible values:
475    /// - `Some(0)`: Reject all matching events
476    /// - `Some(n)`: Allow up to n events per time window
477    /// - `None`: Unlimited quota (counts but doesn't limit)
478    ///
479    /// Requires `window` to be set if the limit is not 0.
480    #[serde(default)]
481    pub limit: Option<u64>,
482
483    /// The time window in seconds for quota enforcement.
484    ///
485    /// Required in all cases except `limit` = 0, since those quotas
486    /// are not measured over time.
487    #[serde(default, skip_serializing_if = "Option::is_none")]
488    pub window: Option<u64>,
489
490    /// The metric namespace this quota applies to.
491    ///
492    /// If `None`, it matches any namespace.
493    pub namespace: Option<MetricNamespace>,
494
495    /// A machine-readable reason code returned when this quota is exceeded.
496    ///
497    /// Required for all quotas except those with `limit` = None, since
498    /// unlimited quotas can never be exceeded.
499    #[serde(default, skip_serializing_if = "Option::is_none")]
500    pub reason_code: Option<ReasonCode>,
501}
502
503impl Quota {
504    /// Returns whether this quota is valid for tracking.
505    ///
506    /// A quota is considered invalid if any of the following conditions are true:
507    ///  - The quota only applies to [`DataCategory::Unknown`] data categories.
508    ///  - The quota is counted (not limit `0`) but specifies categories with different units.
509    ///  - The quota references an unsupported namespace.
510    pub fn is_valid(&self) -> bool {
511        if self.namespace == Some(MetricNamespace::Unsupported) {
512            return false;
513        }
514
515        let mut units = self
516            .categories
517            .iter()
518            .filter_map(CategoryUnit::from_category);
519
520        match units.next() {
521            // There are only unknown categories, which is always invalid
522            None if !self.categories.is_empty() => false,
523            // This is a reject all quota, which is always valid
524            _ if self.limit == Some(0) => true,
525            // Applies to all categories, which implies multiple units
526            None => false,
527            // There are multiple categories, which must all have the same units
528            Some(unit) => units.all(|u| u == unit),
529        }
530    }
531
532    /// Checks whether this quota's scope matches the given item scoping.
533    ///
534    /// This quota matches, if:
535    ///  - there is no `scope_id` constraint
536    ///  - the `scope_id` constraint is not numeric
537    ///  - the scope identifier matches the one from ascoping and the scope is known
538    fn matches_scope(&self, scoping: ItemScoping) -> bool {
539        // Check for a scope identifier constraint. If there is no constraint, this means that the
540        // quota matches any scope. In case the scope is unknown, it will be coerced to the most
541        // specific scope later.
542        let Some(scope_id) = self.scope_id.as_ref() else {
543            return true;
544        };
545
546        // Check if the scope identifier in the quota is parseable. If not, this means we cannot
547        // fulfill the constraint, so the quota does not match.
548        let Ok(parsed) = scope_id.parse::<u64>() else {
549            return false;
550        };
551
552        // At this stage, require that the scope is known since we have to fulfill the constraint.
553        scoping.scope_id(self.scope) == Some(parsed)
554    }
555
556    /// Checks whether the quota's constraints match the current item.
557    ///
558    /// This method determines if this quota should be applied to a given item
559    /// based on its scope, categories, and namespace.
560    pub fn matches(&self, scoping: ItemScoping) -> bool {
561        self.matches_scope(scoping)
562            && scoping.matches_categories(self.categories)
563            && scoping.matches_namespaces(&self.namespace)
564    }
565}
566
567#[cfg(test)]
568mod tests {
569    use relay_base_schema::data_category::UnknownDataCategory;
570
571    use super::*;
572
573    #[test]
574    fn test_parse_quota_reject_all() {
575        let json = r#"{
576            "limit": 0,
577            "reasonCode": "not_yet"
578        }"#;
579
580        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
581
582        insta::assert_ron_snapshot!(quota, @r###"
583        Quota(
584          id: None,
585          categories: [],
586          scope: organization,
587          limit: Some(0),
588          namespace: None,
589          reasonCode: Some(ReasonCode("not_yet")),
590        )
591        "###);
592    }
593
594    #[test]
595    fn test_parse_quota_reject_transactions() {
596        let json = r#"{
597            "limit": 0,
598            "categories": ["transaction"],
599            "reasonCode": "not_yet"
600        }"#;
601
602        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
603
604        insta::assert_ron_snapshot!(quota, @r#"
605        Quota(
606          id: None,
607          categories: [
608            "transaction",
609          ],
610          scope: organization,
611          limit: Some(0),
612          namespace: None,
613          reasonCode: Some(ReasonCode("not_yet")),
614        )
615        "#);
616    }
617
618    #[test]
619    fn test_parse_quota_limited() {
620        let json = r#"{
621            "id": "o",
622            "limit": 4711,
623            "window": 42,
624            "reasonCode": "not_so_fast"
625        }"#;
626
627        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
628
629        insta::assert_ron_snapshot!(quota, @r###"
630        Quota(
631          id: Some("o"),
632          categories: [],
633          scope: organization,
634          limit: Some(4711),
635          window: Some(42),
636          namespace: None,
637          reasonCode: Some(ReasonCode("not_so_fast")),
638        )
639        "###);
640    }
641
642    #[test]
643    fn test_parse_quota_project() {
644        let json = r#"{
645            "id": "p",
646            "scope": "project",
647            "scopeId": "1",
648            "limit": 4711,
649            "window": 42,
650            "reasonCode": "not_so_fast"
651        }"#;
652
653        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
654
655        insta::assert_ron_snapshot!(quota, @r###"
656        Quota(
657          id: Some("p"),
658          categories: [],
659          scope: project,
660          scopeId: Some("1"),
661          limit: Some(4711),
662          window: Some(42),
663          namespace: None,
664          reasonCode: Some(ReasonCode("not_so_fast")),
665        )
666        "###);
667    }
668
669    #[test]
670    fn test_parse_quota_project_large() {
671        let json = r#"{
672            "id": "p",
673            "scope": "project",
674            "scopeId": "1",
675            "limit": 4294967296,
676            "window": 42,
677            "reasonCode": "not_so_fast"
678        }"#;
679
680        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
681
682        insta::assert_ron_snapshot!(quota, @r###"
683        Quota(
684          id: Some("p"),
685          categories: [],
686          scope: project,
687          scopeId: Some("1"),
688          limit: Some(4294967296),
689          window: Some(42),
690          namespace: None,
691          reasonCode: Some(ReasonCode("not_so_fast")),
692        )
693        "###);
694    }
695
696    #[test]
697    fn test_parse_quota_key() {
698        let json = r#"{
699            "id": "k",
700            "scope": "key",
701            "scopeId": "1",
702            "limit": 4711,
703            "window": 42,
704            "reasonCode": "not_so_fast"
705        }"#;
706
707        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
708
709        insta::assert_ron_snapshot!(quota, @r###"
710        Quota(
711          id: Some("k"),
712          categories: [],
713          scope: key,
714          scopeId: Some("1"),
715          limit: Some(4711),
716          window: Some(42),
717          namespace: None,
718          reasonCode: Some(ReasonCode("not_so_fast")),
719        )
720        "###);
721    }
722
723    #[test]
724    fn test_parse_quota_unknown_variants() {
725        let json = r#"{
726            "id": "f",
727            "categories": ["future"],
728            "scope": "future",
729            "scopeId": "1",
730            "limit": 4711,
731            "window": 42,
732            "reasonCode": "not_so_fast"
733        }"#;
734
735        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
736
737        insta::assert_ron_snapshot!(quota, @r#"
738        Quota(
739          id: Some("f"),
740          categories: [
741            "unknown",
742          ],
743          scope: unknown,
744          scopeId: Some("1"),
745          limit: Some(4711),
746          window: Some(42),
747          namespace: None,
748          reasonCode: Some(ReasonCode("not_so_fast")),
749        )
750        "#);
751    }
752
753    #[test]
754    fn test_parse_quota_unlimited() {
755        let json = r#"{
756            "id": "o",
757            "window": 42
758        }"#;
759
760        let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
761
762        insta::assert_ron_snapshot!(quota, @r###"
763        Quota(
764          id: Some("o"),
765          categories: [],
766          scope: organization,
767          limit: None,
768          window: Some(42),
769          namespace: None,
770        )
771        "###);
772    }
773
774    #[test]
775    fn test_quota_valid_reject_all() {
776        let quota = Quota {
777            id: None,
778            categories: Default::default(),
779            scope: QuotaScope::Organization,
780            scope_id: None,
781            limit: Some(0),
782            window: None,
783            reason_code: None,
784            namespace: None,
785        };
786
787        assert!(quota.is_valid());
788    }
789
790    #[test]
791    fn test_quota_invalid_only_unknown() {
792        let quota = Quota {
793            id: None,
794            categories: [DataCategory::Unknown, DataCategory::Unknown].into(),
795            scope: QuotaScope::Organization,
796            scope_id: None,
797            limit: Some(0),
798            window: None,
799            reason_code: None,
800            namespace: None,
801        };
802
803        assert!(!quota.is_valid());
804    }
805
806    #[test]
807    fn test_quota_valid_reject_all_mixed() {
808        let quota = Quota {
809            id: None,
810            categories: [DataCategory::Error, DataCategory::Attachment].into(),
811            scope: QuotaScope::Organization,
812            scope_id: None,
813            limit: Some(0),
814            window: None,
815            reason_code: None,
816            namespace: None,
817        };
818
819        assert!(quota.is_valid());
820    }
821
822    #[test]
823    fn test_quota_invalid_limited_mixed() {
824        let quota = Quota {
825            id: None,
826            categories: [DataCategory::Error, DataCategory::Attachment].into(),
827            scope: QuotaScope::Organization,
828            scope_id: None,
829            limit: Some(1000),
830            window: None,
831            reason_code: None,
832            namespace: None,
833        };
834
835        // This category is limited and counted, but has multiple units.
836        assert!(!quota.is_valid());
837    }
838
839    #[test]
840    fn test_quota_invalid_unlimited_mixed() {
841        let quota = Quota {
842            id: None,
843            categories: [DataCategory::Error, DataCategory::Attachment].into(),
844            scope: QuotaScope::Organization,
845            scope_id: None,
846            limit: None,
847            window: None,
848            reason_code: None,
849            namespace: None,
850        };
851
852        // This category is unlimited and counted, but has multiple units.
853        assert!(!quota.is_valid());
854    }
855
856    #[test]
857    fn test_quota_matches_no_categories() {
858        let quota = Quota {
859            id: None,
860            categories: Default::default(),
861            scope: QuotaScope::Organization,
862            scope_id: None,
863            limit: None,
864            window: None,
865            reason_code: None,
866            namespace: None,
867        };
868
869        assert!(quota.matches(ItemScoping {
870            category: DataCategory::Error,
871            scoping: Scoping {
872                organization_id: OrganizationId::new(42),
873                project_id: ProjectId::new(21),
874                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
875                key_id: Some(17),
876            },
877            namespace: MetricNamespaceScoping::None,
878        }));
879    }
880
881    #[test]
882    fn test_quota_matches_unknown_category() {
883        let quota = Quota {
884            id: None,
885            categories: [DataCategory::Unknown].into(),
886            scope: QuotaScope::Organization,
887            scope_id: None,
888            limit: None,
889            window: None,
890            reason_code: None,
891            namespace: None,
892        };
893
894        assert!(!quota.matches(ItemScoping {
895            category: DataCategory::Error,
896            scoping: Scoping {
897                organization_id: OrganizationId::new(42),
898                project_id: ProjectId::new(21),
899                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
900                key_id: Some(17),
901            },
902            namespace: MetricNamespaceScoping::None,
903        }));
904    }
905
906    #[test]
907    fn test_quota_matches_multiple_categores() {
908        let quota = Quota {
909            id: None,
910            categories: [DataCategory::Unknown, DataCategory::Error].into(),
911            scope: QuotaScope::Organization,
912            scope_id: None,
913            limit: None,
914            window: None,
915            reason_code: None,
916            namespace: None,
917        };
918
919        assert!(quota.matches(ItemScoping {
920            category: DataCategory::Error,
921            scoping: Scoping {
922                organization_id: OrganizationId::new(42),
923                project_id: ProjectId::new(21),
924                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
925                key_id: Some(17),
926            },
927            namespace: MetricNamespaceScoping::None,
928        }));
929
930        assert!(!quota.matches(ItemScoping {
931            category: DataCategory::Transaction,
932            scoping: Scoping {
933                organization_id: OrganizationId::new(42),
934                project_id: ProjectId::new(21),
935                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
936                key_id: Some(17),
937            },
938            namespace: MetricNamespaceScoping::None,
939        }));
940    }
941
942    #[test]
943    fn test_quota_matches_no_invalid_scope() {
944        let quota = Quota {
945            id: None,
946            categories: Default::default(),
947            scope: QuotaScope::Organization,
948            scope_id: Some("not_a_number".into()),
949            limit: None,
950            window: None,
951            reason_code: None,
952            namespace: None,
953        };
954
955        assert!(!quota.matches(ItemScoping {
956            category: DataCategory::Error,
957            scoping: Scoping {
958                organization_id: OrganizationId::new(42),
959                project_id: ProjectId::new(21),
960                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
961                key_id: Some(17),
962            },
963            namespace: MetricNamespaceScoping::None,
964        }));
965    }
966
967    #[test]
968    fn test_quota_matches_organization_scope() {
969        let quota = Quota {
970            id: None,
971            categories: Default::default(),
972            scope: QuotaScope::Organization,
973            scope_id: Some("42".into()),
974            limit: None,
975            window: None,
976            reason_code: None,
977            namespace: None,
978        };
979
980        assert!(quota.matches(ItemScoping {
981            category: DataCategory::Error,
982            scoping: Scoping {
983                organization_id: OrganizationId::new(42),
984                project_id: ProjectId::new(21),
985                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
986                key_id: Some(17),
987            },
988            namespace: MetricNamespaceScoping::None,
989        }));
990
991        assert!(!quota.matches(ItemScoping {
992            category: DataCategory::Error,
993            scoping: Scoping {
994                organization_id: OrganizationId::new(0),
995                project_id: ProjectId::new(21),
996                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
997                key_id: Some(17),
998            },
999            namespace: MetricNamespaceScoping::None,
1000        }));
1001    }
1002
1003    #[test]
1004    fn test_quota_matches_project_scope() {
1005        let quota = Quota {
1006            id: None,
1007            categories: Default::default(),
1008            scope: QuotaScope::Project,
1009            scope_id: Some("21".into()),
1010            limit: None,
1011            window: None,
1012            reason_code: None,
1013            namespace: None,
1014        };
1015
1016        assert!(quota.matches(ItemScoping {
1017            category: DataCategory::Error,
1018            scoping: Scoping {
1019                organization_id: OrganizationId::new(42),
1020                project_id: ProjectId::new(21),
1021                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1022                key_id: Some(17),
1023            },
1024            namespace: MetricNamespaceScoping::None,
1025        }));
1026
1027        assert!(!quota.matches(ItemScoping {
1028            category: DataCategory::Error,
1029            scoping: Scoping {
1030                organization_id: OrganizationId::new(42),
1031                project_id: ProjectId::new(0),
1032                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1033                key_id: Some(17),
1034            },
1035            namespace: MetricNamespaceScoping::None,
1036        }));
1037    }
1038
1039    #[test]
1040    fn test_quota_matches_key_scope() {
1041        let quota = Quota {
1042            id: None,
1043            categories: Default::default(),
1044            scope: QuotaScope::Key,
1045            scope_id: Some("17".into()),
1046            limit: None,
1047            window: None,
1048            reason_code: None,
1049            namespace: None,
1050        };
1051
1052        assert!(quota.matches(ItemScoping {
1053            category: DataCategory::Error,
1054            scoping: Scoping {
1055                organization_id: OrganizationId::new(42),
1056                project_id: ProjectId::new(21),
1057                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1058                key_id: Some(17),
1059            },
1060            namespace: MetricNamespaceScoping::None,
1061        }));
1062
1063        assert!(!quota.matches(ItemScoping {
1064            category: DataCategory::Error,
1065            scoping: Scoping {
1066                organization_id: OrganizationId::new(42),
1067                project_id: ProjectId::new(21),
1068                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1069                key_id: Some(0),
1070            },
1071            namespace: MetricNamespaceScoping::None,
1072        }));
1073
1074        assert!(!quota.matches(ItemScoping {
1075            category: DataCategory::Error,
1076            scoping: Scoping {
1077                organization_id: OrganizationId::new(42),
1078                project_id: ProjectId::new(21),
1079                project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
1080                key_id: None,
1081            },
1082            namespace: MetricNamespaceScoping::None,
1083        }));
1084    }
1085
1086    #[test]
1087    fn test_data_categories_sorted_deduplicated() {
1088        let a = DataCategories::from([
1089            DataCategory::Transaction,
1090            DataCategory::Span,
1091            DataCategory::Transaction,
1092        ]);
1093        let b = DataCategories::from([
1094            DataCategory::Span,
1095            DataCategory::Transaction,
1096            DataCategory::Span,
1097        ]);
1098        let c = DataCategories::from([DataCategory::Span, DataCategory::Transaction]);
1099
1100        assert_eq!(a, b);
1101        assert_eq!(b, c);
1102        assert_eq!(a, c);
1103    }
1104
1105    #[test]
1106    fn test_data_categories_serde() {
1107        let s: DataCategories = serde_json::from_str(r#"["span", "transaction", "span"]"#).unwrap();
1108        insta::assert_json_snapshot!(s, @r#"
1109        [
1110          "transaction",
1111          "span"
1112        ]
1113        "#);
1114    }
1115
1116    #[test]
1117    fn test_data_categories_add() {
1118        let c = DataCategories::new();
1119        let c = c.add(DataCategory::Span).unwrap();
1120        assert!(c.add(DataCategory::Span).is_none());
1121        let c = c.add(DataCategory::Transaction).unwrap();
1122        assert_eq!(c, [DataCategory::Span, DataCategory::Transaction].into());
1123    }
1124
1125    #[test]
1126    fn test_reserve_last_bit_data_category() {
1127        // Ensure we don't accidentally use 63 as a valid DataCategory.
1128        let r: Result<DataCategory, UnknownDataCategory> = 63u32.try_into();
1129        assert!(r.is_err());
1130    }
1131
1132    #[test]
1133    fn test_bitmapping_identity() {
1134        assert_eq!(
1135            DataCategories::category_to_mask(DataCategory::Unknown),
1136            1u64 << 63
1137        );
1138        for i in 0..63u32 {
1139            let r: Result<DataCategory, UnknownDataCategory> = i.try_into();
1140            if let Ok(dc) = r {
1141                assert_eq!(DataCategories::bit_number_to_category(i as u8), dc);
1142                assert_eq!(DataCategories::category_to_mask(dc), 1 << i);
1143            } else {
1144                break;
1145            }
1146        }
1147    }
1148}
relay_quotas/quota.rs

relay_quotas/
quota.rs
