feat(metric-stats): Emit accepted volume

getsentry · Mar 18, 2024 · 58982c8 · 58982c8
1 parent d991a13
commit 58982c8
Show file tree

Hide file tree

Showing 10 changed files with 242 additions and 36 deletions.
diff --git a/relay-base-schema/src/metrics/mri.rs b/relay-base-schema/src/metrics/mri.rs
@@ -136,26 +136,31 @@ impl MetricNamespace {
     /// Returns all namespaces/variants of this enum.
     pub fn all() -> [Self; 7] {
         [
-            MetricNamespace::Sessions,
-            MetricNamespace::Transactions,
-            MetricNamespace::Spans,
-            MetricNamespace::Profiles,
-            MetricNamespace::Custom,
-            MetricNamespace::Stats,
-            MetricNamespace::Unsupported,
+            Self::Sessions,
+            Self::Transactions,
+            Self::Spans,
+            Self::Profiles,
+            Self::Custom,
+            Self::Stats,
+            Self::Unsupported,
         ]
     }
 
+    /// Returns `true` if metric stats are enabled for this namespace.
+    pub fn has_metric_stats(&self) -> bool {
+        matches!(self, Self::Custom)
+    }
+
     /// Returns the string representation for this metric type.
     pub fn as_str(&self) -> &'static str {
         match self {
-            MetricNamespace::Sessions => "sessions",
-            MetricNamespace::Transactions => "transactions",
-            MetricNamespace::Spans => "spans",
-            MetricNamespace::Profiles => "profiles",
-            MetricNamespace::Custom => "custom",
-            MetricNamespace::Stats => "metric_stats",
-            MetricNamespace::Unsupported => "unsupported",
+            Self::Sessions => "sessions",
+            Self::Transactions => "transactions",
+            Self::Spans => "spans",
+            Self::Profiles => "profiles",
+            Self::Custom => "custom",
+            Self::Stats => "metric_stats",
+            Self::Unsupported => "unsupported",
         }
     }
 }

diff --git a/relay-dynamic-config/src/feature.rs b/relay-dynamic-config/src/feature.rs
@@ -45,6 +45,11 @@ 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.
     ///

diff --git a/relay-metrics/src/bucket.rs b/relay-metrics/src/bucket.rs
@@ -1,7 +1,7 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::hash::Hash;
 use std::iter::FusedIterator;
-use std::num::NonZeroU64;
+use std::num::NonZeroU32;
 use std::{fmt, mem};
 
 use hash32::{FnvHasher, Hasher as _};
