feat(cogs): Track categories in COGS (#4496)

Makes it possible to split a COGS measurement into multiple categories.

Also instruments the transaction processing with more categories. The
initial plan was to instrument normalization in multiple smaller
categories, I decided against that in this PR because that requires an
API change with huge effects in tests. If we realize normalization takes
up a significant amount, we can revisit this.

Ref: getsentry/team-ingest#637

Author: Dav1dde

Cargo.lock

@@ -12,3 +12,6 @@ autobenches = false
 
 [lints]
 workspace = true
+
+[dev-dependencies]
+insta = { workspace = true }

relay-cogs/Cargo.toml

@@ -1,10 +1,10 @@
-use core::fmt;
+use crate::time::Instant;
 use std::collections::BTreeMap;
+use std::fmt;
 use std::num::NonZeroUsize;
 use std::sync::Arc;
-use std::time::Instant;
 
-use crate::{AppFeature, ResourceId};
+use crate::{AppFeature, Measurements, ResourceId};
 use crate::{CogsMeasurement, CogsRecorder, Value};
 
 /// COGS measurements collector.
@@ -64,8 +64,8 @@ impl Cogs {
         Token {
             resource,
             features: weights.into(),
-            start: Instant::now(),
-            recorder: Arc::clone(&self.recorder),
+            measurements: Measurements::start(),
+            recorder: Some(Arc::clone(&self.recorder)),
         }
     }
 }
@@ -77,17 +77,43 @@ impl Cogs {
 pub struct Token {
     resource: ResourceId,
     features: FeatureWeights,
-    start: Instant,
-    recorder: Arc<dyn CogsRecorder>,
+    measurements: Measurements,
+    recorder: Option<Arc<dyn CogsRecorder>>,
 }
 
 impl Token {
+    /// Creates a new no-op token, which records nothing.
+    ///
+    /// This is primarily useful for testing.
+    pub fn noop() -> Self {
+        Self {
+            resource: ResourceId::Relay,
+            features: FeatureWeights::none(),
+            measurements: Measurements::start(),
+            recorder: None,
+        }
+    }
+
     /// Cancels the COGS measurement.
     pub fn cancel(&mut self) {
         // No features -> nothing gets attributed.
         self.update(FeatureWeights::none());
     }
 
+    /// Starts a categorized measurement.
+    ///
+    /// The measurement is finalized when the returned [`CategoryToken`] is dropped.
+    ///
+    /// Instead of manually starting a categorized measurement, the [`crate::with`]
+    /// macro can be used.
+    pub fn start_category(&mut self, category: impl Category) -> CategoryToken<'_> {
+        CategoryToken {
+            parent: self,
+            start: Instant::now(),
+            category: category.name(),
+        }
+    }
+
     /// Updates the app features to which the active measurement is attributed to.
     ///
     /// # Example:
@@ -115,15 +141,22 @@ impl Token {
 
 impl Drop for Token {
     fn drop(&mut self) {
-        let elapsed = self.start.elapsed();
-
-        for (feature, ratio) in self.features.weights() {
-            let time = elapsed.mul_f32(ratio);
-            self.recorder.record(CogsMeasurement {
-                resource: self.resource,
-                feature,
-                value: Value::Time(time),
-            });
+        let Some(recorder) = self.recorder.as_mut() else {
+            return;
+        };
+
+        let measurements = self.measurements.finish();
+
+        for measurement in measurements {
+            for (feature, ratio) in self.features.weights() {
+                let time = measurement.duration.mul_f32(ratio);
+                recorder.record(CogsMeasurement {
+                    resource: self.resource,
+                    feature,
+                    category: measurement.category,
+                    value: Value::Time(time),
+                });
+            }
         }
     }
 }
@@ -137,6 +170,46 @@ impl fmt::Debug for Token {
     }
 }
 
+/// A COGS category.
+pub trait Category {
+    /// String representation of the category.
+    fn name(&self) -> &'static str;
+}
+
+impl Category for &'static str {
+    fn name(&self) -> &'static str {
+        self
+    }
+}
+
+/// A categorized COGS measurement.
+///
+/// Must be started with [`Token::start_category`].
+#[must_use]
+pub struct CategoryToken<'a> {
+    parent: &'a mut Token,
+    start: Instant,
+    category: &'static str,
+}
+
+impl Drop for CategoryToken<'_> {
+    fn drop(&mut self) {
+        self.parent
+            .measurements
+            .add(self.start.elapsed(), self.category);
+    }
+}
+
+impl fmt::Debug for CategoryToken<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("CategoryToken")
+            .field("resource", &self.parent.resource)
+            .field("features", &self.parent.features)
+            .field("category", &self.category)
+            .finish()
+    }
+}
+
 /// A collection of weighted [app features](AppFeature).
 ///
 /// Used to attribute a single COGS measurement to multiple features.
