tracing_tunnel/receiver/
arena.rs

1//! Simple string arena.
2
3use std::{
4    borrow::Cow,
5    collections::{hash_map::DefaultHasher, HashMap, HashSet},
6    hash::{Hash, Hasher},
7    ops,
8    sync::RwLock,
9};
10
11use once_cell::sync::{Lazy, OnceCell};
12use tracing_core::{field::FieldSet, Callsite, Interest, Kind, Level, Metadata};
13
14use crate::types::{CallSiteData, CallSiteKind, TracingLevel};
15
16// An emulation of a hash map with keys equivalent to `CallSiteData` (obviously,
17// we don't want to store `CallSiteData` explicitly because of its size).
18type MetadataMap = HashMap<u64, Vec<&'static Metadata<'static>>>;
19
20impl From<TracingLevel> for Level {
21    fn from(level: TracingLevel) -> Self {
22        match level {
23            TracingLevel::Error => Self::ERROR,
24            TracingLevel::Warn => Self::WARN,
25            TracingLevel::Info => Self::INFO,
26            TracingLevel::Debug => Self::DEBUG,
27            TracingLevel::Trace => Self::TRACE,
28        }
29    }
30}
31
32impl From<CallSiteKind> for Kind {
33    fn from(kind: CallSiteKind) -> Self {
34        match kind {
35            CallSiteKind::Span => Self::SPAN,
36            CallSiteKind::Event => Self::EVENT,
37        }
38    }
39}
40
41#[derive(Debug, Default)]
42struct DynamicCallSite {
43    metadata: OnceCell<&'static Metadata<'static>>,
44}
45
46impl Callsite for DynamicCallSite {
47    fn set_interest(&self, _interest: Interest) {
48        // Does nothing
49    }
50
51    fn metadata(&self) -> &Metadata<'_> {
52        self.metadata
53            .get()
54            .copied()
55            .expect("metadata not initialized")
56    }
57}
58
59#[derive(Debug, Default)]
60pub(crate) struct Arena {
61    strings: RwLock<HashSet<&'static str>>,
62    metadata: RwLock<MetadataMap>,
63}
64
65impl Arena {
66    fn leak(s: Cow<'static, str>) -> &'static str {
67        match s {
68            Cow::Borrowed(s) => s,
69            Cow::Owned(string) => Box::leak(string.into_boxed_str()),
70        }
71    }
72
73    fn new_call_site() -> &'static DynamicCallSite {
74        let call_site = Box::default();
75        Box::leak(call_site)
76    }
77
78    fn lock_strings(&self) -> impl ops::Deref<Target = HashSet<&'static str>> + '_ {
79        self.strings.read().unwrap()
80    }
81
82    fn lock_strings_mut(&self) -> impl ops::DerefMut<Target = HashSet<&'static str>> + '_ {
83        self.strings.write().unwrap()
84    }
85
86    fn alloc_string(&self, s: Cow<'static, str>) -> &'static str {
87        if let Some(existing) = self.lock_strings().get(s.as_ref()).copied() {
88            return existing;
89        }
90
91        let mut lock = self.lock_strings_mut();
92        if let Some(existing) = lock.get(s.as_ref()).copied() {
93            return existing;
94        }
95        let leaked = Self::leak(s);
96        lock.insert(leaked);
97        leaked
98    }
99
100    fn leak_fields(&self, fields: Vec<Cow<'static, str>>) -> &'static [&'static str] {
101        let fields: Box<[_]> = fields
102            .into_iter()
103            .map(|field| self.alloc_string(field))
104            .collect();
105        Box::leak(fields)
106    }
107
108    fn leak_metadata(&self, data: CallSiteData) -> &'static Metadata<'static> {
109        let call_site = Self::new_call_site();
110        let call_site_id = tracing_core::identify_callsite!(call_site);
111        let fields = FieldSet::new(self.leak_fields(data.fields), call_site_id);
112        let metadata = Metadata::new(
113            self.alloc_string(data.name),
114            self.alloc_string(data.target),
115            data.level.into(),
116            data.file.map(|file| self.alloc_string(file)),
117            data.line,
118            data.module_path.map(|path| self.alloc_string(path)),
119            fields,
120            data.kind.into(),
121        );
122
123        let metadata = Box::leak(Box::new(metadata)) as &_;
124        call_site.metadata.set(metadata).unwrap();
125        metadata
126    }
127
128    fn lock_metadata(&self) -> impl ops::Deref<Target = MetadataMap> + '_ {
129        self.metadata.read().unwrap()
130    }
131
132    fn lock_metadata_mut(&self) -> impl ops::DerefMut<Target = MetadataMap> + '_ {
133        self.metadata.write().unwrap()
134    }
135
136    /// Returns the metadata and a flag whether it was allocated in this call.
137    pub(super) fn alloc_metadata(&self, data: CallSiteData) -> (&'static Metadata<'static>, bool) {
138        let hash_value = Self::hash_metadata(&data);
139        let scanned_bucket_len = {
140            let lock = self.lock_metadata();
141            if let Some(bucket) = lock.get(&hash_value) {
142                for &metadata in bucket {
143                    if Self::eq_metadata(&data, metadata) {
144                        return (metadata, false);
145                    }
146                }
147                bucket.len()
148            } else {
149                0
150            }
151        };
152
153        let mut lock = self.lock_metadata_mut();
154        let bucket = lock.entry(hash_value).or_default();
155        for &metadata in &bucket[scanned_bucket_len..] {
156            if Self::eq_metadata(&data, metadata) {
157                return (metadata, false);
158            }
159        }
160
161        // Finally, we need to actually leak metadata.
162        let metadata = self.leak_metadata(data);
163        bucket.push(metadata);
164        (metadata, true)
165    }
166
167    // The returned hash doesn't necessarily match the hash of `Metadata`, but it is the same
168    // for the equivalent `(kind, data)` tuples, which is what we need.
169    fn hash_metadata(data: &CallSiteData) -> u64 {
170        let mut hasher = DefaultHasher::new();
171        data.hash(&mut hasher);
172        hasher.finish()
173    }
174
175    fn eq_metadata(data: &CallSiteData, metadata: &Metadata<'_>) -> bool {
176        // number comparisons go first
177        matches!(data.kind, CallSiteKind::Span) == metadata.is_span()
178            && Level::from(data.level) == *metadata.level()
179            && data.line == metadata.line()
180            // ...then, string comparisons
181            && data.name == metadata.name()
182            && data.target == metadata.target()
183            && data.module_path.as_ref().map(Cow::as_ref) == metadata.module_path()
184            && data.file.as_ref().map(Cow::as_ref) == metadata.file()
185            // ...and finally, comparison of fields
186            && data
187                .fields
188                .iter()
189                .map(Cow::as_ref)
190                .eq(metadata.fields().iter().map(|field| field.name()))
191    }
192}
193
194pub(crate) static ARENA: Lazy<Arena> = Lazy::new(Arena::default);
tracing_tunnel/receiver/arena.rs

tracing_tunnel/receiver/
arena.rs
