vendor/tracing-subscriber/src/registry/extensions.rs

   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.debug_struct("Extensions")
 193             .field("len", &self.map.len())
 194             .field("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 }