From 72b5d37cf9a00a004f851ec7c1e0f8cecdbc0fc0 Mon Sep 17 00:00:00 2001
From: colin-sentry <161344340+colin-sentry@users.noreply.github.com>
Date: Tue, 30 Jul 2024 15:22:32 -0400
Subject: [PATCH] Send client sample rate as a span measurement (#3872)

This will be consumed by snuba spans consumer
---
 CHANGELOG.md                                  |  2 +
 .../src/normalize/span/tag_extraction.rs      | 52 +++++++++++++++++--
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 482b8084d5..3bd86a9957 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,8 @@
 **Internal**:
 
 - Add `EnvelopeStack` and `SQLiteEnvelopeStack` to manage envelopes on disk. ([#3855](https://github.com/getsentry/relay/pull/3855))
+- Add `client_sample_rate` to spans, pulled from the trace context ([#3872](https://github.com/getsentry/relay/pull/3872)).
+
 
 ## 24.7.1
 
diff --git a/relay-event-normalization/src/normalize/span/tag_extraction.rs b/relay-event-normalization/src/normalize/span/tag_extraction.rs
index 7e987ea400..ef3e1e0676 100644
--- a/relay-event-normalization/src/normalize/span/tag_extraction.rs
+++ b/relay-event-normalization/src/normalize/span/tag_extraction.rs
@@ -8,12 +8,12 @@ use std::ops::ControlFlow;
 
 use once_cell::sync::Lazy;
 use regex::Regex;
-use relay_base_schema::metrics::{DurationUnit, InformationUnit, MetricUnit};
+use relay_base_schema::metrics::{DurationUnit, FractionUnit, InformationUnit, MetricUnit};
 use relay_event_schema::protocol::{
-    AppContext, BrowserContext, Event, Measurement, OsContext, ProfileContext, Span, Timestamp,
-    TraceContext,
+    AppContext, BrowserContext, Event, Measurement, Measurements, OsContext, ProfileContext, Span,
+    Timestamp, TraceContext,
 };
-use relay_protocol::{Annotated, Value};
+use relay_protocol::{Annotated, Empty, Value};
 use sqlparser::ast::Visit;
 use sqlparser::ast::{ObjectName, Visitor};
 use url::Url;
@@ -194,6 +194,7 @@ pub fn extract_span_tags(event: &Event, spans: &mut [Annotated<Span>], max_tag_v
     // TODO: To prevent differences between metrics and payloads, we should not extract tags here
     // when they have already been extracted by a downstream relay.
     let shared_tags = extract_shared_tags(event);
+    let shared_measurements = extract_shared_measurements(event);
     let is_mobile = shared_tags
         .get(&SpanTagKey::Mobile)
         .is_some_and(|v| v.as_str() == "true");
@@ -218,6 +219,17 @@ pub fn extract_span_tags(event: &Event, spans: &mut [Annotated<Span>], max_tag_v
                 .collect(),
         );
 
+        if !shared_measurements.is_empty() {
+            match span.measurements.value_mut() {
+                Some(left) => shared_measurements.iter().for_each(|(key, val)| {
+                    left.insert(key.clone(), val.clone());
+                }),
+                None => span
+                    .measurements
+                    .set_value(Some(shared_measurements.clone())),
+            }
+        }
+
         extract_measurements(span, is_mobile);
     }
 }
@@ -794,6 +806,27 @@ fn value_to_f64(val: Option<&Value>) -> Option<f64> {
     }
 }
 
+fn extract_shared_measurements(event: &Event) -> Measurements {
+    let mut measurements = Measurements::default();
+
+    if let Some(trace_context) = event.context::<TraceContext>() {
+        if let Some(client_sample_rate) = trace_context.client_sample_rate.value() {
+            if *client_sample_rate > 0. {
+                measurements.insert(
+                    "client_sample_rate".into(),
+                    Measurement {
+                        value: (*client_sample_rate).into(),
+                        unit: MetricUnit::Fraction(FractionUnit::Ratio).into(),
+                    }
+                    .into(),
+                );
+            }
+        }
+    }
+
+    measurements
+}
+
 /// Copies specific numeric values from span data to span measurements.
 pub fn extract_measurements(span: &mut Span, is_mobile: bool) {
     let Some(span_op) = span.op.as_str() else {
@@ -1623,6 +1656,11 @@ LIMIT 1
     fn test_ai_extraction() {
         let json = r#"
             {
+                "contexts": {
+                    "trace": {
+                        "client_sample_rate": 0.1
+                    }
+                },
                 "spans": [
                     {
                         "timestamp": 1694732408.3145,
@@ -1683,6 +1721,12 @@ LIMIT 1
                     value: 300.0,
                     unit: None,
                 },
+                "client_sample_rate": Measurement {
+                    value: 0.1,
+                    unit: Fraction(
+                        Ratio,
+                    ),
+                },
             },
         )
         "###);