From 230d6d8a613930b3d987a3db3543dfffff8b0b28 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Fri, 29 Nov 2024 19:54:20 +1100
Subject: [PATCH] Delay sampling of span to `finish`

---
 sentry-core/src/performance.rs      | 17 ++++++++---------
 sentry-tracing/tests/breadcrumbs.rs |  4 ++--
 sentry-tracing/tests/shared.rs      |  4 ++--
 sentry-tracing/tests/smoke.rs       |  2 +-
 4 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/sentry-core/src/performance.rs b/sentry-core/src/performance.rs
index 3255f5db..c95f8386 100644
--- a/sentry-core/src/performance.rs
+++ b/sentry-core/src/performance.rs
@@ -561,8 +561,8 @@ impl<'a> TransactionData<'a> {
 
 impl Transaction {
     #[cfg(feature = "client")]
-    fn new(mut client: Option<Arc<Client>>, ctx: TransactionContext) -> Self {
-        let (sampled, mut transaction) = match client.as_ref() {
+    fn new(client: Option<Arc<Client>>, ctx: TransactionContext) -> Self {
+        let (sampled, transaction) = match client.as_ref() {
             Some(client) => (
                 client.is_transaction_sampled(&ctx),
                 Some(protocol::Transaction {
@@ -581,13 +581,6 @@ impl Transaction {
             ..Default::default()
         };
 
-        // throw away the transaction here, which means there is nothing to send
-        // on `finish`.
-        if !sampled {
-            transaction = None;
-            client = None;
-        }
-
         Self {
             inner: Arc::new(Mutex::new(TransactionInner {
                 client,
@@ -699,6 +692,12 @@ impl Transaction {
     pub fn finish(self) {
         with_client_impl! {{
             let mut inner = self.inner.lock().unwrap();
+
+            // Discard `Transaction` unless sampled.
+            if !inner.sampled {
+                return;
+            }
+
             if let Some(mut transaction) = inner.transaction.take() {
                 if let Some(client) = inner.client.take() {
                     transaction.finish();
diff --git a/sentry-tracing/tests/breadcrumbs.rs b/sentry-tracing/tests/breadcrumbs.rs
index f62e13b6..f1eab18d 100644
--- a/sentry-tracing/tests/breadcrumbs.rs
+++ b/sentry-tracing/tests/breadcrumbs.rs
@@ -2,12 +2,12 @@ mod shared;
 
 #[test]
 fn breadcrumbs_should_capture_span_fields() {
-    let transport = shared::init_sentry();
+    let transport = shared::init_sentry(0.0); // This test should work even if we are not sampling transactions.
 
     foo();
 
     let data = transport.fetch_and_clear_envelopes();
-    assert_eq!(data.len(), 2);
+    assert_eq!(data.len(), 1);
 
     let event = data.first().expect("should have 1 event");
     let event = match event.items().next().unwrap() {
diff --git a/sentry-tracing/tests/shared.rs b/sentry-tracing/tests/shared.rs
index 9b4409ed..4deadac5 100644
--- a/sentry-tracing/tests/shared.rs
+++ b/sentry-tracing/tests/shared.rs
@@ -3,7 +3,7 @@ use sentry_core::test::TestTransport;
 
 use std::sync::Arc;
 
-pub fn init_sentry() -> Arc<TestTransport> {
+pub fn init_sentry(traces_sample_rate: f32) -> Arc<TestTransport> {
     use tracing_subscriber::prelude::*;
 
     let transport = TestTransport::new();
@@ -11,7 +11,7 @@ pub fn init_sentry() -> Arc<TestTransport> {
         dsn: Some("https://test@sentry-tracing.com/test".parse().unwrap()),
         transport: Some(Arc::new(transport.clone())),
         sample_rate: 1.0,
-        traces_sample_rate: 1.0,
+        traces_sample_rate,
         ..ClientOptions::default()
     };
     Hub::current().bind_client(Some(Arc::new(options.into())));
diff --git a/sentry-tracing/tests/smoke.rs b/sentry-tracing/tests/smoke.rs
index 05c3f9b5..5eb19c0f 100644
--- a/sentry-tracing/tests/smoke.rs
+++ b/sentry-tracing/tests/smoke.rs
@@ -7,7 +7,7 @@ fn function_with_tags(value: i32) {
 
 #[test]
 fn should_instrument_function_with_event() {
-    let transport = shared::init_sentry();
+    let transport = shared::init_sentry(1.0); // Sample all spans.
 
     function_with_tags(1);