getsentry · jan-auer · Mar 28, 2024 · Mar 27, 2024 · Mar 27, 2024 · Mar 27, 2024
@@ -25,6 +25,7 @@
 - Collect duration for all spans. ([#3322](https://github.com/getsentry/relay/pull/3322))
 - Add `project_id` as part of the span Kafka message headers. ([#3320](https://github.com/getsentry/relay/pull/3320))
 - Stop producing to sessions topic, the feature is now fully migrated to metrics. ([#3271](https://github.com/getsentry/relay/pull/3271))
+- Support and expose namespaces for metric rate limit propagation via the `x-sentry-rate-limits` header. ([#3347](https://github.com/getsentry/relay/pull/3347))
 
 ## 24.3.0
 

@@ -252,7 +252,7 @@ mod tests {
     use relay_common::time::UnixTimestamp;
     use relay_redis::{RedisConfigOptions, RedisPool};
 
-    use crate::{DataCategories, ItemScoping, Quota, QuotaScope, Scoping};
+    use crate::{DataCategories, Quota, QuotaScope, Scoping};
 
     fn build_redis_pool() -> RedisPool {
         let url = std::env::var("RELAY_REDIS_URL")
@@ -284,11 +284,7 @@ mod tests {
     }
 
     fn build_redis_quota<'a>(quota: &'a Quota, scoping: &'a Scoping) -> RedisQuota<'a> {
-        let scoping = ItemScoping {
-            category: DataCategory::MetricBucket,
-            scoping,
-            namespace: None,
-        };
+        let scoping = scoping.item(DataCategory::MetricBucket);
         RedisQuota::new(quota, scoping, UnixTimestamp::now()).unwrap()
     }
 
@@ -470,18 +466,14 @@ mod tests {
         let ts = UnixTimestamp::now();
         let quota = build_quota(window, limit);
         let scoping = build_scoping();
-        let scoping = ItemScoping {
-            category: DataCategory::MetricBucket,
-            scoping: &scoping,
-            namespace: None,
-        };
+        let item_scoping = scoping.item(DataCategory::MetricBucket);
 
         let pool = build_redis_pool();
         let mut client = pool.client().unwrap();
 
         let rl = GlobalRateLimits::default();
 
-        let redis_quota = [RedisQuota::new(&quota, scoping, ts).unwrap()];
+        let redis_quota = [RedisQuota::new(&quota, item_scoping, ts).unwrap()];
         assert!(rl
             .filter_rate_limited(&mut client, &redis_quota, 200)
             .unwrap()
@@ -494,7 +486,10 @@ mod tests {
 
         // Fast forward time.
         let redis_quota =
-            [RedisQuota::new(&quota, scoping, ts + Duration::from_secs(window + 1)).unwrap()];
+            [
+                RedisQuota::new(&quota, item_scoping, ts + Duration::from_secs(window + 1))
+                    .unwrap(),
+            ];
         assert!(rl
             .filter_rate_limited(&mut client, &redis_quota, 200)
             .unwrap()
@@ -513,11 +508,8 @@ mod tests {
         let timestamp = UnixTimestamp::now();
 
         let mut quota = build_quota(100, limit);
-        let scoping = ItemScoping {
-            category: DataCategory::MetricBucket,
-            scoping: &build_scoping(),
-            namespace: None,
-        };
+        let scoping = build_scoping();
+        let item_scoping = scoping.item(DataCategory::MetricBucket);
 
         let pool = build_redis_pool();
         let mut client = pool.client().unwrap();
@@ -527,7 +519,7 @@ mod tests {
         let quantity = 2;
         let redis_threshold = (quantity as f32 / DEFAULT_BUDGET_RATIO) as u64;
         for _ in 0..redis_threshold + 10 {
-            let redis_quota = RedisQuota::new(&quota, scoping, timestamp).unwrap();
+            let redis_quota = RedisQuota::new(&quota, item_scoping, timestamp).unwrap();
             assert!(rl
                 .filter_rate_limited(&mut client, &[redis_quota], quantity)
                 .unwrap()
@@ -539,7 +531,7 @@ mod tests {
         let rl = GlobalRateLimits::default();
 
         quota.limit = Some(redis_threshold);
-        let redis_quota = RedisQuota::new(&quota, scoping, timestamp).unwrap();
+        let redis_quota = RedisQuota::new(&quota, item_scoping, timestamp).unwrap();
 
         assert!(!rl
             .filter_rate_limited(&mut client, &[redis_quota], quantity)

@@ -36,11 +36,63 @@ impl Scoping {
         ItemScoping {
             category,
             scoping: self,
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
+        }
+    }
+
+    /// Returns an `ItemScoping` for metric buckets in this scope.
+    ///
+    /// The item scoping will contain a reference to this scope, the category
+    /// [`MetricBucket](DataCategory::MetricBucket), and the information passed to this function.
+    /// This is a cheap operation to allow rate limiting for an individual item.
+    pub fn metric_bucket(&self, namespace: MetricNamespace) -> ItemScoping<'_> {
+        ItemScoping {
+            category: DataCategory::MetricBucket,
+            scoping: self,
+            namespace: MetricNamespaceScoping::Some(namespace),
+        }
+    }
+}
+
+/// Item scoping of metric namespaces.
+///
+/// This enum is used in [`ItemScoping`] to declare the contents of an item's metric namespace.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Hash, PartialOrd)]
+pub enum MetricNamespaceScoping {
+    /// This item does not contain metrics of any namespace. This should only be used for non-metric
+    /// items.
+    #[default]
+    None,
+
+    /// This item contains metrics of a specific namespace.
+    Some(MetricNamespace),
+
+    /// This item contains metrics of any namespace.
+    ///
+    /// The namespace of metrics contained in this item is not known. This can be used to check rate
+    /// limits or quotas of any namespace.
+    Any,
+}
+
+impl MetricNamespaceScoping {
+    /// Returns `true` if the given namespace matches the namespace of the item.
+    ///
+    /// If the self is `Any`, this method returns `true` for any namespace.
+    pub fn matches(&self, namespace: MetricNamespace) -> bool {
+        match self {
+            Self::None => false,
+            Self::Some(ns) => *ns == namespace,
+            Self::Any => true,
         }
     }
 }
 
+impl From<MetricNamespace> for MetricNamespaceScoping {
+    fn from(namespace: MetricNamespace) -> Self {
+        Self::Some(namespace)
+    }
+}
+
 /// Data categorization and scoping information.
 ///
 /// `ItemScoping` is always attached to a `Scope` and references it internally. It is a cheap,
@@ -54,8 +106,8 @@ pub struct ItemScoping<'a> {
     /// Scoping of the data.
     pub scoping: &'a Scoping,
 
-    /// Namespace if applicable. `None` matches any namespace.
-    pub namespace: Option<MetricNamespace>,
+    /// Namespace for metric items, requiring [`DataCategory::MetricBucket`].
+    pub namespace: MetricNamespaceScoping,
 }
 
 impl AsRef<Scoping> for ItemScoping<'_> {
@@ -92,6 +144,27 @@ impl ItemScoping<'_> {
         // not support yet.
         categories.is_empty() || categories.iter().any(|cat| *cat == self.category)
     }
+
+    /// Returns `true` if the rate limit namespace matches the namespace of the item.
+    ///
+    /// Matching behavior depends on the passed namespaces and the namespace of the scoping:
+    ///  - If the list of namespaces is empty, this check always returns `true`.
+    ///  - If the list of namespaces contains at least one namespace, a namespace on the scoping is
+    ///    required. [`MetricNamespaceScoping::None`] will not match.
+    ///  - If the namespace of this scoping is [`MetricNamespaceScoping::Any`], this check will
+    ///    always return true.
+    ///  - Otherwise, an exact match of the scoping's namespace must be found in the list.
+    ///
+    /// `namespace` can be either a slice, an iterator, or a reference to an
+    /// `Option<MetricNamespace>`. In case of `None`, this method behaves like an empty list and
+    /// permits any namespace.
+    pub(crate) fn matches_namespaces<'a, I>(&self, namespaces: I) -> bool
+    where
+        I: IntoIterator<Item = &'a MetricNamespace>,
+    {
+        let mut iter = namespaces.into_iter().peekable();
+        iter.peek().is_none() || iter.any(|ns| self.namespace.matches(*ns))
+    }
 }
 
 /// The unit in which a data category is measured.
@@ -312,11 +385,6 @@ impl Quota {
     ///  - the `scope_id` constraint is not numeric
     ///  - the scope identifier matches the one from ascoping and the scope is known
     fn matches_scope(&self, scoping: ItemScoping<'_>) -> bool {
-        // Accept all types of namespaces if none are configured in the quota.
-        if self.namespace.is_some() && self.namespace != scoping.namespace {
-            return false;
-        }
-
         if self.scope == QuotaScope::Global {
             return true;
         }
@@ -340,7 +408,9 @@ impl Quota {
 
     /// Checks whether the quota's constraints match the current item.
     pub fn matches(&self, scoping: ItemScoping<'_>) -> bool {
-        self.matches_scope(scoping) && scoping.matches_categories(&self.categories)
+        self.matches_scope(scoping)
+            && scoping.matches_categories(&self.categories)
+            && scoping.matches_namespaces(&self.namespace)
     }
 }
 
@@ -654,7 +724,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 
@@ -679,7 +749,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 
@@ -704,7 +774,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
 
         assert!(!quota.matches(ItemScoping {
@@ -715,7 +785,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 
@@ -740,7 +810,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 
@@ -765,7 +835,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
 
         assert!(!quota.matches(ItemScoping {
@@ -776,7 +846,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 
@@ -801,7 +871,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
 
         assert!(!quota.matches(ItemScoping {
@@ -812,7 +882,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 
@@ -837,7 +907,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(17),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
 
         assert!(!quota.matches(ItemScoping {
@@ -848,7 +918,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: Some(0),
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
 
         assert!(!quota.matches(ItemScoping {
@@ -859,7 +929,7 @@ mod tests {
                 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
                 key_id: None,
             },
-            namespace: None,
+            namespace: MetricNamespaceScoping::None,
         }));
     }
 }