global config and tests

getsentry · Dav1dde · Mar 19, 2024 · Mar 18, 2024 · Mar 19, 2024 · Mar 19, 2024
commit 74e0747cc83315e4f024528d5dc44331a28d8ca1
diff --git a/relay-dynamic-config/src/feature.rs b/relay-dynamic-config/src/feature.rs
@@ -45,11 +45,6 @@ pub enum Feature {
     /// Serialized as `organizations:custom-metrics`.
     #[serde(rename = "organizations:custom-metrics")]
     CustomMetrics,
-    /// Enables metric stats.
-    ///
-    /// Serialized as `organizations:metric-stats`.
-    #[serde(rename = "organizations:metric-stats")]
-    MetricStats,
 
     /// Enable processing profiles.
     ///

@@ -145,6 +145,17 @@ pub struct Options {
     )]
     pub metric_bucket_dist_encodings: BucketEncodings,
 
+    /// Rollout rate for metric stats.
+    ///
+    /// Rate needs to be between `0.0` and `1.0`.
+    /// If set to `1.0` all organizations will have metric stats enabled.
+    #[serde(
+        rename = "relay.metric-stats.rollout-rate",
+        deserialize_with = "default_on_error",
+        skip_serializing_if = "is_default"
+    )]
+    pub metric_stats_rollout_rate: f32,
+
     /// All other unknown options.
     #[serde(flatten)]
     other: HashMap<String, Value>,

