Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add LocalSpan::add_property #200

Merged
merged 1 commit into from
Feb 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

## Unreleased

- Add `LocalSpan::add_property` and `LocalSpan::add_properties`.

## v0.6.3

- Add `LocalSpans::to_span_records()`.
Expand Down
58 changes: 58 additions & 0 deletions minitrace/src/local/local_span.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
19 changes: 16 additions & 3 deletions minitrace/src/local/local_span_line.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
}
Expand Down
2 changes: 1 addition & 1 deletion minitrace/src/local/local_span_stack.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}
Expand Down
43 changes: 27 additions & 16 deletions minitrace/src/local/span_queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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()),
Expand Down
28 changes: 28 additions & 0 deletions minitrace/tests/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
3 changes: 3 additions & 0 deletions test-statically-disable/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading