tracing_tunnel/receiver/
arena.rs

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

tracing_tunnel/receiver/
arena.rs
