Vendor things

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

@@ -0,0 +1,274 @@
+// taken from https://github.com/hyperium/http/blob/master/src/extensions.rs.
+
+use crate::sync::{RwLockReadGuard, RwLockWriteGuard};
+use std::{
+    any::{Any, TypeId},
+    collections::HashMap,
+    fmt,
+    hash::{BuildHasherDefault, Hasher},
+};
+
+#[allow(warnings)]
+type AnyMap = HashMap<TypeId, Box<dyn Any + Send + Sync>, BuildHasherDefault<IdHasher>>;
+
+/// With TypeIds as keys, there's no need to hash them. They are already hashes
+/// themselves, coming from the compiler. The IdHasher holds the u64 of
+/// the TypeId, and then returns it, instead of doing any bit fiddling.
+#[derive(Default, Debug)]
+struct IdHasher(u64);
+
+impl Hasher for IdHasher {
+    fn write(&mut self, _: &[u8]) {
+        unreachable!("TypeId calls write_u64");
+    }
+
+    #[inline]
+    fn write_u64(&mut self, id: u64) {
+        self.0 = id;
+    }
+
+    #[inline]
+    fn finish(&self) -> u64 {
+        self.0
+    }
+}
+
+/// An immutable, read-only reference to a Span's extensions.
+#[derive(Debug)]
+#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
+pub struct Extensions<'a> {
+    inner: RwLockReadGuard<'a, ExtensionsInner>,
+}
+
+impl<'a> Extensions<'a> {
+    #[cfg(feature = "registry")]
+    pub(crate) fn new(inner: RwLockReadGuard<'a, ExtensionsInner>) -> Self {
+        Self { inner }
+    }
+
+    /// Immutably borrows a type previously inserted into this `Extensions`.
+    pub fn get<T: 'static>(&self) -> Option<&T> {
+        self.inner.get::<T>()
+    }
+}
+
+/// An mutable reference to a Span's extensions.
+#[derive(Debug)]
+#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
+pub struct ExtensionsMut<'a> {
+    inner: RwLockWriteGuard<'a, ExtensionsInner>,
+}
+
+impl<'a> ExtensionsMut<'a> {
+    #[cfg(feature = "registry")]
+    pub(crate) fn new(inner: RwLockWriteGuard<'a, ExtensionsInner>) -> Self {
+        Self { inner }
+    }
+
+    /// Insert a type into this `Extensions`.
+    ///
+    /// Note that extensions are _not_
+    /// `Layer`-specific—they are _span_-specific. This means that
+    /// other layers can access and mutate extensions that
+    /// a different Layer recorded. For example, an application might
+    /// have a layer that records execution timings, alongside a layer
+    /// that reports spans and events to a distributed
+    /// tracing system that requires timestamps for spans.
+    /// Ideally, if one layer records a timestamp _x_, the other layer
+    /// should be able to reuse timestamp _x_.
+    ///
+    /// Therefore, extensions should generally be newtypes, rather than common
+    /// types like [`String`](std::string::String), to avoid accidental
+    /// cross-`Layer` clobbering.
+    ///
+    /// ## Panics
+    ///
+    /// If `T` is already present in `Extensions`, then this method will panic.
+    pub fn insert<T: Send + Sync + 'static>(&mut self, val: T) {
+        assert!(self.replace(val).is_none())
+    }
+
+    /// Replaces an existing `T` into this extensions.
+    ///
+    /// If `T` is not present, `Option::None` will be returned.
+    pub fn replace<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
+        self.inner.insert(val)
+    }
+
+    /// Get a mutable reference to a type previously inserted on this `ExtensionsMut`.
+    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
+        self.inner.get_mut::<T>()
+    }
+
+    /// Remove a type from this `Extensions`.
+    ///
+    /// If a extension of this type existed, it will be returned.
+    pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
+        self.inner.remove::<T>()
+    }
+}
+
+/// A type map of span extensions.
+///
+/// [ExtensionsInner] is used by `SpanData` to store and
+/// span-specific data. A given `Layer` can read and write
+/// data that it is interested in recording and emitting.
+#[derive(Default)]
+pub(crate) struct ExtensionsInner {
+    map: AnyMap,
+}
+
+impl ExtensionsInner {
+    /// Create an empty `Extensions`.
+    #[cfg(any(test, feature = "registry"))]
+    #[inline]
+    #[cfg(any(test, feature = "registry"))]
+    pub(crate) fn new() -> ExtensionsInner {
+        ExtensionsInner {
+            map: AnyMap::default(),
+        }
+    }
+
+    /// Insert a type into this `Extensions`.
+    ///
+    /// If a extension of this type already existed, it will
+    /// be returned.
+    pub(crate) fn insert<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
+        self.map
+            .insert(TypeId::of::<T>(), Box::new(val))
+            .and_then(|boxed| {
+                #[allow(warnings)]
+                {
+                    (boxed as Box<Any + 'static>)
+                        .downcast()
+                        .ok()
+                        .map(|boxed| *boxed)
+                }
+            })
+    }
+
+    /// Get a reference to a type previously inserted on this `Extensions`.
+    pub(crate) fn get<T: 'static>(&self) -> Option<&T> {
+        self.map
+            .get(&TypeId::of::<T>())
+            .and_then(|boxed| (&**boxed as &(dyn Any + 'static)).downcast_ref())
+    }
+
+    /// Get a mutable reference to a type previously inserted on this `Extensions`.
+    pub(crate) fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
+        self.map
+            .get_mut(&TypeId::of::<T>())
+            .and_then(|boxed| (&mut **boxed as &mut (dyn Any + 'static)).downcast_mut())
+    }
+
+    /// Remove a type from this `Extensions`.
+    ///
+    /// If a extension of this type existed, it will be returned.
+    pub(crate) fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
+        self.map.remove(&TypeId::of::<T>()).and_then(|boxed| {
+            #[allow(warnings)]
+            {
+                (boxed as Box<Any + 'static>)
+                    .downcast()
+                    .ok()
+                    .map(|boxed| *boxed)
+            }
+        })
+    }
+
+    /// Clear the `ExtensionsInner` in-place, dropping any elements in the map but
+    /// retaining allocated capacity.
+    ///
+    /// This permits the hash map allocation to be pooled by the registry so
+    /// that future spans will not need to allocate new hashmaps.
+    #[cfg(any(test, feature = "registry"))]
+    pub(crate) fn clear(&mut self) {
+        self.map.clear();
+    }
+}
+
+impl fmt::Debug for ExtensionsInner {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Extensions")
+            .field("len", &self.map.len())
+            .field("capacity", &self.map.capacity())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[derive(Debug, PartialEq)]
+    struct MyType(i32);
+
+    #[test]
+    fn test_extensions() {
+        let mut extensions = ExtensionsInner::new();
+
+        extensions.insert(5i32);
+        extensions.insert(MyType(10));
+
+        assert_eq!(extensions.get(), Some(&5i32));
+        assert_eq!(extensions.get_mut(), Some(&mut 5i32));
+
+        assert_eq!(extensions.remove::<i32>(), Some(5i32));
+        assert!(extensions.get::<i32>().is_none());
+
+        assert_eq!(extensions.get::<bool>(), None);
+        assert_eq!(extensions.get(), Some(&MyType(10)));
+    }
+
+    #[test]
+    fn clear_retains_capacity() {
+        let mut extensions = ExtensionsInner::new();
+        extensions.insert(5i32);
+        extensions.insert(MyType(10));
+        extensions.insert(true);
+
+        assert_eq!(extensions.map.len(), 3);
+        let prev_capacity = extensions.map.capacity();
+        extensions.clear();
+
+        assert_eq!(
+            extensions.map.len(),
+            0,
+            "after clear(), extensions map should have length 0"
+        );
+        assert_eq!(
+            extensions.map.capacity(),
+            prev_capacity,
+            "after clear(), extensions map should retain prior capacity"
+        );
+    }
+
+    #[test]
+    fn clear_drops_elements() {
+        use std::sync::Arc;
+        struct DropMePlease(Arc<()>);
+        struct DropMeTooPlease(Arc<()>);
+
+        let mut extensions = ExtensionsInner::new();
+        let val1 = DropMePlease(Arc::new(()));
+        let val2 = DropMeTooPlease(Arc::new(()));
+
+        let val1_dropped = Arc::downgrade(&val1.0);
+        let val2_dropped = Arc::downgrade(&val2.0);
+        extensions.insert(val1);
+        extensions.insert(val2);
+
+        assert!(val1_dropped.upgrade().is_some());
+        assert!(val2_dropped.upgrade().is_some());
+
+        extensions.clear();
+        assert!(
+            val1_dropped.upgrade().is_none(),
+            "after clear(), val1 should be dropped"
+        );
+        assert!(
+            val2_dropped.upgrade().is_none(),
+            "after clear(), val2 should be dropped"
+        );
+    }
+}