tracing_subscriber/registry/
extensions.rs

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    /// [subscriber]-specific—they are _span_-specific. This means that
72    /// other subscribers can access and mutate extensions that
73    /// a different Subscriber recorded. For example, an application might
74    /// have a subscriber that records execution timings, alongside a subscriber
75    /// that reports spans and events to a distributed
76    /// tracing system that requires timestamps for spans.
77    /// Ideally, if one subscriber records a timestamp _x_, the other subscriber
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-`Subscriber` clobbering.
83    ///
84    /// ## Panics
85    ///
86    /// If `T` is already present in `Extensions`, then this method will panic.
87    ///
88    /// [subscriber]: crate::subscribe::Subscribe
89    pub fn insert<T: Send + Sync + 'static>(&mut self, val: T) {
90        assert!(self.replace(val).is_none())
91    }
92
93    /// Replaces an existing `T` into this extensions.
94    ///
95    /// If `T` is not present, `Option::None` will be returned.
96    pub fn replace<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
97        self.inner.insert(val)
98    }
99
100    /// Get a mutable reference to a type previously inserted on this `ExtensionsMut`.
101    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
102        self.inner.get_mut::<T>()
103    }
104
105    /// Remove a type from this `Extensions`.
106    ///
107    /// If a extension of this type existed, it will be returned.
108    pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
109        self.inner.remove::<T>()
110    }
111}
112
113/// A type map of span extensions.
114///
115/// [ExtensionsInner] is used by [Data] to store and
116/// span-specific data. A given [Subscriber] can read and write
117/// data that it is interested in recording and emitting.
118#[derive(Default)]
119pub(crate) struct ExtensionsInner {
120    map: AnyMap,
121}
122
123impl ExtensionsInner {
124    /// Create an empty `Extensions`.
125    #[cfg(any(test, feature = "registry"))]
126    #[inline]
127    pub(crate) fn new() -> ExtensionsInner {
128        ExtensionsInner {
129            map: AnyMap::default(),
130        }
131    }
132
133    /// Insert a type into this `Extensions`.
134    ///
135    /// If a extension of this type already existed, it will
136    /// be returned.
137    pub(crate) fn insert<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
138        self.map
139            .insert(TypeId::of::<T>(), Box::new(val))
140            .and_then(|boxed| {
141                #[allow(warnings)]
142                {
143                    (boxed as Box<dyn Any + 'static>)
144                        .downcast()
145                        .ok()
146                        .map(|boxed| *boxed)
147                }
148            })
149    }
150
151    /// Get a reference to a type previously inserted on this `Extensions`.
152    pub(crate) fn get<T: 'static>(&self) -> Option<&T> {
153        self.map
154            .get(&TypeId::of::<T>())
155            .and_then(|boxed| (&**boxed as &(dyn Any + 'static)).downcast_ref())
156    }
157
158    /// Get a mutable reference to a type previously inserted on this `Extensions`.
159    pub(crate) fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
160        self.map
161            .get_mut(&TypeId::of::<T>())
162            .and_then(|boxed| (&mut **boxed as &mut (dyn Any + 'static)).downcast_mut())
163    }
164
165    /// Remove a type from this `Extensions`.
166    ///
167    /// If a extension of this type existed, it will be returned.
168    pub(crate) fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
169        self.map.remove(&TypeId::of::<T>()).and_then(|boxed| {
170            #[allow(warnings)]
171            {
172                (boxed as Box<dyn Any + 'static>)
173                    .downcast()
174                    .ok()
175                    .map(|boxed| *boxed)
176            }
177        })
178    }
179
180    /// Clear the `ExtensionsInner` in-place, dropping any elements in the map but
181    /// retaining allocated capacity.
182    ///
183    /// This permits the hash map allocation to be pooled by the registry so
184    /// that future spans will not need to allocate new hashmaps.
185    #[cfg(any(test, feature = "registry"))]
186    pub(crate) fn clear(&mut self) {
187        self.map.clear();
188    }
189}
190
191impl fmt::Debug for ExtensionsInner {
192    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193        f.debug_struct("Extensions")
194            .field("len", &self.map.len())
195            .field("capacity", &self.map.capacity())
196            .finish()
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    #[derive(Debug, PartialEq)]
205    struct MyType(i32);
206
207    #[test]
208    fn test_extensions() {
209        let mut extensions = ExtensionsInner::new();
210
211        extensions.insert(5i32);
212        extensions.insert(MyType(10));
213
214        assert_eq!(extensions.get(), Some(&5i32));
215        assert_eq!(extensions.get_mut(), Some(&mut 5i32));
216
217        assert_eq!(extensions.remove::<i32>(), Some(5i32));
218        assert!(extensions.get::<i32>().is_none());
219
220        assert_eq!(extensions.get::<bool>(), None);
221        assert_eq!(extensions.get(), Some(&MyType(10)));
222    }
223
224    #[test]
225    fn clear_retains_capacity() {
226        let mut extensions = ExtensionsInner::new();
227        extensions.insert(5i32);
228        extensions.insert(MyType(10));
229        extensions.insert(true);
230
231        assert_eq!(extensions.map.len(), 3);
232        let prev_capacity = extensions.map.capacity();
233        extensions.clear();
234
235        assert_eq!(
236            extensions.map.len(),
237            0,
238            "after clear(), extensions map should have length 0"
239        );
240        assert_eq!(
241            extensions.map.capacity(),
242            prev_capacity,
243            "after clear(), extensions map should retain prior capacity"
244        );
245    }
246
247    #[test]
248    fn clear_drops_elements() {
249        use std::sync::Arc;
250        struct DropMePlease(Arc<()>);
251        struct DropMeTooPlease(Arc<()>);
252
253        let mut extensions = ExtensionsInner::new();
254        let val1 = DropMePlease(Arc::new(()));
255        let val2 = DropMeTooPlease(Arc::new(()));
256
257        let val1_dropped = Arc::downgrade(&val1.0);
258        let val2_dropped = Arc::downgrade(&val2.0);
259        extensions.insert(val1);
260        extensions.insert(val2);
261
262        assert!(val1_dropped.upgrade().is_some());
263        assert!(val2_dropped.upgrade().is_some());
264
265        extensions.clear();
266        assert!(
267            val1_dropped.upgrade().is_none(),
268            "after clear(), val1 should be dropped"
269        );
270        assert!(
271            val2_dropped.upgrade().is_none(),
272            "after clear(), val2 should be dropped"
273        );
274    }
275}
tracing_subscriber/registry/extensions.rs

tracing_subscriber/registry/
extensions.rs
