From 2c83f670f8f2910e27462d3bd1250b1513e97e6b Mon Sep 17 00:00:00 2001 From: Alex Saveau Date: Sun, 31 Dec 2023 12:02:52 -0800 Subject: [PATCH] Eliminate extensions allocations Signed-off-by: Alex Saveau --- tracing-subscriber/src/lib.rs | 1 + tracing-subscriber/src/registry/extensions.rs | 188 +++++++++++++++--- 2 files changed, 161 insertions(+), 28 deletions(-) diff --git a/tracing-subscriber/src/lib.rs b/tracing-subscriber/src/lib.rs index 4b1ce72d82..2b18324d6c 100644 --- a/tracing-subscriber/src/lib.rs +++ b/tracing-subscriber/src/lib.rs @@ -171,6 +171,7 @@ // "needless". #![allow(clippy::needless_update)] #![cfg_attr(not(feature = "std"), no_std)] +#![cfg_attr(nightly, feature(ptr_mask, strict_provenance))] #[cfg(feature = "alloc")] extern crate alloc; diff --git a/tracing-subscriber/src/registry/extensions.rs b/tracing-subscriber/src/registry/extensions.rs index ecac177bff..b2e1510e5e 100644 --- a/tracing-subscriber/src/registry/extensions.rs +++ b/tracing-subscriber/src/registry/extensions.rs @@ -1,15 +1,134 @@ // taken from https://github.com/hyperium/http/blob/master/src/extensions.rs. -use crate::sync::{RwLockReadGuard, RwLockWriteGuard}; use std::{ - any::{Any, TypeId}, - collections::HashMap, + any::TypeId, + collections::{hash_map::Entry, HashMap}, fmt, hash::{BuildHasherDefault, Hasher}, + mem, ptr, }; +use crate::registry::extensions::boxed_entry::{Aligned, BoxedEntry, BoxedEntryOp}; +use crate::sync::{RwLockReadGuard, RwLockWriteGuard}; + +#[allow(unreachable_pub)] +mod boxed_entry { + use std::mem::{transmute, MaybeUninit}; + use std::ptr; + use std::ptr::NonNull; + + /// Used to give ourselves one free bit for liveness checks. + #[repr(align(2))] + pub struct Aligned(pub T); + + /// Wrapper around an untyped pointer to avoid casting bugs. + /// + /// This type holds the pointer to a `Box>`. + #[derive(Copy, Clone)] + pub struct ExtPtr(NonNull<()>); + + impl ExtPtr { + fn new(val: T) -> Self { + let mut this = { + let ptr = Box::into_raw(Box::new(Aligned(val))); + Self(unsafe { NonNull::new_unchecked(ptr) }.cast()) + }; + this.set_live(true); + this + } + + pub fn typed(self) -> NonNull> { + #[cfg(not(nightly))] + let ptr = { + let addr = self.0.as_ptr() as usize; + let addr = addr & !1; // Zero out the live bit + addr as *mut _ + }; + #[cfg(nightly)] + let ptr = self.0.as_ptr().mask(!1).cast(); + unsafe { NonNull::new_unchecked(ptr) } + } + + pub fn live(&self) -> bool { + #[cfg(not(nightly))] + let addr = self.0.as_ptr() as usize; + #[cfg(nightly)] + let addr = self.0.as_ptr().addr(); + (addr & 1) == 1 + } + + pub fn set_live(&mut self, live: bool) { + let update = |addr: usize| if live { addr | 1 } else { addr & !1 }; + #[cfg(not(nightly))] + let ptr = update(self.0.as_ptr() as usize) as *mut _; + #[cfg(nightly)] + let ptr = self.0.as_ptr().map_addr(update); + self.0 = unsafe { NonNull::new_unchecked(ptr) }; + } + } + + /// Extension storage + pub struct BoxedEntry { + pub ptr: ExtPtr, + op_fn: NonNull, + } + + unsafe impl Send for BoxedEntry {} + + unsafe impl Sync for BoxedEntry {} + + pub enum BoxedEntryOp { + DropLiveValue, + Drop, + } + + unsafe fn run_boxed_entry_op(ext: ExtPtr, op: BoxedEntryOp) { + let ptr = ext.typed::().as_ptr(); + match op { + BoxedEntryOp::DropLiveValue => unsafe { + ptr::drop_in_place(ptr); + }, + BoxedEntryOp::Drop => { + if ext.live() { + unsafe { + drop(Box::from_raw(ptr)); + } + } else { + unsafe { + drop(Box::from_raw(ptr.cast::>>())); + } + } + } + } + } + + impl BoxedEntry { + pub fn new(val: T) -> Self { + Self { + ptr: ExtPtr::new(val), + op_fn: { + let fn_ptr = run_boxed_entry_op:: as *mut _; + unsafe { NonNull::new_unchecked(fn_ptr) } + }, + } + } + + pub fn op_fn(&self) -> unsafe fn(ExtPtr, BoxedEntryOp) { + unsafe { transmute(self.op_fn.as_ptr()) } + } + } + + impl Drop for BoxedEntry { + fn drop(&mut self) { + unsafe { + self.op_fn()(self.ptr, BoxedEntryOp::Drop); + } + } + } +} + #[allow(warnings)] -type AnyMap = HashMap, BuildHasherDefault>; +type AnyMap = HashMap>; /// 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 @@ -135,46 +254,54 @@ impl ExtensionsInner { /// If a extension of this type already existed, it will /// be returned. pub(crate) fn insert(&mut self, val: T) -> Option { - self.map - .insert(TypeId::of::(), Box::new(val)) - .and_then(|boxed| { - #[allow(warnings)] - { - (boxed as Box) - .downcast() - .ok() - .map(|boxed| *boxed) + match self.map.entry(TypeId::of::()) { + Entry::Occupied(mut entry) => { + let entry = entry.get_mut(); + let mut ptr = entry.ptr.typed::(); + if entry.ptr.live() { + Some(mem::replace(unsafe { ptr.as_mut() }, Aligned(val)).0) + } else { + unsafe { + ptr::write(ptr.as_ptr(), Aligned(val)); + } + entry.ptr.set_live(true); + None } - }) + } + Entry::Vacant(entry) => { + entry.insert(BoxedEntry::new(val)); + None + } + } } /// Get a reference to a type previously inserted on this `Extensions`. pub(crate) fn get(&self) -> Option<&T> { self.map .get(&TypeId::of::()) - .and_then(|boxed| (&**boxed as &(dyn Any + 'static)).downcast_ref()) + .filter(|boxed| boxed.ptr.live()) + .map(|boxed| &unsafe { boxed.ptr.typed::().as_ref() }.0) } /// Get a mutable reference to a type previously inserted on this `Extensions`. pub(crate) fn get_mut(&mut self) -> Option<&mut T> { self.map .get_mut(&TypeId::of::()) - .and_then(|boxed| (&mut **boxed as &mut (dyn Any + 'static)).downcast_mut()) + .filter(|boxed| boxed.ptr.live()) + .map(|boxed| &mut unsafe { boxed.ptr.typed::().as_mut() }.0) } /// Remove a type from this `Extensions`. /// /// If a extension of this type existed, it will be returned. pub(crate) fn remove(&mut self) -> Option { - self.map.remove(&TypeId::of::()).and_then(|boxed| { - #[allow(warnings)] - { - (boxed as Box) - .downcast() - .ok() - .map(|boxed| *boxed) - } - }) + self.map + .get_mut(&TypeId::of::()) + .filter(|boxed| boxed.ptr.live()) + .map(|boxed| { + boxed.ptr.set_live(false); + unsafe { ptr::read(boxed.ptr.typed::().as_ptr()) }.0 + }) } /// Clear the `ExtensionsInner` in-place, dropping any elements in the map but @@ -184,7 +311,12 @@ impl ExtensionsInner { /// that future spans will not need to allocate new hashmaps. #[cfg(any(test, feature = "registry"))] pub(crate) fn clear(&mut self) { - self.map.clear(); + for boxed in self.map.values_mut().filter(|boxed| boxed.ptr.live()) { + boxed.ptr.set_live(false); + unsafe { + boxed.op_fn()(boxed.ptr, BoxedEntryOp::DropLiveValue); + } + } } } @@ -234,8 +366,8 @@ mod tests { assert_eq!( extensions.map.len(), - 0, - "after clear(), extensions map should have length 0" + 3, + "after clear(), extensions map should have length 3" ); assert_eq!( extensions.map.capacity(),