fix(server): Reorder DSC validation for correct client sample rate

CHANGELOG.md

@@ -12,6 +12,7 @@
 **Bug Fixes**:
 
 - Separates profiles into backend and ui profiles. ([#4595](https://github.com/getsentry/relay/pull/4595))
+- Normalize trace context information before writing it into transaction and span data. This ensures the correct sampling rates are stored for extrapolation in Sentry. ([#4625](https://github.com/getsentry/relay/pull/4625))
 
 **Internal**:
 

relay-server/src/envelope.rs

@@ -1363,6 +1363,11 @@ impl Envelope {
         self.headers.trace = Some(ErrorBoundary::Ok(dsc));
     }
 
+    /// Removes the dynamic sampling context from envelope headers.
+    pub fn remove_dsc(&mut self) {
+        self.headers.trace = None;
+    }
+
     /// Features required to process this envelope.
     pub fn required_features(&self) -> &[Feature] {
         &self.headers.required_features

relay-server/src/services/processor.rs

@@ -1560,7 +1560,7 @@ impl EnvelopeProcessorService {
         managed_envelope: &mut TypedEnvelope<ErrorGroup>,
         project_id: ProjectId,
         project_info: Arc<ProjectInfo>,
-        sampling_project_info: Option<Arc<ProjectInfo>>,
+        mut sampling_project_info: Option<Arc<ProjectInfo>>,
         rate_limits: Arc<RateLimits>,
     ) -> Result<Option<ProcessingExtractedMetrics>, ProcessingError> {
         let mut event_fully_normalized = EventFullyNormalized::new(managed_envelope.envelope());
@@ -1601,6 +1601,12 @@ impl EnvelopeProcessorService {
             }
         });
 
+        sampling_project_info = dynamic_sampling::validate_and_set_dsc(
+            managed_envelope,
+            &mut event,
+            project_info.clone(),
+            sampling_project_info,
+        );
         event::finalize(
             managed_envelope,
             &mut event,
@@ -1722,6 +1728,15 @@ impl EnvelopeProcessorService {
             profile::transfer_id(&mut event, profile_id);
         });
 
+        relay_cogs::with!(cogs, "dynamic_sampling_dsc", {
+            sampling_project_info = dynamic_sampling::validate_and_set_dsc(
+                managed_envelope,
+                &mut event,
+                project_info.clone(),
+                sampling_project_info,
+            );
+        });
+
         relay_cogs::with!(cogs, "event_finalize", {
             event::finalize(
                 managed_envelope,
@@ -1741,15 +1756,6 @@ impl EnvelopeProcessorService {
             )?;
         });
 
-        relay_cogs::with!(cogs, "dynamic_sampling_dsc", {
-            sampling_project_info = dynamic_sampling::validate_and_set_dsc(
-                managed_envelope,
-                &mut event,
-                project_info.clone(),
-                sampling_project_info.clone(),
-            );
-        });
-
         relay_cogs::with!(cogs, "filter", {
             let filter_run = event::filter(
                 managed_envelope,
@@ -3977,6 +3983,8 @@ mod tests {
     #[tokio::test]
     #[cfg(feature = "processing")]
     async fn test_materialize_dsc() {
+        use crate::services::projects::project::PublicKeyConfig;
+
         let dsn = "https://e12d836b15bb49d7bbf99e64295d995b:@sentry.io/42"
             .parse()
             .unwrap();
@@ -4002,11 +4010,18 @@ mod tests {
             ProcessingGroup::Error,
         );
 
+        let mut project_info = ProjectInfo::default();
+        project_info.public_keys.push(PublicKeyConfig {
+            public_key: ProjectKey::parse("e12d836b15bb49d7bbf99e64295d995b").unwrap(),
+            numeric_id: Some(1),
+        });
+        let project_info = Arc::new(project_info);
+
         let process_message = ProcessEnvelope {
             envelope: managed_envelope,
-            project_info: Arc::new(ProjectInfo::default()),
+            project_info: project_info.clone(),
             rate_limits: Default::default(),
-            sampling_project_info: None,
+            sampling_project_info: Some(project_info),
             reservoir_counters: ReservoirCounters::default(),
         };
 

relay-server/src/services/processor/dynamic_sampling.rs

@@ -44,31 +44,36 @@ use crate::utils::{self, SamplingResult, TypedEnvelope};
 /// The function will return the sampling project information of the root project for the event. If
 /// no sampling project information is specified, the project information of the event’s project
 /// will be returned.
-pub fn validate_and_set_dsc(
-    managed_envelope: &mut TypedEnvelope<TransactionGroup>,
+pub fn validate_and_set_dsc<T>(
+    managed_envelope: &mut TypedEnvelope<T>,
     event: &mut Annotated<Event>,
     project_info: Arc<ProjectInfo>,
     sampling_project_info: Option<Arc<ProjectInfo>>,
 ) -> Option<Arc<ProjectInfo>> {
-    if managed_envelope.envelope().dsc().is_some() && sampling_project_info.is_some() {
+    let original_dsc = managed_envelope.envelope().dsc();
+    if original_dsc.is_some() && sampling_project_info.is_some() {
         return sampling_project_info;
     }
 
     // The DSC can only be computed if there's a transaction event. Note that `dsc_from_event`
     // below already checks for the event type.
-    let Some(event) = event.value() else {
-        return sampling_project_info;
-    };
-    let Some(key_config) = project_info.get_public_key_config() else {
-        return sampling_project_info;
-    };
-
-    if let Some(dsc) = utils::dsc_from_event(key_config.public_key, event) {
-        managed_envelope.envelope_mut().set_dsc(dsc);
-        return Some(project_info.clone());
+    if let Some(event) = event.value() {
+        if let Some(key_config) = project_info.get_public_key_config() {
+            if let Some(mut dsc) = utils::dsc_from_event(key_config.public_key, event) {
+                // All other information in the DSC must be discarded, but the sample rate was
+                // actually applied by the client and is therefore correct.
+                let original_sample_rate = original_dsc.and_then(|dsc| dsc.sample_rate);
+                dsc.sample_rate = dsc.sample_rate.or(original_sample_rate);
+
+                managed_envelope.envelope_mut().set_dsc(dsc);
+                return Some(project_info.clone());
+            }
+        }
     }
 
-    sampling_project_info
+    // If we cannot compute a new DSC but the old one is incorrect, we need to remove it.
+    managed_envelope.envelope_mut().remove_dsc();
+    None
 }
 
 /// Computes the sampling decision on the incoming event

tests/integration/test_basic.py

@@ -355,6 +355,8 @@ def send_transaction_with_dsc(mini_sentry, relay, project_id, sampling_project_k
         trace_info={
             "public_key": sampling_project_key,
             "trace_id": "1234F60C11214EB38604F4AE0781BFB2",
+            "release": "[email protected]",
+            "sample_rate": 0.5,
         },
     )
 
@@ -366,12 +368,12 @@ def test_root_project_disabled(mini_sentry, relay):
     mini_sentry.add_full_project_config(project_id)
     disabled_dsn = "00000000000000000000000000000000"
     txn = send_transaction_with_dsc(mini_sentry, relay, project_id, disabled_dsn)
-    assert txn
+    assert txn["contexts"]["trace"].get("client_sample_rate") == 0.5
 
 
 def test_root_project_same(mini_sentry, relay):
     project_id = 42
     mini_sentry.add_full_project_config(project_id)
     same_dsn = mini_sentry.get_dsn_public_key(project_id)
     txn = send_transaction_with_dsc(mini_sentry, relay, project_id, same_dsn)
-    assert txn
+    assert txn["contexts"]["trace"]["client_sample_rate"] == 0.5