@@ -9,6 +9,7 @@ use relay_metrics::{
 use relay_quotas::Scoping;
 use relay_system::Addr;
 
+use crate::services::global_config::GlobalConfigHandle;
 use crate::services::outcome::Outcome;
 
 fn volume_metric_mri() -> Arc<str> {
@@ -27,18 +28,27 @@ fn volume_metric_mri() -> Arc<str> {
 #[derive(Clone, Debug)]
 pub struct MetricStats {
     config: Arc<Config>,
+    global_config: GlobalConfigHandle,
     aggregator: Addr<Aggregator>,
 }
 
 impl MetricStats {
     /// Creates a new [`MetricStats`] instance.
-    pub fn new(config: Arc<Config>, aggregator: Addr<Aggregator>) -> Self {
-        Self { config, aggregator }
+    pub fn new(
+        config: Arc<Config>,
+        global_config: GlobalConfigHandle,
+        aggregator: Addr<Aggregator>,
+    ) -> Self {
+        Self {
+            config,
+            global_config,
+            aggregator,
+        }
     }
 
     /// Tracks the metric volume and outcome for the bucket.
     pub fn track(&self, scoping: Scoping, bucket: &BucketView<'_>, outcome: Outcome) {
-        if !self.config.processing_enabled() {
+        if !self.config.processing_enabled() || !self.is_rolled_out(scoping.organization_id) {
             return;
         }
 
@@ -56,6 +66,16 @@ impl MetricStats {
             .send(MergeBuckets::new(scoping.project_key, vec![volume]));
     }
 
+    fn is_rolled_out(&self, organization_id: u64) -> bool {
+        let rate = self
+            .global_config
+            .current()
+            .options
+            .metric_stats_rollout_rate;
+
+        ((organization_id % 100000) as f32 / 100000.0f32) <= rate
+    }
+
     fn to_volume_metric(&self, bucket: &BucketView<'_>, outcome: &Outcome) -> Option<Bucket> {
         let volume = bucket.metadata().merges.get();
         if volume == 0 {
@@ -92,3 +112,138 @@ impl MetricStats {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use relay_base_schema::project::{ProjectId, ProjectKey};
+    use relay_dynamic_config::GlobalConfig;
+    use relay_quotas::ReasonCode;
+    use tokio::sync::mpsc::UnboundedReceiver;
+
+    use super::*;
+
+    fn create_metric_stats(rollout_rate: f32) -> (MetricStats, UnboundedReceiver<Aggregator>) {
+        let config = Config::from_json_value(serde_json::json!({
+            "processing": {
+                "enabled": true,
+                "kafka_config": [],
+            }
+        }))
+        .unwrap();
+
+        let mut global_config = GlobalConfig::default();
+        global_config.options.metric_stats_rollout_rate = rollout_rate;
+        let global_config = GlobalConfigHandle::fixed(global_config);
+
+        let (addr, receiver) = Addr::custom();
+        let ms = MetricStats::new(Arc::new(config), global_config, addr);
+
+        (ms, receiver)
+    }
+
+    fn scoping() -> Scoping {
+        Scoping {
+            organization_id: 42,
+            project_id: ProjectId::new(21),
+            project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
+            key_id: Some(17),
+        }
+    }
+
+    macro_rules! tags {
+        ($(($key:expr, $value:expr),)*) => {
+            BTreeMap::from([
+                $(($key.to_owned(), $value.to_owned())),*
+            ])
+        }
+    }
+
+    #[test]
+    fn test_metric_stats_volume() {
+        let (ms, mut receiver) = create_metric_stats(1.0);
+
+        let scoping = scoping();
+        let mut bucket = Bucket::parse(b"rt@millisecond:57|d", UnixTimestamp::now()).unwrap();
+
+        ms.track(scoping, &BucketView::from(&bucket), Outcome::Accepted);
+
+        bucket.metadata.merges = bucket.metadata.merges.saturating_add(41);
+        ms.track(
+            scoping,
+            &BucketView::from(&bucket),
+            Outcome::RateLimited(Some(ReasonCode::new("foobar"))),
+        );
+
+        drop(ms);
+
+        let Aggregator::MergeBuckets(mb) = receiver.blocking_recv().unwrap() else {
+            panic!();
+        };
+        assert_eq!(mb.project_key(), scoping.project_key);
+
+        let mut buckets = mb.buckets();
+        assert_eq!(buckets.len(), 1);
+        let bucket = buckets.pop().unwrap();
+
+        assert_eq!(&*bucket.name, "c:metric_stats/volume@none");
+        assert_eq!(bucket.value, BucketValue::Counter(1.into()));
+        assert_eq!(
+            bucket.tags,
+            tags!(
+                ("mri", "d:custom/rt@millisecond"),
+                ("mri.namespace", "custom"),
+                ("outcome.id", "0"),
+            )
+        );
+
+        let Aggregator::MergeBuckets(mb) = receiver.blocking_recv().unwrap() else {
+            panic!();
+        };
+        assert_eq!(mb.project_key(), scoping.project_key);
+
+        let mut buckets = mb.buckets();
+        assert_eq!(buckets.len(), 1);
+        let bucket = buckets.pop().unwrap();
+
+        assert_eq!(&*bucket.name, "c:metric_stats/volume@none");
+        assert_eq!(bucket.value, BucketValue::Counter(42.into()));
+        assert_eq!(
+            bucket.tags,
+            tags!(
+                ("mri", "d:custom/rt@millisecond"),
+                ("mri.namespace", "custom"),
+                ("outcome.id", "2"),
+                ("outcome.reason", "foobar"),
+            )
+        );
+
+        assert!(receiver.blocking_recv().is_none());
+    }
+
+    #[test]
+    fn test_metric_stats_rollout_rate_disabled() {
+        let (ms, mut receiver) = create_metric_stats(0.0);
+
+        let scoping = scoping();
+        let bucket = Bucket::parse(b"rt@millisecond:57|d", UnixTimestamp::now()).unwrap();
+        ms.track(scoping, &BucketView::from(&bucket), Outcome::Accepted);
+
+        drop(ms);
+
+        assert!(receiver.blocking_recv().is_none());
+    }
+
+    #[test]
+    fn test_metric_stats_disabled_namespace() {
+        let (ms, mut receiver) = create_metric_stats(1.0);
+
+        let scoping = scoping();
+        let bucket =
+            Bucket::parse(b"transactions/rt@millisecond:57|d", UnixTimestamp::now()).unwrap();
+        ms.track(scoping, &BucketView::from(&bucket), Outcome::Accepted);
+
+        drop(ms);
+
+        assert!(receiver.blocking_recv().is_none());
+    }
+}
@@ -137,7 +137,11 @@ impl ServiceState {
         )
         .start_in(&runtimes.aggregator);
 
-        let metric_stats = MetricStats::new(config.clone(), aggregator.clone());
+        let metric_stats = MetricStats::new(
+            config.clone(),
+            global_config_handle.clone(),
+            aggregator.clone(),
+        );
 
         #[cfg(feature = "processing")]
         let store = match runtimes.store {

@@ -1248,7 +1248,7 @@ fn is_metric_namespace_valid(state: &ProjectState, namespace: &MetricNamespace)
         }
         MetricNamespace::Profiles => true,
         MetricNamespace::Custom => state.has_feature(Feature::CustomMetrics),
-        MetricNamespace::Stats => state.has_feature(Feature::MetricStats),
+        MetricNamespace::Stats => true,
         MetricNamespace::Unsupported => false,
     }
 }

@@ -25,6 +25,8 @@ def metric_stats_by_mri(metrics_consumer, count, timeout=None):
 
 
 def test_metric_stats_simple(mini_sentry, relay_with_processing, metrics_consumer):
+    mini_sentry.global_config["options"]["relay.metric-stats.rollout-rate"] = 1.0
+
     metrics_consumer = metrics_consumer()
     relay = relay_with_processing(options=TEST_CONFIG)