1 | // taken from https://github.com/hyperium/http/blob/master/src/extensions.rs. |
2 | |
3 | use crate::sync::{RwLockReadGuard, RwLockWriteGuard}; |
4 | use std::{ |
5 | any::{Any, TypeId}, |
6 | collections::HashMap, |
7 | fmt, |
8 | hash::{BuildHasherDefault, Hasher}, |
9 | }; |
10 | |
11 | #[allow (warnings)] |
12 | type 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)] |
18 | struct IdHasher(u64); |
19 | |
20 | impl 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" )))] |
39 | pub struct Extensions<'a> { |
40 | inner: RwLockReadGuard<'a, ExtensionsInner>, |
41 | } |
42 | |
43 | impl<'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" )))] |
58 | pub struct ExtensionsMut<'a> { |
59 | inner: RwLockWriteGuard<'a, ExtensionsInner>, |
60 | } |
61 | |
62 | impl<'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)] |
117 | pub(crate) struct ExtensionsInner { |
118 | map: AnyMap, |
119 | } |
120 | |
121 | impl 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 | |
190 | impl fmt::Debug for ExtensionsInner { |
191 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
192 | f&mut DebugStruct<'_, '_>.debug_struct("Extensions" ) |
193 | .field("len" , &self.map.len()) |
194 | .field(name:"capacity" , &self.map.capacity()) |
195 | .finish() |
196 | } |
197 | } |
198 | |
199 | #[cfg (test)] |
200 | mod 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 | } |
275 | |