feat(outcomes): Have relay generate metric billing outcomes

Makefile

@@ -75,6 +75,12 @@ test-integration: build setup-venv ## run integration tests
 	.venv/bin/pytest tests -n $(PYTEST_N) -v
 .PHONY: test-integration
 
+PYTEST_N ?= auto
+TEST ?= tests
+test-custom: build setup-venv ## run arbitrary tests
+	.venv/bin/pytest -n $(PYTEST_N) $(TEST) -vv
+.PHONY: test-custom
+
 # Documentation
 
 doc: doc-rust ## generate all API docs

relay-dynamic-config/src/feature.rs

@@ -142,6 +142,10 @@ pub enum Feature {
     #[serde(rename = "organizations:standalone-span-ingestion")]
     DeprecatedStandaloneSpanIngestion,
 
+    /// Enable relay billing outcome generation.
+    #[serde(rename = "organizations:relay-generate-billing-outcome")]
+    GenerateBillingOutcome,
+
     /// Forward compatibility.
     #[doc(hidden)]
     #[serde(other)]

relay-server/src/metrics/outcomes.rs

@@ -58,6 +58,43 @@ impl MetricOutcomes {
             }
         }
     }
+
+    /// Tracks billing-related outcomes for the list of buckets, adding the
+    /// "billing_outcome_accepted" tag to the bucket if that bucket is accepted.
+    pub fn track_billing_outcome(&self, scoping: Scoping, buckets: &mut [Bucket]) {
+        let timestamp = Utc::now();
+        for bucket in buckets {
+            let summary = bucket.summary();
+            match summary {
+                BucketSummary::Spans {
+                    count,
+                    is_segment,
+                    was_transaction: _,
+                } if count > 0 => {
+                    let all_categories = [DataCategory::Span, DataCategory::Transaction];
+                    let num_categories = if is_segment { 2 } else { 1 };
+                    let categories = &all_categories[0..num_categories];
+
+                    bucket
+                        .tags
+                        .insert("billing_outcome_accepted".to_owned(), "true".to_owned());
+
+                    for category in categories {
+                        self.outcomes.send(TrackOutcome {
+                            timestamp,
+                            scoping,
+                            outcome: Outcome::Accepted,
+                            event_id: None,
+                            remote_addr: None,
+                            category: *category,
+                            quantity: count as u32,
+                        });
+                    }
+                }
+                _ => continue,
+            };
+        }
+    }
 }
 
 /// The return value of [`TrackableBucket::summary`].

relay-server/src/services/processor.rs

@@ -19,6 +19,7 @@ use relay_base_schema::project::{ProjectId, ProjectKey};
 use relay_cogs::{AppFeature, Cogs, FeatureWeights, ResourceId, Token};
 use relay_common::time::UnixTimestamp;
 use relay_config::{Config, HttpEncoding, UpstreamDescriptor};
+use relay_dynamic_config::Feature;
 use relay_event_normalization::{ClockDriftProcessor, GeoIpLookup};
 use relay_event_schema::processor::ProcessingAction;
 use relay_event_schema::protocol::ClientReport;
@@ -1171,6 +1172,7 @@ impl EnvelopeProcessorService {
     ///
     /// This function runs the following steps:
     ///  - rate limiting
+    ///  - emit billing outcomes
     ///  - submit to `StoreForwarder`
     #[cfg(feature = "processing")]
     async fn encode_metrics_processing(
@@ -1188,14 +1190,25 @@ impl EnvelopeProcessorService {
             ..
         } in message.buckets.into_values()
         {
-            let buckets = self
+            let mut buckets = self
                 .rate_limit_buckets(scoping, &project_info, buckets)
                 .await;
 
             if buckets.is_empty() {
                 continue;
             }
 
+            if project_info
+                .config
+                .features
+                .has(Feature::GenerateBillingOutcome)
+            {
+                // Emit metric billing outcomes.
+                self.inner
+                    .metric_outcomes
+                    .track_billing_outcome(scoping, &mut buckets);
+            }
+
             let retention = project_info
                 .config
                 .event_retention

tests/integration/test_ai.py

@@ -32,8 +32,10 @@ def test_ai_spans_example_transaction(
     outcomes_consumer = outcomes_consumer()
 
     project_id = 42
-    mini_sentry.add_full_project_config(project_id)
-
+    project = mini_sentry.add_full_project_config(project_id)
+    project["config"].setdefault("features", []).extend(
+        ["organizations:relay-generate-billing-outcome"]
+    )
     mini_sentry.global_config["aiModelMetadata"] = {
         "version": 1,
         "models": {
@@ -1298,12 +1300,47 @@ def test_ai_spans_example_transaction(
 
     if relay_emits_accepted_outcome:
         assert outcomes_consumer.get_aggregated_outcomes() == [
+            {
+                "category": DataCategory.TRANSACTION.value,
+                "key_id": 123,
+                "org_id": 1,
+                "outcome": 0,
+                "project_id": 42,
+                "quantity": 1,
+            },
+            {
+                "category": DataCategory.SPAN.value,
+                "key_id": 123,
+                "org_id": 1,
+                "outcome": 0,
+                "project_id": 42,
+                "quantity": 10,
+            },
             {
                 "category": DataCategory.SPAN_INDEXED.value,
                 "key_id": 123,
                 "org_id": 1,
                 "outcome": 0,
                 "project_id": 42,
                 "quantity": 10,
-            }
+            },
+        ]
+    else:
+        assert outcomes_consumer.get_aggregated_outcomes() == [
+            {
+                "category": DataCategory.TRANSACTION.value,
+                "key_id": 123,
+                "org_id": 1,
+                "outcome": 0,
+                "project_id": 42,
+                "quantity": 1,
+            },
+            {
+                "category": DataCategory.SPAN.value,
+                "key_id": 123,
+                "org_id": 1,
+                "outcome": 0,
+                "project_id": 42,
+                "quantity": 10,
+            },
         ]

tests/integration/test_attachment_ref.py

@@ -254,8 +254,11 @@ def test_attachment_ref_validation(
     event_id = "515539018c9b4260a6f999572f1661ee"
     project_id = 42
     project_config = mini_sentry.add_full_project_config(project_id)
-    project_config["config"].setdefault("features", []).append(
-        "projects:relay-upload-endpoint"
+    project_config["config"].setdefault("features", []).extend(
+        [
+            "projects:relay-upload-endpoint",
+            "organizations:relay-generate-billing-outcome",
+        ]
     )
 
     relay = relay_with_processing()
@@ -275,7 +278,7 @@ def test_attachment_ref_validation(
 
     relay.send_envelope(project_id, envelope)
 
-    outcomes = outcomes_consumer.get_outcomes(n=3 if event_type == "transaction" else 2)
+    outcomes = outcomes_consumer.get_outcomes(n=5 if event_type == "transaction" else 2)
     o = {DataCategory(o["category"]): o for o in outcomes}
     assert o[DataCategory.ATTACHMENT]["reason"] == "invalid_placeholder_attachment"
     assert o[DataCategory.ATTACHMENT]["quantity"] == expected_bytes_quantity

tests/integration/test_attachmentsv2.py

@@ -353,6 +353,7 @@ def test_attachment_with_matching_span_store(
     project_config["config"]["features"] = [
         "projects:span-v2-experimental-processing",
         "projects:span-v2-attachment-processing",
+        "organizations:relay-generate-billing-outcome",
     ]
     relay = relay_with_processing()
 
@@ -424,8 +425,16 @@ def test_attachment_with_matching_span_store(
     objectstore = objectstore(usecase="trace_attachments", project_id=project_id)
     assert objectstore.get(metadata["attachment_id"]).payload.read() == body
 
-    outcomes = outcomes_consumer.get_aggregated_outcomes(n=3)
+    outcomes = outcomes_consumer.get_aggregated_outcomes(n=5)
     assert outcomes == [
+        {
+            "category": DataCategory.TRANSACTION.value,
+            "key_id": 123,
+            "org_id": 1,
+            "outcome": 0,
+            "project_id": 42,
+            "quantity": 1,
+        },
         {
             "category": DataCategory.ATTACHMENT.value,
             "key_id": 123,
@@ -434,6 +443,14 @@ def test_attachment_with_matching_span_store(
             "project_id": 42,
             "quantity": 23,
         },
+        {
+            "category": DataCategory.SPAN.value,
+            "key_id": 123,
+            "org_id": 1,
+            "outcome": 0,
+            "project_id": 42,
+            "quantity": 1,
+        },
         {
             "category": DataCategory.SPAN_INDEXED.value,
             "key_id": 123,

tests/integration/test_otlp_logs.py

@@ -187,6 +187,7 @@ def test_otlp_logs_multiple_records(
     project_config = mini_sentry.add_full_project_config(project_id)
     project_config["config"]["features"] = [
         "organizations:ourlogs-ingestion",
+        "organizations:relay-generate-billing-outcome",
     ]
     project_config["config"]["retentions"] = {
         "log": {"standard": 30, "downsampled": 13 * 30},