@@ -278,7 +351,7 @@ impl FeatureWeightsBuilder {
 
 #[cfg(test)]
 mod tests {
-    use std::{collections::HashMap, time::Duration};
+    use std::collections::HashMap;
 
     use super::*;
     use crate::test::TestRecorder;
@@ -291,17 +364,25 @@ mod tests {
         drop(cogs.timed(ResourceId::Relay, AppFeature::Spans));
 
         let measurements = recorder.measurements();
-        assert_eq!(measurements.len(), 1);
-        assert_eq!(measurements[0].resource, ResourceId::Relay);
-        assert_eq!(measurements[0].feature, AppFeature::Spans);
+        insta::assert_debug_snapshot!(measurements, @r###"
+        [
+            CogsMeasurement {
+                resource: Relay,
+                feature: Spans,
+                category: None,
+                value: Time(
+                    100ns,
+                ),
+            },
+        ]
+        "###);
     }
 
     #[test]
     fn test_cogs_multiple_weights() {
         let recorder = TestRecorder::default();
         let cogs = Cogs::new(recorder.clone());
 
-        let start = Instant::now();
         let f = FeatureWeights::builder()
             .weight(AppFeature::Spans, 1)
             .weight(AppFeature::Transactions, 1)
@@ -311,20 +392,114 @@ mod tests {
             .build();
         {
             let _token = cogs.timed(ResourceId::Relay, f);
-            std::thread::sleep(Duration::from_millis(50));
+            crate::time::advance_millis(50);
+        }
+
+        let measurements = recorder.measurements();
+        insta::assert_debug_snapshot!(measurements, @r###"
+        [
+            CogsMeasurement {
+                resource: Relay,
+                feature: Spans,
+                category: None,
+                value: Time(
+                    25ms,
+                ),
+            },
+            CogsMeasurement {
+                resource: Relay,
+                feature: MetricsSpans,
+                category: None,
+                value: Time(
+                    25ms,
+                ),
+            },
+        ]
+        "###);
+    }
+
+    #[test]
+    fn test_cogs_categorized() {
+        let recorder = TestRecorder::default();
+        let cogs = Cogs::new(recorder.clone());
+
+        let features = FeatureWeights::builder()
+            .weight(AppFeature::Spans, 1)
+            .weight(AppFeature::Errors, 1)
+            .build();
+
+        {
+            let mut token = cogs.timed(ResourceId::Relay, features);
+            crate::time::advance_millis(10);
+            crate::with!(token, "s1", {
+                crate::time::advance_millis(6);
+            });
+            crate::time::advance_millis(20);
+            let _category = token.start_category("s2");
+            crate::time::advance_millis(12);
         }
-        let elapsed = start.elapsed();
 
         let measurements = recorder.measurements();
-        assert_eq!(measurements.len(), 2);
-        assert_eq!(measurements[0].resource, ResourceId::Relay);
-        assert_eq!(measurements[0].feature, AppFeature::Spans);
-        assert_eq!(measurements[1].resource, ResourceId::Relay);
-        assert_eq!(measurements[1].feature, AppFeature::MetricsSpans);
-        assert_eq!(measurements[0].value, measurements[1].value);
-        let Value::Time(time) = measurements[0].value;
-        assert!(time >= Duration::from_millis(25), "{time:?}");
-        assert!(time <= elapsed.div_f32(1.99), "{time:?}");
+        insta::assert_debug_snapshot!(measurements, @r###"
+        [
+            CogsMeasurement {
+                resource: Relay,
+                feature: Errors,
+                category: None,
+                value: Time(
+                    15ms,
+                ),
+            },
+            CogsMeasurement {
+                resource: Relay,
+                feature: Spans,
+                category: None,
+                value: Time(
+                    15ms,
+                ),
+            },
+            CogsMeasurement {
+                resource: Relay,
+                feature: Errors,
+                category: Some(
+                    "s1",
+                ),
+                value: Time(
+                    3ms,
+                ),
+            },
+            CogsMeasurement {
+                resource: Relay,
+                feature: Spans,
+                category: Some(
+                    "s1",
+                ),
+                value: Time(
+                    3ms,
+                ),
+            },
+            CogsMeasurement {
+                resource: Relay,
+                feature: Errors,
+                category: Some(
+                    "s2",
+                ),
+                value: Time(
+                    6ms,
+                ),
+            },
+            CogsMeasurement {
+                resource: Relay,
+                feature: Spans,
+                category: Some(
+                    "s2",
+                ),
+                value: Time(
+                    6ms,
+                ),
+            },
+        ]
+        "###);
     }
 
     #[test]