feat: add LocalSpan::add_property (#200)

Signed-off-by: Andy Lok <[email protected]>

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add `LocalSpan::add_property` and `LocalSpan::add_properties`.
+
 ## v0.6.3
 
 - Add `LocalSpans::to_span_records()`.

minitrace/src/local/local_span.rs

@@ -54,6 +54,64 @@ impl LocalSpan {
         }
     }
 
+    /// Add a single property to the current local parent. If the local parent is a [`Span`],
+    /// the property will not be added to the `Span`.
+    ///
+    /// A property is an arbitrary key-value pair associated with a span.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use minitrace::prelude::*;
+    ///
+    /// LocalSpan::add_property(|| ("key", "value"));
+    /// ```
+    ///
+    /// [`Span`]: crate::Span
+    #[inline]
+    pub fn add_property<K, V, F>(property: F)
+    where
+        K: Into<Cow<'static, str>>,
+        V: Into<Cow<'static, str>>,
+        F: FnOnce() -> (K, V),
+    {
+        Self::add_properties(|| [property()])
+    }
+
+    /// Add multiple properties to the current local parent. If the local parent is a [`Span`],
+    /// the properties will not be added to the `Span`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use minitrace::prelude::*;
+    ///
+    /// LocalSpan::add_properties(|| [("key1", "value1"), ("key2", "value2")]);
+    /// ```
+    ///
+    /// [`Span`]: crate::Span
+    #[inline]
+    pub fn add_properties<K, V, I, F>(properties: F)
+    where
+        K: Into<Cow<'static, str>>,
+        V: Into<Cow<'static, str>>,
+        I: IntoIterator<Item = (K, V)>,
+        F: FnOnce() -> I,
+    {
+        #[cfg(feature = "enable")]
+        {
+            LOCAL_SPAN_STACK
+                .try_with(|s| {
+                    let span_stack = &mut *s.borrow_mut();
+                    let span_line = span_stack.current_span_line()?;
+                    let parent_handle = span_line.current_parent_handle()?;
+                    span_line.add_properties(&parent_handle, properties);
+                    Some(())
+                })
+                .ok();
+        }
+    }
+
     /// Add a single property to the `LocalSpan` and return the modified `LocalSpan`.
     ///
     /// A property is an arbitrary key-value pair associated with a span.

minitrace/src/local/local_span_line.rs

@@ -77,14 +77,27 @@ impl SpanLine {
                 .iter()
                 .map(|item| CollectTokenItem {
                     trace_id: item.trace_id,
-                    parent_id: self.span_queue.current_span_id().unwrap_or(item.parent_id),
+                    parent_id: self
+                        .span_queue
+                        .current_parent_id()
+                        .unwrap_or(item.parent_id),
                     collect_id: item.collect_id,
                     is_root: false,
                 })
                 .collect()
         })
     }
 