@@ -754,7 +754,7 @@ pub struct BucketMetadata {
     ///
     /// For example: Merging two un-merged buckets will yield a total
     /// of `2` merges.
-    pub merges: NonZeroU64,
+    pub merges: NonZeroU32,
 }
 
 impl BucketMetadata {
@@ -763,14 +763,14 @@ impl BucketMetadata {
     /// The new metadata is initialized with `1` merge.
     pub fn new() -> Self {
         Self {
-            merges: NonZeroU64::MIN,
+            merges: NonZeroU32::MIN,
         }
     }
 
     /// Whether the metadata does not contain more information than the default.
     pub fn is_default(&self) -> bool {
         let Self { merges } = self;
-        *merges == NonZeroU64::MIN
+        *merges == NonZeroU32::MIN
     }
 
     /// Merges another metadata object into the current one.

diff --git a/relay-server/src/service.rs b/relay-server/src/service.rs
@@ -16,6 +16,7 @@ use tokio::runtime::Runtime;
 use crate::services::cogs::{CogsService, CogsServiceRecorder};
 use crate::services::global_config::{GlobalConfigManager, GlobalConfigService};
 use crate::services::health_check::{HealthCheck, HealthCheckService};
+use crate::services::metric_stats::{MetricStatsService, TrackMetric};
 use crate::services::outcome::{OutcomeProducer, OutcomeProducerService, TrackOutcome};
 use crate::services::outcome_aggregator::OutcomeAggregator;
 use crate::services::processor::{EnvelopeProcessor, EnvelopeProcessorService};
@@ -50,6 +51,7 @@ pub struct Registry {
     pub health_check: Addr<HealthCheck>,
     pub outcome_producer: Addr<OutcomeProducer>,
     pub outcome_aggregator: Addr<TrackOutcome>,
+    pub metric_stats: Addr<TrackMetric>,
     pub processor: Addr<EnvelopeProcessor>,
     pub test_store: Addr<TestStore>,
     pub relay_cache: Addr<RelayCache>,
@@ -136,13 +138,16 @@ impl ServiceState {
         )
         .start_in(&runtimes.aggregator);
 
+        let metric_stats = MetricStatsService::new(config.clone(), aggregator.clone()).start();
+
         #[cfg(feature = "processing")]
         let store = match runtimes.store {
             Some(ref rt) => Some(
                 StoreService::create(
                     config.clone(),
                     global_config_handle.clone(),
                     outcome_aggregator.clone(),
+                    metric_stats.clone(),
                 )?
                 .start_in(rt),
             ),
@@ -214,6 +219,7 @@ impl ServiceState {
             health_check,
             outcome_producer,
             outcome_aggregator,
+            metric_stats,
             test_store,
             relay_cache,
             global_config,

diff --git a/relay-server/src/services/metric_stats.rs b/relay-server/src/services/metric_stats.rs
@@ -0,0 +1,114 @@
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use relay_config::Config;
+use relay_metrics::{
+    Aggregator, Bucket, BucketValue, BucketView, MergeBuckets, MetricResourceIdentifier,
+    UnixTimestamp,
+};
+use relay_quotas::Scoping;
+use relay_system::{Addr, FromMessage, Interface, NoResponse, Service};
+
+use crate::services::outcome::Outcome;
+
+#[derive(Debug)]
+pub struct TrackMetric {
+    pub timestamp: UnixTimestamp,
+    pub scoping: Scoping,
+    pub mri: String,
+    pub outcome: Outcome,
+    pub volume: u32,
+}
+
+impl TrackMetric {
+    pub fn volume(scoping: Scoping, bucket: &BucketView<'_>, outcome: Outcome) -> Self {
+        Self {
+            timestamp: UnixTimestamp::now(),
+            scoping,
+            mri: bucket.name().to_owned(),
+            outcome,
+            volume: bucket.metadata().merges.get(),
+        }
+    }
+
+    fn to_volume_metric(&self) -> Option<Bucket> {
+        if self.volume == 0 {
+            return None;
+        }
+
+        let namespace = MetricResourceIdentifier::parse(&self.mri).ok()?.namespace;
+        if !namespace.has_metric_stats() {
+            return None;
+        }
+
+        let mut tags = BTreeMap::from([
+            ("mri".to_owned(), self.mri.to_owned()),
+            ("mri.namespace".to_owned(), namespace.to_string()),
+            (
+                "outcome.id".to_owned(),
+                self.outcome.to_outcome_id().as_u8().to_string(),
+            ),
+        ]);
+
+        if let Some(reason) = self.outcome.to_reason() {
+            tags.insert("outcome.reason".to_owned(), reason.into_owned());
+        }
+
+        Some(Bucket {
+            timestamp: self.timestamp,
+            width: 0,
+            name: "c:metric_stats/volume@none".to_owned(),
+            value: BucketValue::Counter(self.volume.into()),
+            tags,
+            metadata: Default::default(),
+        })
+    }
+}
+
+impl Interface for TrackMetric {}
+
+impl FromMessage<Self> for TrackMetric {
+    type Response = NoResponse;
+
+    fn from_message(message: Self, _: ()) -> Self {
+        message
+    }
+}
+
+pub struct MetricStatsService {
+    config: Arc<Config>,
+    aggregator: Addr<Aggregator>,
+}
+
+impl MetricStatsService {
+    pub fn new(config: Arc<Config>, aggregator: Addr<Aggregator>) -> Self {
+        Self { config, aggregator }
+    }
+
+    fn handle_track_metric(&self, tm: TrackMetric) {
+        if !self.config.processing_enabled() {
+            return;
+        }
+
+        let Some(volume) = tm.to_volume_metric() else {
+            return;
+        };
+
+        relay_log::trace!("Tracking volume of {} for mri '{}'", tm.volume, tm.mri);
+        self.aggregator
+            .send(MergeBuckets::new(tm.scoping.project_key, vec![volume]));
+    }
+}
+
+impl Service for MetricStatsService {
+    type Interface = TrackMetric;
+
+    fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) {
+        tokio::spawn(async move {
+            while let Some(message) = rx.recv().await {
+                self.handle_track_metric(message);
+            }
+            dbg!("END");
+        });
+    }
+}
diff --git a/relay-server/src/services/mod.rs b/relay-server/src/services/mod.rs
@@ -30,6 +30,7 @@
 pub mod cogs;
 pub mod global_config;
 pub mod health_check;
+pub mod metric_stats;
 pub mod outcome;
 pub mod outcome_aggregator;
 pub mod processor;

diff --git a/relay-server/src/services/outcome.rs b/relay-server/src/services/outcome.rs
@@ -81,6 +81,10 @@ impl OutcomeId {
     const INVALID: OutcomeId = OutcomeId(3);
     const ABUSE: OutcomeId = OutcomeId(4);
     const CLIENT_DISCARD: OutcomeId = OutcomeId(5);
+
+    pub fn as_u8(self) -> u8 {
+        self.0
+    }
 }
 
 trait TrackOutcomeLike {
@@ -180,7 +184,7 @@ pub enum Outcome {
 
 impl Outcome {
     /// Returns the raw numeric value of this outcome for the JSON and Kafka schema.
-    fn to_outcome_id(&self) -> OutcomeId {
+    pub fn to_outcome_id(&self) -> OutcomeId {
         match self {
             Outcome::Filtered(_) | Outcome::FilteredSampling(_) => OutcomeId::FILTERED,
             Outcome::RateLimited(_) => OutcomeId::RATE_LIMITED,
@@ -194,7 +198,7 @@ impl Outcome {
     }
 
     /// Returns the `reason` code field of this outcome.
-    fn to_reason(&self) -> Option<Cow<str>> {
+    pub fn to_reason(&self) -> Option<Cow<str>> {
         match self {
             Outcome::Invalid(discard_reason) => Some(Cow::Borrowed(discard_reason.name())),
             Outcome::Filtered(filter_key) => Some(filter_key.clone().name()),

diff --git a/relay-server/src/services/project.rs b/relay-server/src/services/project.rs
@@ -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 => false,
+        MetricNamespace::Stats => state.has_feature(Feature::MetricStats),
         MetricNamespace::Unsupported => false,
     }
 }

diff --git a/relay-server/src/services/store.rs b/relay-server/src/services/store.rs
@@ -32,6 +32,7 @@ use uuid::Uuid;
 
 use crate::envelope::{AttachmentType, Envelope, Item, ItemType, SourceQuantities};
 use crate::services::global_config::GlobalConfigHandle;
+use crate::services::metric_stats::TrackMetric;
 use crate::services::outcome::{DiscardReason, Outcome, TrackOutcome};
 use crate::services::processor::Processed;
 use crate::statsd::RelayCounters;
@@ -140,6 +141,7 @@ pub struct StoreService {
     config: Arc<Config>,
     global_config: GlobalConfigHandle,
     outcome_aggregator: Addr<TrackOutcome>,
+    metric_stats: Addr<TrackMetric>,
     producer: Producer,
 }
 
@@ -148,12 +150,14 @@ impl StoreService {
         config: Arc<Config>,
         global_config: GlobalConfigHandle,
         outcome_aggregator: Addr<TrackOutcome>,
+        metric_stats: Addr<TrackMetric>,
     ) -> anyhow::Result<Self> {
         let producer = Producer::create(&config)?;
         Ok(Self {
             config,
             global_config,
             outcome_aggregator,
+            metric_stats,
             producer,
         })
     }
@@ -250,12 +254,7 @@ impl StoreService {
                         item,
                     )?;
                 }
-                ItemType::MetricBuckets => self.produce_metrics(
-                    scoping.organization_id,
-                    scoping.project_id,
-                    item,
-                    retention,
-                )?,
+                ItemType::MetricBuckets => self.produce_metrics(scoping, item, retention)?,
                 ItemType::Profile => self.produce_profile(
                     scoping.organization_id,
                     scoping.project_id,
@@ -376,11 +375,21 @@ impl StoreService {
                     retention,
                 );
 
-                if let Err(e) =
-                    message.and_then(|message| self.send_metric_message(namespace, message))
-                {
-                    error.get_or_insert(e);
-                    dropped += utils::extract_metric_quantities([view], mode);
+                let result =
+                    message.and_then(|message| self.send_metric_message(namespace, message));
+
+                match result {
+                    Ok(()) => {
+                        self.metric_stats.send(TrackMetric::volume(
+                            scoping,
+                            &view,
+                            Outcome::Accepted,
+                        ));
+                    }
+                    Err(e) => {
+                        error.get_or_insert(e);
+                        dropped += utils::extract_metric_quantities([view], mode);
+                    }
                 }
             }
         }
@@ -806,8 +815,7 @@ impl StoreService {
 
     fn produce_metrics(
         &self,
-        org_id: u64,
-        project_id: ProjectId,
+        scoping: Scoping,
         item: &Item,
         retention: u16,
     ) -> Result<(), StoreError> {
@@ -821,14 +829,18 @@ impl StoreService {
 
             let view = BucketView::new(&bucket);
             let message = self.create_metric_message(
-                org_id,
-                project_id,
+                scoping.organization_id,
+                scoping.project_id,
                 &mut encoder,
                 namespace,
                 &view,
                 retention,
             );
+
             message.and_then(|message| self.send_metric_message(namespace, message))?;
+
+            self.metric_stats
+                .send(TrackMetric::volume(scoping, &view, Outcome::Accepted));
         }
 
         Ok(())