onesignal_tracing_tail_sample/
extensions.rs

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

onesignal_tracing_tail_sample/
extensions.rs
