From 62a3aab998f1e4db96ac95c662275ac92ecd7f8a Mon Sep 17 00:00:00 2001
From: Tom Milligan <tom@reinfer.io>
Date: Fri, 4 Nov 2022 14:08:56 +0000
Subject: [PATCH] sentry-core: support custom user data in TransactionContext

---
 CHANGELOG.md                   |  6 +++
 sentry-core/src/performance.rs | 88 +++++++++++++++++++++++++++++++++-
 sentry/tests/test_client.rs    |  1 +
 3 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ec4a08dc..bc1afb68 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+**Features**:
+
+- Users of `TransactionContext` may now add `custom` context to it. This may be used by `traces_sampler` to decide sampling rates on a per-transaction basis. ([#512](https://github.com/getsentry/sentry-rust/pull/512))
+
 ## 0.28.0
 
 **Breaking Changes**:
diff --git a/sentry-core/src/performance.rs b/sentry-core/src/performance.rs
index 6672d6de..372681a4 100644
--- a/sentry-core/src/performance.rs
+++ b/sentry-core/src/performance.rs
@@ -51,6 +51,13 @@ impl Hub {
 
 // "Context" Types:
 
+/// Arbitrary data passed by the caller, when starting a transaction.
+///
+/// May be inspected by the user in the `traces_sampler` callback, if set.
+///
+/// Represents arbitrary JSON data, the top level of which must be a map.
+pub type CustomTransactionContext = serde_json::Map<String, serde_json::Value>;
+
 /// The Transaction Context used to start a new Performance Monitoring Transaction.
 ///
 /// The Transaction Context defines the metadata for a Performance Monitoring
@@ -63,6 +70,7 @@ pub struct TransactionContext {
     trace_id: protocol::TraceId,
     parent_span_id: Option<protocol::SpanId>,
     sampled: Option<bool>,
+    custom: Option<CustomTransactionContext>,
 }
 
 impl TransactionContext {
@@ -108,6 +116,7 @@ impl TransactionContext {
             trace_id,
             parent_span_id,
             sampled,
+            custom: None,
         }
     }
 
@@ -144,6 +153,7 @@ impl TransactionContext {
             trace_id,
             parent_span_id: Some(parent_span_id),
             sampled,
+            custom: None,
         }
     }
 
@@ -154,6 +164,56 @@ impl TransactionContext {
     pub fn set_sampled(&mut self, sampled: impl Into<Option<bool>>) {
         self.sampled = sampled.into();
     }
+
+    /// Get the sampling decision for this Transaction.
+    pub fn sampled(&self) -> Option<bool> {
+        self.sampled
+    }
+
+    /// Get the name of this Transaction.
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    /// Get the operation of this Transaction.
+    pub fn operation(&self) -> &str {
+        &self.op
+    }
+
+    /// Get the custom context of this Transaction.
+    pub fn custom(&self) -> Option<&CustomTransactionContext> {
+        self.custom.as_ref()
+    }
+
+    /// Update the custom context of this Transaction.
+    ///
+    /// For simply adding a key, use the `set_custom_key` method.
+    pub fn custom_mut(&mut self) -> &mut Option<CustomTransactionContext> {
+        &mut self.custom
+    }
+
+    /// Inserts a key-value pair into the custom context.
+    ///
+    /// If the context did not have this key present, None is returned.
+    ///
+    /// If the context did have this key present, the value is updated, and the old value is returned.
+    pub fn custom_insert(
+        &mut self,
+        key: String,
+        value: serde_json::Value,
+    ) -> Option<serde_json::Value> {
+        // Get the custom context
+        let mut custom = None;
+        std::mem::swap(&mut self.custom, &mut custom);
+
+        // Initialise the context, if not used yet
+        let mut custom = custom.unwrap_or_default();
+
+        // And set our key
+        let existing_value = custom.insert(key, value);
+        std::mem::swap(&mut self.custom, &mut Some(custom));
+        existing_value
+    }
 }
 
 /// A function to be run for each new transaction, to determine the rate at which
@@ -737,7 +797,7 @@ mod tests {
         ctx.set_sampled(false);
         assert_eq!(transaction_sample_rate(Some(&|_| { 0.7 }), &ctx, 0.3), 0.7);
         // But the sampler may choose to inspect parent sampling
-        let sampler = |ctx: &TransactionContext| match ctx.sampled {
+        let sampler = |ctx: &TransactionContext| match ctx.sampled() {
             Some(true) => 0.8,
             Some(false) => 0.4,
             None => 0.6,
@@ -746,5 +806,31 @@ mod tests {
         assert_eq!(transaction_sample_rate(Some(&sampler), &ctx, 0.3), 0.8);
         ctx.set_sampled(None);
         assert_eq!(transaction_sample_rate(Some(&sampler), &ctx, 0.3), 0.6);
+
+        // Can use first-class and custom attributes of the context.
+        let sampler = |ctx: &TransactionContext| {
+            if ctx.name() == "must-name" || ctx.operation() == "must-operation" {
+                return 1.0;
+            }
+
+            if let Some(custom) = ctx.custom() {
+                if let Some(rate) = custom.get("rate") {
+                    if let Some(rate) = rate.as_f64() {
+                        return rate as f32;
+                    }
+                }
+            }
+
+            0.1
+        };
+        // First class attributes
+        let ctx = TransactionContext::new("noop", "must-operation");
+        assert_eq!(transaction_sample_rate(Some(&sampler), &ctx, 0.3), 1.0);
+        let ctx = TransactionContext::new("must-name", "noop");
+        assert_eq!(transaction_sample_rate(Some(&sampler), &ctx, 0.3), 1.0);
+        // Custom data payload
+        let mut ctx = TransactionContext::new("noop", "noop");
+        ctx.custom_insert("rate".to_owned(), serde_json::json!(0.7));
+        assert_eq!(transaction_sample_rate(Some(&sampler), &ctx, 0.3), 0.7);
     }
 }
diff --git a/sentry/tests/test_client.rs b/sentry/tests/test_client.rs
index a3ca737a..9fbd8aa3 100644
--- a/sentry/tests/test_client.rs
+++ b/sentry/tests/test_client.rs
@@ -18,6 +18,7 @@ fn test_into_client() {
         "https://public@example.com/42%21",
         sentry::ClientOptions {
             release: Some("foo@1.0".into()),
+            traces_sampler: Some(Arc::new(|_| 0.1)),
             ..Default::default()
         },
     ));