tracing_subscriber/registry/
extensions.rs

Help
1// taken from https://github.com/hyperium/http/blob/master/src/extensions.rs.
2
3use crate::sync::{RwLockReadGuard, RwLockWriteGuard};
4use std::{
5    any::{Any, TypeId},
6    collections::HashMap,
7    fmt,
8    hash::{BuildHasherDefault, Hasher},
9};
10
11#[allow(warnings)]
12type AnyMap = HashMap<TypeId, Box<dyn Any + Send + Sync>, BuildHasherDefault<IdHasher>>;
13
14/// With TypeIds as keys, there's no need to hash them. They are already hashes
15/// themselves, coming from the compiler. The IdHasher holds the u64 of
16/// the TypeId, and then returns it, instead of doing any bit fiddling.
17#[derive(Default, Debug)]
18struct IdHasher(u64);
19
20impl Hasher for IdHasher {
21    fn write(&mut self, _: &[u8]) {
22        unreachable!("TypeId calls write_u64");
23    }
24
25    #[inline]
26    fn write_u64(&mut self, id: u64) {
27        self.0 = id;
28    }
29
30    #[inline]
31    fn finish(&self) -> u64 {
32        self.0
33    }
34}
35
36/// An immutable, read-only reference to a Span's extensions.
37#[derive(Debug)]
38#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
39pub struct Extensions<'a> {
40    inner: RwLockReadGuard<'a, ExtensionsInner>,
41}
42
43impl<'a> Extensions<'a> {
44    #[cfg(feature = "registry")]
45    pub(crate) fn new(inner: RwLockReadGuard<'a, ExtensionsInner>) -> Self {
46        Self { inner }
47    }
48
49    /// Immutably borrows a type previously inserted into this `Extensions`.
50    pub fn get<T: 'static>(&self) -> Option<&T> {
51        self.inner.get::<T>()
52    }
53}
54
55/// An mutable reference to a Span's extensions.
56#[derive(Debug)]
57#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
58pub struct ExtensionsMut<'a> {
59    inner: RwLockWriteGuard<'a, ExtensionsInner>,
60}
61
62impl<'a> ExtensionsMut<'a> {
63    #[cfg(feature = "registry")]
64    pub(crate) fn new(inner: RwLockWriteGuard<'a, ExtensionsInner>) -> Self {
65        Self { inner }
66    }
67
68    /// Insert a type into this `Extensions`.
69    ///
70    /// Note that extensions are _not_
71    /// `Layer`-specific—they are _span_-specific. This means that
72    /// other layers can access and mutate extensions that
73    /// a different Layer recorded. For example, an application might
74    /// have a layer that records execution timings, alongside a layer
75    /// that reports spans and events to a distributed
76    /// tracing system that requires timestamps for spans.
77    /// Ideally, if one layer records a timestamp _x_, the other layer
78    /// should be able to reuse timestamp _x_.
79    ///
80    /// Therefore, extensions should generally be newtypes, rather than common
81    /// types like [`String`](std::string::String), to avoid accidental
82    /// cross-`Layer` clobbering.
83    ///
84    /// ## Panics
85    ///
86    /// If `T` is already present in `Extensions`, then this method will panic.
87    pub fn insert<T: Send + Sync + 'static>(&mut self, val: T) {
88        assert!(self.replace(val).is_none())
89    }
90
91    /// Replaces an existing `T` into this extensions.
92    ///
93    /// If `T` is not present, `Option::None` will be returned.
94    pub fn replace<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
95        self.inner.insert(val)
96    }
97
98    /// Get a mutable reference to a type previously inserted on this `ExtensionsMut`.
99    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
100        self.inner.get_mut::<T>()
101    }
102
103    /// Remove a type from this `Extensions`.
104    ///
105    /// If a extension of this type existed, it will be returned.
106    pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
107        self.inner.remove::<T>()
108    }
109}
110
111/// A type map of span extensions.
112///
113/// [ExtensionsInner] is used by `SpanData` to store and
114/// span-specific data. A given `Layer` can read and write
115/// data that it is interested in recording and emitting.
116#[derive(Default)]
117pub(crate) struct ExtensionsInner {
118    map: AnyMap,
119}
120
121impl ExtensionsInner {
122    /// Create an empty `Extensions`.
123    #[cfg(any(test, feature = "registry"))]
124    #[inline]
125    #[cfg(any(test, feature = "registry"))]
126    pub(crate) fn new() -> ExtensionsInner {
127        ExtensionsInner {
128            map: AnyMap::default(),
129        }
130    }
131
132    /// Insert a type into this `Extensions`.
133    ///
134    /// If a extension of this type already existed, it will
135    /// be returned.
136    pub(crate) fn insert<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
137        self.map
138            .insert(TypeId::of::<T>(), Box::new(val))
139            .and_then(|boxed| {
140                #[allow(warnings)]
141                {
142                    (boxed as Box<Any + 'static>)
143                        .downcast()
144                        .ok()
145                        .map(|boxed| *boxed)
146                }
147            })
148    }
149
150    /// Get a reference to a type previously inserted on this `Extensions`.
151    pub(crate) fn get<T: 'static>(&self) -> Option<&T> {
152        self.map
153            .get(&TypeId::of::<T>())
154            .and_then(|boxed| (&**boxed as &(dyn Any + 'static)).downcast_ref())
155    }
156
157    /// Get a mutable reference to a type previously inserted on this `Extensions`.
158    pub(crate) fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
159        self.map
160            .get_mut(&TypeId::of::<T>())
161            .and_then(|boxed| (&mut **boxed as &mut (dyn Any + 'static)).downcast_mut())
162    }
163
164    /// Remove a type from this `Extensions`.
165    ///
166    /// If a extension of this type existed, it will be returned.
167    pub(crate) fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
168        self.map.remove(&TypeId::of::<T>()).and_then(|boxed| {
169            #[allow(warnings)]
170            {
171                (boxed as Box<Any + 'static>)
172                    .downcast()
173                    .ok()
174                    .map(|boxed| *boxed)
175            }
176        })
177    }
178
179    /// Clear the `ExtensionsInner` in-place, dropping any elements in the map but
180    /// retaining allocated capacity.
181    ///
182    /// This permits the hash map allocation to be pooled by the registry so
183    /// that future spans will not need to allocate new hashmaps.
184    #[cfg(any(test, feature = "registry"))]
185    pub(crate) fn clear(&mut self) {
186        self.map.clear();
187    }
188}
189
190impl fmt::Debug for ExtensionsInner {
191    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
192        f.debug_struct("Extensions")
193            .field("len", &self.map.len())
194            .field("capacity", &self.map.capacity())
195            .finish()
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[derive(Debug, PartialEq)]
204    struct MyType(i32);
205
206    #[test]
207    fn test_extensions() {
208        let mut extensions = ExtensionsInner::new();
209
210        extensions.insert(5i32);
211        extensions.insert(MyType(10));
212
213        assert_eq!(extensions.get(), Some(&5i32));
214        assert_eq!(extensions.get_mut(), Some(&mut 5i32));
215
216        assert_eq!(extensions.remove::<i32>(), Some(5i32));
217        assert!(extensions.get::<i32>().is_none());
218
219        assert_eq!(extensions.get::<bool>(), None);
220        assert_eq!(extensions.get(), Some(&MyType(10)));
221    }
222
223    #[test]
224    fn clear_retains_capacity() {
225        let mut extensions = ExtensionsInner::new();
226        extensions.insert(5i32);
227        extensions.insert(MyType(10));
228        extensions.insert(true);
229
230        assert_eq!(extensions.map.len(), 3);
231        let prev_capacity = extensions.map.capacity();
232        extensions.clear();
233
234        assert_eq!(
235            extensions.map.len(),
236            0,
237            "after clear(), extensions map should have length 0"
238        );
239        assert_eq!(
240            extensions.map.capacity(),
241            prev_capacity,
242            "after clear(), extensions map should retain prior capacity"
243        );
244    }
245
246    #[test]
247    fn clear_drops_elements() {
248        use std::sync::Arc;
249        struct DropMePlease(Arc<()>);
250        struct DropMeTooPlease(Arc<()>);
251
252        let mut extensions = ExtensionsInner::new();
253        let val1 = DropMePlease(Arc::new(()));
254        let val2 = DropMeTooPlease(Arc::new(()));
255
256        let val1_dropped = Arc::downgrade(&val1.0);
257        let val2_dropped = Arc::downgrade(&val2.0);
258        extensions.insert(val1);
259        extensions.insert(val2);
260
261        assert!(val1_dropped.upgrade().is_some());
262        assert!(val2_dropped.upgrade().is_some());
263
264        extensions.clear();
265        assert!(
266            val1_dropped.upgrade().is_none(),
267            "after clear(), val1 should be dropped"
268        );
269        assert!(
270            val2_dropped.upgrade().is_none(),
271            "after clear(), val2 should be dropped"
272        );
273    }
274}
tracing_subscriber/registry/extensions.rs

tracing_subscriber/registry/
extensions.rs
