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&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)]
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}
275