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 ] |
20 | 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 ] |
27 | 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 ] |
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 ] |
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] |
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] |
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 | |