+    #[inline]
+    pub fn current_parent_handle(&self) -> Option<LocalSpanHandle> {
+        let span_handle = self.span_queue.current_parent_handle()?;
+        let span_line_epoch = self.epoch;
+        Some(LocalSpanHandle {
+            span_handle,
+            span_line_epoch,
+        })
+    }
+
     #[inline]
     pub fn collect(self, span_line_epoch: usize) -> Option<(RawSpans, Option<CollectToken>)> {
         (self.epoch == span_line_epoch)
@@ -158,13 +171,13 @@ span1 []
         assert_eq!(current_token.as_slice(), &[
             CollectTokenItem {
                 trace_id: TraceId(1234),
-                parent_id: span_line.span_queue.current_span_id().unwrap(),
+                parent_id: span_line.span_queue.current_parent_id().unwrap(),
                 collect_id: 42,
                 is_root: false,
             },
             CollectTokenItem {
                 trace_id: TraceId(1235),
-                parent_id: span_line.span_queue.current_span_id().unwrap(),
+                parent_id: span_line.span_queue.current_parent_id().unwrap(),
                 collect_id: 43,
                 is_root: false,
             }

minitrace/src/local/local_span_stack.rs

@@ -121,7 +121,7 @@ impl LocalSpanStack {
     }
 
     #[inline]
-    fn current_span_line(&mut self) -> Option<&mut SpanLine> {
+    pub fn current_span_line(&mut self) -> Option<&mut SpanLine> {
         self.span_lines.last_mut()
     }
 }

minitrace/src/local/span_queue.rs

@@ -104,10 +104,21 @@ impl SpanQueue {
     }
 
     #[inline]
-    pub fn current_span_id(&self) -> Option<SpanId> {
+    pub fn current_parent_id(&self) -> Option<SpanId> {
         self.next_parent_id
     }
 
+    #[inline]
+    pub fn current_parent_handle(&self) -> Option<SpanHandle> {
+        self.next_parent_id.and_then(|id| {
+            let index = self.span_queue.iter().position(|span| span.id == id);
+
+            debug_assert!(index.is_some());
+
+            Some(SpanHandle { index: index? })
+        })
+    }
+
     #[cfg(test)]
     pub fn get_raw_span(&self, handle: &SpanHandle) -> &RawSpan {
         &self.span_queue[handle.index]
@@ -237,81 +248,81 @@ span1 []
     fn last_span_id() {
         let mut queue = SpanQueue::with_capacity(16);
 
-        assert_eq!(queue.current_span_id(), None);
+        assert_eq!(queue.current_parent_id(), None);
         {
             let span1 = queue.start_span("span1").unwrap();
             assert_eq!(
-                queue.current_span_id().unwrap(),
+                queue.current_parent_id().unwrap(),
                 queue.get_raw_span(&span1).id
             );
             queue.finish_span(span1);
-            assert_eq!(queue.current_span_id(), None);
+            assert_eq!(queue.current_parent_id(), None);
         }
         {
             let span2 = queue.start_span("span2").unwrap();
             assert_eq!(
-                queue.current_span_id().unwrap(),
+                queue.current_parent_id().unwrap(),
                 queue.get_raw_span(&span2).id
             );
             {
                 let span3 = queue.start_span("span3").unwrap();
                 assert_eq!(
-                    queue.current_span_id().unwrap(),
+                    queue.current_parent_id().unwrap(),
                     queue.get_raw_span(&span3).id
                 );
                 queue.finish_span(span3);
                 assert_eq!(
-                    queue.current_span_id().unwrap(),
+                    queue.current_parent_id().unwrap(),
                     queue.get_raw_span(&span2).id
                 );
             }
             {
                 let span4 = queue.start_span("span4").unwrap();
                 assert_eq!(
-                    queue.current_span_id().unwrap(),
+                    queue.current_parent_id().unwrap(),
                     queue.get_raw_span(&span4).id
                 );
                 {
                     let span5 = queue.start_span("span5").unwrap();
                     assert_eq!(
-                        queue.current_span_id().unwrap(),
+                        queue.current_parent_id().unwrap(),
                         queue.get_raw_span(&span5).id
                     );
                     {
                         let span6 = queue.start_span("span6").unwrap();
                         assert_eq!(
-                            queue.current_span_id().unwrap(),
+                            queue.current_parent_id().unwrap(),
                             queue.get_raw_span(&span6).id
                         );
                         queue.finish_span(span6);
                         assert_eq!(
-                            queue.current_span_id().unwrap(),
+                            queue.current_parent_id().unwrap(),
                             queue.get_raw_span(&span5).id
                         );
                     }
                     queue.finish_span(span5);
                     assert_eq!(
-                        queue.current_span_id().unwrap(),
+                        queue.current_parent_id().unwrap(),
                         queue.get_raw_span(&span4).id
                     );
                 }
                 queue.finish_span(span4);
                 assert_eq!(
-                    queue.current_span_id().unwrap(),
+                    queue.current_parent_id().unwrap(),
                     queue.get_raw_span(&span2).id
                 );
             }
             queue.finish_span(span2);
-            assert_eq!(queue.current_span_id(), None);
+            assert_eq!(queue.current_parent_id(), None);
         }
         {
             let span7 = queue.start_span("span7").unwrap();
             assert_eq!(
-                queue.current_span_id().unwrap(),
+                queue.current_parent_id().unwrap(),
                 queue.get_raw_span(&span7).id
             );
             queue.finish_span(span7);
-            assert_eq!(queue.current_span_id(), None);
+            assert_eq!(queue.current_parent_id(), None);
         }
         assert_eq!(
             tree_str_from_raw_spans(queue.take_queue()),

minitrace/tests/lib.rs

@@ -656,6 +656,34 @@ fn test_elapsed() {
     minitrace::flush();
 }
 
+#[test]
+#[serial]
+fn test_add_property() {
+    let (reporter, collected_spans) = TestReporter::new();
+    minitrace::set_reporter(reporter, Config::default());
+
+    {
+        let root = Span::root("root", SpanContext::random());
+        let _g = root.set_local_parent();
+        LocalSpan::add_property(|| ("noop", "noop"));
+        LocalSpan::add_properties(|| [("noop", "noop")]);
+        let _span = LocalSpan::enter_with_local_parent("span");
+        LocalSpan::add_property(|| ("k1", "v1"));
+        LocalSpan::add_properties(|| [("k2", "v2"), ("k3", "v3")]);
+    }
+
+    minitrace::flush();
+
+    let expected_graph = r#"
+root []
+    span [("k1", "v1"), ("k2", "v2"), ("k3", "v3")]
+"#;
+    assert_eq!(
+        tree_str_from_span_records(collected_spans.lock().clone()),
+        expected_graph
+    );
+}
+
 #[test]
 #[serial]
 fn test_macro_properties() {

test-statically-disable/src/main.rs

@@ -40,6 +40,9 @@ fn main() {
 
     let _span2 = LocalSpan::enter_with_local_parent("span2");
 
+    LocalSpan::add_property(|| ("k", "v"));
+    LocalSpan::add_properties(|| [("k", "v")]);
+
     let local_collector = LocalCollector::start();
     let _ = LocalSpan::enter_with_local_parent("span3");
     let local_spans = local_collector.collect();