Refactoring to correctly fetch tags from context

Author: DoumanAsh

sentry-core/src/performance.rs

@@ -526,6 +526,22 @@ pub struct Transaction {
     pub(crate) inner: TransactionArc,
 }
 
+/// Iterable for a transaction's [tags attributes](protocol::Transaction::tags).
+pub struct TransactionTags<'a>(MutexGuard<'a, TransactionInner>);
+impl<'a> TransactionTags<'a> {
+    #[inline]
+    /// Returns iterator over available tags, if transaction is sampled
+    pub fn iter(&'a self) -> Option<impl Iterator<Item = (&'a String, &'a String)> + 'a> {
+        self.0.transaction.as_ref().map(|tx| tx.tags.iter())
+    }
+
+    #[inline]
+    /// Converts self into data pointer
+    pub fn into_data(self) -> TransactionData<'a> {
+        TransactionData(self.0)
+    }
+}
+
 /// Iterable for a transaction's [data attributes](protocol::TraceContext::data).
 pub struct TransactionData<'a>(MutexGuard<'a, TransactionInner>);
 
@@ -647,6 +663,12 @@ impl Transaction {
         TransactionData(self.inner.lock().unwrap())
     }
 
+    /// Returns an iterating accessor to the transaction's
+    /// [tags attributes](protocol::TraceContext::tags).
+    pub fn tags(&self) -> TransactionTags {
+        TransactionTags(self.inner.lock().unwrap())
+    }
+
     /// Get the TransactionContext of the Transaction.
     ///
     /// Note that this clones the underlying value.
@@ -756,6 +778,38 @@ impl Transaction {
     }
 }
 
+/// A smart pointer to a span's [`tags` field](protocol::Span::tags).
+pub struct Tags<'a>(MutexGuard<'a, protocol::Span>);
+
+impl<'a> Tags<'a> {
+    /// Set some tag to be sent with this Span.
+    pub fn set_tag(&mut self, key: String, value: String) {
+        self.0.tags.insert(key, value);
+    }
+
+    #[inline]
+    /// Moves pointer to [`data` field](protocol::Span::data)
+    pub fn into_data(self) -> Data<'a> {
+        Data(self.0)
+    }
+}
+
+impl Deref for Tags<'_> {
+    type Target = BTreeMap<String, String>;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0.tags
+    }
+}
+
+impl DerefMut for Tags<'_> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0.tags
+    }
+}
+
 /// A smart pointer to a span's [`data` field](protocol::Span::data).
 pub struct Data<'a>(MutexGuard<'a, protocol::Span>);
 
@@ -826,6 +880,11 @@ impl Span {
         Data(self.span.lock().unwrap())
     }
 
+    /// Returns a smart pointer to the span's [`tags` field](protocol::Span::tags).
+    pub fn tags(&self) -> Tags {
+        Tags(self.span.lock().unwrap())
+    }
+
     /// Get the TransactionContext of the Span.
     ///
     /// Note that this clones the underlying value.

sentry-tracing/src/converters.rs

@@ -75,13 +75,34 @@ where
             if let Some(span_data) = ext.get::<SentrySpanData>() {
                 match &span_data.sentry_span {
                     TransactionOrSpan::Span(span) => {
-                        for (key, value) in span.data().iter() {
-                            visitor.propagate_span_attr(key, value, propagation, name);
+                        let tags = span.tags();
+
+                        if propagation.is_tags_enabled() {
+                            for (key, value) in tags.iter() {
+                                visitor.propagate_span_tag(key, value);
+                            }
+                        }
+
+                        if propagation.is_attrs_enabled() {
+                            for (key, value) in tags.into_data().iter() {
+                                visitor.propagate_span_attr(key, value, name);
+                            }
                         }
                     }
                     TransactionOrSpan::Transaction(transaction) => {
-                        for (key, value) in transaction.data().iter() {
-                            visitor.propagate_span_attr(key, value, propagation, name);
+                        let tags = transaction.tags();
+                        if propagation.is_tags_enabled() {
+                            if let Some(tags) = tags.iter() {
+                                for (key, value) in tags {
+                                    visitor.propagate_span_tag(key, value);
+                                }
+                            }
+                        }
+
+                        if propagation.is_attrs_enabled() {
+                            for (key, value) in tags.into_data().iter() {
+                                visitor.propagate_span_attr(key, value, name);
+                            }
                         }
                     }
                 }
@@ -100,23 +121,16 @@ pub(crate) struct FieldVisitor {
 }
 
 impl FieldVisitor {
-    fn propagate_span_attr(
-        &mut self,
-        key: &str,
-        value: &Value,
-        span_propagation: SpanPropagation,
-        span_name: &str,
-    ) {
+    fn propagate_span_tag(&mut self, key: &str, value: &str) {
+        //Propagate tags as it is, it will be extracted later on
+        let tag = format!("{TAGS_PREFIX}{key}");
+        self.json_values.entry(tag).or_insert_with(|| value.into());
+    }
+
+    fn propagate_span_attr(&mut self, key: &str, value: &Value, span_name: &str) {
         if key != "message" {
-            if span_propagation.is_tags_enabled() && key.starts_with(TAGS_PREFIX) {
-                //Propagate tags as it is, it will be extracted later on
-                if !self.json_values.contains_key(key) {
-                    self.json_values.insert(key.to_owned(), value.clone());
-                }
-            } else if span_propagation.is_attrs_enabled() {
-                let key = format!("{}:{}", span_name, key);
-                self.json_values.insert(key, value.clone());
-            }
+            let key = format!("{}:{}", span_name, key);
+            self.json_values.insert(key, value.clone());
         }
     }
 

sentry-tracing/tests/shared.rs

@@ -1,5 +1,6 @@
 use sentry::{ClientOptions, Hub};
 use sentry_core::test::TestTransport;
+use sentry_tracing::SpanPropagation;
 
 use std::sync::Arc;
 
@@ -17,7 +18,7 @@ pub fn init_sentry(traces_sample_rate: f32) -> Arc<TestTransport> {
     Hub::current().bind_client(Some(Arc::new(options.into())));
 
     let _ = tracing_subscriber::registry()
-        .with(sentry_tracing::layer().enable_span_attributes())
+        .with(sentry_tracing::layer().enable_span_propagation(SpanPropagation::All))
         .try_init();
 
     transport

sentry-tracing/tests/smoke.rs

@@ -19,6 +19,13 @@ fn should_instrument_function_with_event() {
         unexpected => panic!("Expected event, but got {:#?}", unexpected),
     };
 
+    //Event must be get tag attached
+    let event_tag = event
+        .tags
+        .get("tag")
+        .expect("event should be associated with span's tag");
+    assert_eq!(event_tag, "key");
+
     //Validate transaction is created
     let trace = match event.contexts.get("trace").expect("to get 'trace' context") {
         sentry::protocol::Context::Trace(trace) => trace,