tracing_subscriber/registry/
stack.rs

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