1 pub(crate) use tracing_core::span::Id;
2 
3 #[derive(Debug)]
4 struct ContextId {
5     id: Id,
6     duplicate: bool,
7 }
8 
9 /// `SpanStack` tracks what spans are currently executing on a thread-local basis.
10 ///
11 /// A "separate current span" for each thread is a semantic choice, as each span
12 /// can be executing in a different thread.
13 #[derive(Debug, Default)]
14 pub(crate) struct SpanStack {
15     stack: Vec<ContextId>,
16 }
17 
18 impl SpanStack {
19     #[inline]
push(&mut self, id: Id) -> bool20     pub(super) fn push(&mut self, id: Id) -> bool {
21         let duplicate = self.stack.iter().any(|i| i.id == id);
22         self.stack.push(ContextId { id, duplicate });
23         !duplicate
24     }
25 
26     #[inline]
pop(&mut self, expected_id: &Id) -> bool27     pub(super) fn pop(&mut self, expected_id: &Id) -> bool {
28         if let Some((idx, _)) = self
29             .stack
30             .iter()
31             .enumerate()
32             .rev()
33             .find(|(_, ctx_id)| ctx_id.id == *expected_id)
34         {
35             let ContextId { id: _, duplicate } = self.stack.remove(idx);
36             return !duplicate;
37         }
38         false
39     }
40 
41     #[inline]
iter(&self) -> impl Iterator<Item = &Id>42     pub(crate) fn iter(&self) -> impl Iterator<Item = &Id> {
43         self.stack
44             .iter()
45             .rev()
46             .filter_map(|ContextId { id, duplicate }| if !*duplicate { Some(id) } else { None })
47     }
48 
49     #[inline]
current(&self) -> Option<&Id>50     pub(crate) fn current(&self) -> Option<&Id> {
51         self.iter().next()
52     }
53 }
54 
55 #[cfg(test)]
56 mod tests {
57     use super::{Id, SpanStack};
58 
59     #[test]
pop_last_span()60     fn pop_last_span() {
61         let mut stack = SpanStack::default();
62         let id = Id::from_u64(1);
63         stack.push(id.clone());
64 
65         assert!(stack.pop(&id));
66     }
67 
68     #[test]
pop_first_span()69     fn pop_first_span() {
70         let mut stack = SpanStack::default();
71         stack.push(Id::from_u64(1));
72         stack.push(Id::from_u64(2));
73 
74         let id = Id::from_u64(1);
75         assert!(stack.pop(&id));
76     }
77 }
78