From eb847d6db58dc75f1a607b1e954b44851096ef13 Mon Sep 17 00:00:00 2001 From: Alex Saveau Date: Sat, 6 Jan 2024 12:18:55 -0800 Subject: [PATCH] Tighten boxed_entry abstractions Signed-off-by: Alex Saveau --- tracing-subscriber/src/registry/extensions.rs | 98 +++++++++++-------- 1 file changed, 56 insertions(+), 42 deletions(-) diff --git a/tracing-subscriber/src/registry/extensions.rs b/tracing-subscriber/src/registry/extensions.rs index 4c47a57010..98bbb57542 100644 --- a/tracing-subscriber/src/registry/extensions.rs +++ b/tracing-subscriber/src/registry/extensions.rs @@ -5,27 +5,26 @@ use std::{ collections::{hash_map::Entry, HashMap}, fmt, hash::{BuildHasherDefault, Hasher}, - mem, ptr, }; -use crate::registry::extensions::boxed_entry::{Aligned, BoxedEntry, BoxedEntryOp}; +use crate::registry::extensions::boxed_entry::BoxedEntry; use crate::sync::{RwLockReadGuard, RwLockWriteGuard}; #[allow(unreachable_pub)] mod boxed_entry { use std::mem::MaybeUninit; - use std::ptr; use std::ptr::NonNull; + use std::{mem, ptr}; /// Used to give ourselves one free bit for liveness checks. #[repr(align(2))] - pub struct Aligned(pub T); + 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<()>); + struct ExtPtr(NonNull<()>); impl ExtPtr { fn new(val: T) -> Self { @@ -37,7 +36,7 @@ mod boxed_entry { this } - pub fn typed(self) -> NonNull> { + fn typed(self) -> NonNull> { #[cfg(not(nightly))] let ptr = { let addr = self.0.as_ptr() as usize; @@ -49,7 +48,7 @@ mod boxed_entry { unsafe { NonNull::new_unchecked(ptr) } } - pub fn live(&self) -> bool { + fn live(&self) -> bool { #[cfg(not(nightly))] let addr = self.0.as_ptr() as usize; #[cfg(nightly)] @@ -57,7 +56,7 @@ mod boxed_entry { (addr & 1) == 1 } - pub fn set_live(&mut self, live: bool) { + 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 _; @@ -68,18 +67,16 @@ mod boxed_entry { } /// Extension storage - #[allow(clippy::manual_non_exhaustive)] pub struct BoxedEntry { - pub ptr: ExtPtr, - pub op_fn: unsafe fn(ExtPtr, BoxedEntryOp), - _private: (), + ptr: ExtPtr, + op_fn: unsafe fn(ExtPtr, BoxedEntryOp), } unsafe impl Send for BoxedEntry {} unsafe impl Sync for BoxedEntry {} - pub enum BoxedEntryOp { + enum BoxedEntryOp { DropLiveValue, Drop, } @@ -109,7 +106,47 @@ mod boxed_entry { Self { ptr: ExtPtr::new(val), op_fn: run_boxed_entry_op::, - _private: (), + } + } + + pub fn insert(&mut self, val: T) -> Option { + let mut ptr = self.ptr.typed::(); + if self.ptr.live() { + Some(mem::replace(unsafe { ptr.as_mut() }, Aligned(val)).0) + } else { + unsafe { + ptr::write(ptr.as_ptr(), Aligned(val)); + } + self.ptr.set_live(true); + None + } + } + + pub fn as_ref(&self) -> Option<&T> { + self.ptr + .live() + .then(|| &unsafe { self.ptr.typed::().as_ref() }.0) + } + + pub fn as_mut(&mut self) -> Option<&mut T> { + self.ptr + .live() + .then(|| &mut unsafe { self.ptr.typed::().as_mut() }.0) + } + + pub fn remove(&mut self) -> Option { + self.ptr.live().then(|| { + self.ptr.set_live(false); + unsafe { ptr::read(self.ptr.typed::().as_ptr()) }.0 + }) + } + + pub fn clear(&mut self) { + if self.ptr.live() { + self.ptr.set_live(false); + unsafe { + (self.op_fn)(self.ptr, BoxedEntryOp::DropLiveValue); + } } } } @@ -251,19 +288,7 @@ impl ExtensionsInner { /// be returned. pub(crate) fn insert(&mut self, val: T) -> Option { 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::Occupied(mut entry) => entry.get_mut().insert::(val), Entry::Vacant(entry) => { entry.insert(BoxedEntry::new(val)); None @@ -275,16 +300,14 @@ impl ExtensionsInner { pub(crate) fn get(&self) -> Option<&T> { self.map .get(&TypeId::of::()) - .filter(|boxed| boxed.ptr.live()) - .map(|boxed| &unsafe { boxed.ptr.typed::().as_ref() }.0) + .and_then(BoxedEntry::as_ref::) } /// 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::()) - .filter(|boxed| boxed.ptr.live()) - .map(|boxed| &mut unsafe { boxed.ptr.typed::().as_mut() }.0) + .and_then(BoxedEntry::as_mut::) } /// Remove a type from this `Extensions`. @@ -293,11 +316,7 @@ impl ExtensionsInner { pub(crate) fn remove(&mut self) -> Option { 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 - }) + .and_then(BoxedEntry::remove::) } /// Clear the `ExtensionsInner` in-place, dropping any elements in the map but @@ -307,12 +326,7 @@ impl ExtensionsInner { /// that future spans will not need to allocate new hashmaps. #[cfg(any(test, feature = "registry"))] pub(crate) fn clear(&mut self) { - 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); - } - } + self.map.values_mut().for_each(BoxedEntry::clear); } }