Merge branch 'master' into ci/leaner-snuba-image

Author: aldy505

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
     "sentry-arroyo>=2.33.0",
     "sentry-conventions>=0.3.0",
     "sentry-kafka-schemas>=2.1.2",
-    "sentry-protos>=0.4.3",
+    "sentry-protos>=0.4.7",
     "sentry-redis-tools>=0.5.1",
     "sentry-relay>=0.9.19",
     "sentry-sdk>=2.35.0",

rust_snuba/Cargo.lock

@@ -42,7 +42,7 @@ reqwest = { version = "0.11.11", features = ["stream"] }
 schemars = { version = "0.8.16", features = ["uuid1"] }
 sentry = { version = "0.32.0", features = ["anyhow", "tracing"] }
 sentry-kafka-schemas = "2.0.2"
-sentry_protos = "0.4.3"
+sentry_protos = "0.4.7"
 sentry_arroyo = { version = "2.33.0", features = ["ssl"] }
 sentry_usage_accountant = { version = "0.1.0", features = ["kafka"] }
 seq-macro = "0.3"

rust_snuba/Cargo.toml

@@ -13,7 +13,7 @@ use sentry_protos::snuba::v1::{ArrayValue, TraceItem};
 
 use crate::config::ProcessorConfig;
 use crate::processors::utils::enforce_retention;
-use crate::types::{InsertBatch, KafkaMessageMetadata};
+use crate::types::{InsertBatch, ItemTypeMetrics, KafkaMessageMetadata};
 
 use crate::runtime_config::get_str_config;
 
@@ -36,6 +36,10 @@ pub fn process_message(
     } else {
         retention_days
     };
+
+    // Capture the item_type before consuming trace_item
+    let item_type = trace_item.item_type as u8;
+
     let mut eap_item = EAPItem::try_from(trace_item)?;
 
     eap_item.retention_days = retention_days;
@@ -45,7 +49,12 @@ pub fn process_message(
         Utc::now().timestamp_millis(),
     );
 
-    InsertBatch::from_rows([eap_item], origin_timestamp)
+    let mut item_type_metrics = ItemTypeMetrics::new();
+    item_type_metrics.record_item(item_type);
+
+    let mut batch = InsertBatch::from_rows([eap_item], origin_timestamp)?;
+    batch.item_type_metrics = Some(item_type_metrics);
+    Ok(batch)
 }
 
 #[derive(Debug, Default, Serialize)]
@@ -346,6 +355,88 @@ mod tests {
         assert_eq!(item.downsampled_retention_days, 365);
     }
 
+    #[test]
+    fn test_item_type_metrics_accumulated() {
+        let item_id = Uuid::new_v4();
+        let mut trace_item = generate_trace_item(item_id);
+
+        // Set the item type to Span (value 1)
+        trace_item.item_type = TraceItemType::Span.into();
+
+        let mut payload = Vec::new();
+        trace_item.encode(&mut payload).unwrap();
+
+        let payload = KafkaPayload::new(None, None, Some(payload));
+        let meta = KafkaMessageMetadata {
+            partition: 0,
+            offset: 1,
+            timestamp: DateTime::from(SystemTime::now()),
+        };
+
+        let batch = process_message(payload, meta, &ProcessorConfig::default())
+            .expect("The message should be processed");
+
+        // Verify that item_type_metrics is populated
+        assert!(batch.item_type_metrics.is_some());
+
+        let metrics = batch.item_type_metrics.unwrap();
+
+        // Verify that the item_type (Span = 1) has a count of 1
+        assert_eq!(metrics.counts.len(), 1);
+        assert_eq!(metrics.counts.get(&(TraceItemType::Span as u8)), Some(&1));
+    }
+
+    #[test]
+    fn test_item_type_metrics_different_types() {
+        // Process first message with item type 1 (Span)
+        let item_id_1 = Uuid::new_v4();
+        let mut trace_item_1 = generate_trace_item(item_id_1);
+        trace_item_1.item_type = TraceItemType::Span.into();
+
+        let mut payload_1 = Vec::new();
+        trace_item_1.encode(&mut payload_1).unwrap();
+
+        let payload_1 = KafkaPayload::new(None, None, Some(payload_1));
+        let meta_1 = KafkaMessageMetadata {
+            partition: 0,
+            offset: 1,
+            timestamp: DateTime::from(SystemTime::now()),
+        };
+
+        let batch_1 = process_message(payload_1, meta_1, &ProcessorConfig::default())
+            .expect("The message should be processed");
+
+        // Process second message with item type 2 (Transaction)
+        let item_id_2 = Uuid::new_v4();
+        let mut trace_item_2 = generate_trace_item(item_id_2);
+        trace_item_2.item_type = 2; // Transaction type
+
+        let mut payload_2 = Vec::new();
+        trace_item_2.encode(&mut payload_2).unwrap();
+
+        let payload_2 = KafkaPayload::new(None, None, Some(payload_2));
+        let meta_2 = KafkaMessageMetadata {
+            partition: 0,
+            offset: 2,
+            timestamp: DateTime::from(SystemTime::now()),
+        };
+
+        let batch_2 = process_message(payload_2, meta_2, &ProcessorConfig::default())
+            .expect("The message should be processed");
+
+        // Merge the metrics as would happen in the pipeline
+        let mut merged_metrics = batch_1.item_type_metrics.unwrap();
+        merged_metrics.merge(batch_2.item_type_metrics.unwrap());
+
+        // Verify that both item types are present with count 1 each
+        assert_eq!(merged_metrics.counts.len(), 2);
+        assert_eq!(
+            merged_metrics.counts.get(&(TraceItemType::Span as u8)),
+            Some(&1)
+        );
+        assert_eq!(merged_metrics.counts.get(&2u8), Some(&1));
+    }
+
     #[test]
     fn test_insert_arrays() {
         let item_id = Uuid::new_v4();

rust_snuba/src/processors/eap_items.rs

@@ -72,6 +72,7 @@ pub fn process_message_with_replacement(
                 rows: RowData::from_rows([row])?,
                 sentry_received_timestamp: None,
                 cogs_data: None,
+                item_type_metrics: None,
             }))
         }
         ("insert", None, _) => {