From 2c28cc9683a69e342ddfe1aa24b90001f6203b98 Mon Sep 17 00:00:00 2001 From: TalBarYakar Date: Wed, 23 Sep 2026 15:20:16 +0300 Subject: [PATCH 1/3] use unique pointers as hash set entries --- src/lib.rs | 5 +- src/unsafe_string.rs | 391 ++++++++++++++++++++++++++++++++----------- src/value.rs | 5 +- 3 files changed, 298 insertions(+), 103 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index b7f2335e..d17f44f9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -77,9 +77,8 @@ pub trait Defrag { /// Defrag implementation fn defrag(self, defrag_allocator: &mut A) -> Self; } -/// Reinitialized the shared strings cache. -/// Any json that still uses a shared string will continue using it. -/// But new strings will be reinitialized instead of reused the old ones. +/// Shrinks the shared string cache while preserving live entries and ownership. +/// Existing strings remain interned; shared allocations stay pinned during defrag. pub fn reinit_shared_string_cache() { unsafe_string::reinit_cache(); } diff --git a/src/unsafe_string.rs b/src/unsafe_string.rs index 4e46729a..625afe39 100644 --- a/src/unsafe_string.rs +++ b/src/unsafe_string.rs @@ -3,13 +3,14 @@ use hashbrown::HashSet; use std::alloc::{alloc, dealloc, Layout, LayoutError}; use std::borrow::Borrow; +use std::cell::{Cell, UnsafeCell}; use std::cmp::Ordering; use std::fmt::{self, Debug, Formatter}; use std::hash::Hash; -use std::mem::{self, transmute}; +use std::mem::transmute; use std::ops::Deref; use std::ptr::{addr_of_mut, copy_nonoverlapping, NonNull}; -use std::sync::{atomic::AtomicU32, Mutex, MutexGuard, OnceLock}; +use std::sync::{Mutex, MutexGuard, OnceLock}; use crate::{ thin::{ThinMut, ThinMutExt, ThinRef, ThinRefExt}, @@ -17,9 +18,8 @@ use crate::{ Defrag, DefragAllocator, IValue, }; -#[repr(align(8))] +#[repr(C)] struct Header { - rc: AtomicU32, // We use 32 bits for the length, which allows up to 4 GiB (safely covers 512MB) len: u32, } @@ -62,21 +62,14 @@ fn can_inline_string(s: &str) -> bool { enum StringCache { ThreadSafe(Mutex>), - ThreadUnsafe(HashSet), + ThreadUnsafe(UnsafeCell>), } static mut STRING_CACHE: OnceLock = OnceLock::new(); pub(crate) fn reinit_cache() { - let s_c = get_cache_mut(); - match s_c { - StringCache::ThreadUnsafe(s_c) => *s_c = HashSet::new(), - StringCache::ThreadSafe(s_c) => { - let mut s_c: std::sync::MutexGuard<'_, HashSet> = - s_c.lock().expect("Mutex lock should succeed"); - *s_c = HashSet::new(); - } - } + // The cache now owns reference counts; live entries must survive a reset. + get_cache_guard().shrink(); } pub(crate) fn init_cache(thread_safe: bool) -> Result<(), String> { @@ -84,22 +77,16 @@ pub(crate) fn init_cache(thread_safe: bool) -> Result<(), String> { s_c.set(if thread_safe { StringCache::ThreadSafe(Mutex::new(HashSet::new())) } else { - StringCache::ThreadUnsafe(HashSet::new()) + StringCache::ThreadUnsafe(UnsafeCell::new(HashSet::new())) }) .map_err(|_| "Cache is already initialized".to_owned()) } -fn get_cache_mut() -> &'static mut StringCache { - let s_c = unsafe { &mut *addr_of_mut!(STRING_CACHE) }; - s_c.get_or_init(|| StringCache::ThreadUnsafe(HashSet::new())); - s_c.get_mut().unwrap() -} - -fn is_thread_safe() -> bool { - match get_cache_mut() { - StringCache::ThreadSafe(_) => true, - StringCache::ThreadUnsafe(_) => false, - } +fn get_cache() -> &'static StringCache { + // SAFETY: the static is never assigned after initialization. OnceLock and + // Mutex synchronize the thread-safe path without aliasing mutable references. + let s_c = unsafe { &*addr_of_mut!(STRING_CACHE) }; + s_c.get_or_init(|| StringCache::ThreadUnsafe(UnsafeCell::new(HashSet::new()))) } enum CacheGuard { @@ -141,7 +128,6 @@ impl CacheGuard { } } - #[cfg(test)] fn shrink(&mut self) { match self { CacheGuard::ThreadSafe(c_g) => c_g.shrink_to_fit(), @@ -151,17 +137,28 @@ impl CacheGuard { } fn get_cache_guard() -> CacheGuard { - let s_c = get_cache_mut(); + let s_c = get_cache(); match s_c { - StringCache::ThreadUnsafe(s_c) => CacheGuard::ThreadUnsafe(s_c), + StringCache::ThreadUnsafe(s_c) => { + // SAFETY: callers selecting this mode must externally serialize all + // cache operations, as required by init_shared_string_cache(false). + CacheGuard::ThreadUnsafe(unsafe { &mut *s_c.get() }) + } StringCache::ThreadSafe(s_c) => { CacheGuard::ThreadSafe(s_c.lock().expect("Mutex lock should succeed")) } } } +struct SharedHeader { + data: NonNull
, + rc: Cell, +} + struct WeakIString { - ptr: NonNull
, + // Low bit set: SharedHeader; otherwise a unique Header. Both are aligned. + // Mutated only with the cache guard held; string content/hash never changes. + ptr: Cell>, } impl PartialEq for WeakIString { @@ -189,34 +186,92 @@ impl Borrow for WeakIString { } } impl WeakIString { - fn header(&self) -> ThinMut<'_, Header> { - // Safety: pointer is always valid - unsafe { ThinMut::new(self.ptr.as_ptr()) } + fn shared(&self) -> Option<&SharedHeader> { + let ptr = self.ptr.get().as_ptr(); + if ptr.addr() & 1 == 0 { + None + } else { + // SAFETY: the tagged pointer owns a live Box; the cache + // guard prevents promotion/demotion while this reference is used. + Some(unsafe { &*ptr.with_addr(ptr.addr() & !1).cast::() }) + } } - fn upgrade(&self) -> IString { + + fn data(&self) -> NonNull
{ + self.shared() + .map_or_else(|| self.ptr.get().cast(), |s| s.data) + } + + fn header(&self) -> ThinRef<'_, Header> { + // SAFETY: the cache entry always refers to a live immutable string. + unsafe { ThinRef::new(self.data().as_ptr()) } + } + + fn value(&self) -> IString { + // SAFETY: the caller has registered ownership of this aligned allocation. unsafe { - self.header() - .rc - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); IString(IValue::new_ptr( - self.ptr.as_ptr().cast::(), + self.data().as_ptr().cast(), TypeTag::StringOrNull, )) } } + + fn upgrade(&self) -> IString { + if let Some(shared) = self.shared() { + shared.rc.set( + shared + .rc + .get() + .checked_add(1) + .expect("string reference count overflow"), + ); + } else { + let shared = Box::new(SharedHeader { + data: self.data(), + rc: Cell::new(2), + }); + let ptr = Box::into_raw(shared).cast::(); + // SAFETY: Box pointers are non-null; bit 0 is free by alignment. + self.ptr + .set(unsafe { NonNull::new_unchecked(ptr.with_addr(ptr.addr() | 1)) }); + } + self.value() + } + + // Returns true when the departing owner was the last one. + fn release(&self) -> bool { + let Some(shared) = self.shared() else { + return true; + }; + let remaining = shared.rc.get() - 1; + shared.rc.set(remaining); + if remaining == 1 { + let ptr = self.ptr.get().as_ptr(); + self.ptr.set(shared.data.cast()); + // SAFETY: no other code accesses the descriptor under this guard. + // The string remains alive, and the table entry is now unique. + unsafe { + drop(Box::from_raw( + ptr.with_addr(ptr.addr() & !1).cast::(), + )); + } + } + false + } } /// The `IString` type is an interned, immutable string, and is where this crate /// gets its name. /// -/// Cloning an `IString` is cheap, and it can be easily converted from `&str` or -/// `String` types. Comparisons between `IString`s is a simple pointer +/// Cloning an `IString` looks up its ownership in the cache. It can be converted +/// from `&str` or `String`. Comparisons between `IString`s use a pointer /// comparison. /// /// The memory backing an `IString` is reference counted, so that unlike many /// string interning libraries, memory is not leaked as new strings are interned. -/// Interning uses `DashSet`, an implementation of a concurrent hash-set, allowing -/// many strings to be interned concurrently without becoming a bottleneck. +/// One hash-set stores unique pointers or shared descriptors. In thread-safe +/// mode, interning, cloning and dropping take the cache mutex. /// /// Given the nature of `IString` it is better to intern a string once and reuse /// it, rather than continually convert from `&str` to `IString`. @@ -226,16 +281,16 @@ pub struct IString(pub(crate) IValue); value_subtype_impls!(IString, into_string, as_string, as_string_mut); -static EMPTY_HEADER: Header = Header { - len: 0, - rc: AtomicU32::new(0), -}; +#[repr(align(8))] +struct EmptyHeader(Header); +static EMPTY_HEADER: EmptyHeader = EmptyHeader(Header { len: 0 }); impl IString { fn layout(len: usize) -> Result { Ok(Layout::new::
() .extend(Layout::array::(len)?)? .0 + .align_to(ALIGNMENT)? .pad_to_align()) } @@ -246,9 +301,11 @@ impl IString { Self::layout(s.len()).expect("layout is expected to return a valid value"), ) .cast::
(); + let ptr = NonNull::new(ptr) + .unwrap_or_else(|| std::alloc::handle_alloc_error(Self::layout(s.len()).unwrap())) + .as_ptr(); ptr.write(Header { len: s.len() as u32, - rc: AtomicU32::new(0), }); let hd = ThinMut::new(ptr); copy_nonoverlapping(s.as_ptr(), hd.str_ptr_mut(), s.len()); @@ -270,14 +327,20 @@ impl IString { } let mut cache = get_cache_guard(); + if let Some(existing) = cache.get_val(s) { + return existing.upgrade(); + } let k = cache.get_or_insert( s, Box::new(|s| WeakIString { - ptr: unsafe { NonNull::new_unchecked(Self::alloc(s, allocator)) }, + ptr: Cell::new( + NonNull::new(Self::alloc(s, allocator).cast()) + .expect("string allocation failed"), + ), }), ); - k.upgrade() + k.value() } /// Create an inline string by storing bytes in upper bits @@ -311,8 +374,9 @@ impl IString { (self.0.ptr_usize() % ALIGNMENT) == TypeTag::InlineString as usize } - fn header(&self) -> ThinMut<'_, Header> { - unsafe { ThinMut::new(self.0.ptr().cast()) } + fn header(&self) -> ThinRef<'_, Header> { + // SAFETY: a non-inline string owns a live immutable header. + unsafe { ThinRef::new(self.0.ptr().cast()) } } /// Returns the length (in bytes) of this string. @@ -360,7 +424,7 @@ impl IString { /// Returns the empty string. #[must_use] pub fn new() -> Self { - unsafe { IString(IValue::new_ref(&EMPTY_HEADER, TypeTag::StringOrNull)) } + unsafe { IString(IValue::new_ref(&EMPTY_HEADER.0, TypeTag::StringOrNull)) } } pub(crate) fn clone_impl(&self) -> IValue { @@ -369,65 +433,52 @@ impl IString { } else if self.is_inline() { unsafe { self.0.raw_copy() } } else { - self.header() - .rc - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - unsafe { self.0.raw_copy() } + let cache = get_cache_guard(); + cache + .get_val(self.as_str()) + .expect("live string missing from cache") + .upgrade() + .0 } } - fn drop_impl_with_deallocator(&mut self, deallocator: D) { + pub(crate) fn drop_impl(&mut self) { if !self.is_empty() && !self.is_inline() { - let hd = self.header(); - - if is_thread_safe() { - // Optimization for the thread safe case, we want to avoid locking the cache if the ref count - // is not potentially going to reach zero. - let mut rc = hd.rc.load(std::sync::atomic::Ordering::Relaxed); - while rc > 1 { - match hd.rc.compare_exchange_weak( - rc, - rc - 1, - std::sync::atomic::Ordering::Relaxed, - std::sync::atomic::Ordering::Relaxed, - ) { - Ok(_) => return, - Err(new_rc) => rc = new_rc, - } - } - } - let mut cache = get_cache_guard(); - if hd.rc.fetch_sub(1, std::sync::atomic::Ordering::Relaxed) == 1 { - // Reference count reached zero, free the string - if let Some(element) = cache.get_val(hd.str()) { - // we can not simply remove the element from the cache, while we - // perform active defrag, the element might be in the cache but will - // point to another (newer) value. In this case we do not want to remove it. - if element.ptr.as_ptr().cast() == unsafe { self.0.ptr() } { - cache.remove_val(hd.str()); - } - } + if cache + .get_val(self.as_str()) + .expect("live string missing from cache") + .release() + { + cache.remove_val(self.as_str()); // Shrink the cache if it is empty in tests to verify no memory leaks #[cfg(test)] if cache.check_if_empty() { cache.shrink(); } - Self::dealloc(unsafe { self.0.ptr().cast() }, deallocator); + // SAFETY: the cache entry was removed after its last owner departed. + Self::dealloc(unsafe { self.0.ptr().cast() }, |ptr, layout| unsafe { + dealloc(ptr, layout) + }); } } } - pub(crate) fn drop_impl(&mut self) { - self.drop_impl_with_deallocator(|ptr, layout| unsafe { dealloc(ptr, layout) }); - } - pub(crate) fn mem_allocated(&self) -> usize { if self.is_empty() || self.is_inline() { 0 } else { + let cache = get_cache_guard(); + let entry = cache + .get_val(self.as_str()) + .expect("live string missing from cache"); Self::layout(self.len()).unwrap().size() + + if entry.shared().is_some() { + std::mem::size_of::() + } else { + 0 + } } } } @@ -554,14 +605,28 @@ impl Debug for IString { impl Defrag for IString { fn defrag(mut self, defrag_allocator: &mut A) -> Self { - let new = Self::intern_with_allocator(self.as_str(), |layout| unsafe { - defrag_allocator.alloc(layout) - }); - self.drop_impl_with_deallocator(|ptr, layout| unsafe { - defrag_allocator.free(ptr, layout) - }); - mem::forget(self); - new + if self.is_empty() || self.is_inline() { + return self; + } + let cache = get_cache_guard(); + let entry = cache + .get_val(self.as_str()) + .expect("live string missing from cache"); + if entry.shared().is_none() { + // SAFETY: this is the sole owner, and the cache guard prevents a new + // owner. Relocation updates both the cache and this value together. + unsafe { + let ptr = + defrag_allocator.realloc_ptr(self.0.ptr(), Self::layout(self.len()).unwrap()); + entry + .ptr + .set(NonNull::new(ptr).expect("defrag allocation failed")); + self.0.set_ptr(ptr); + } + } + // ponytail: shared records stay pinned; moving them needs stable handles + // or a forwarding mechanism to preserve every owner's pointer. + self } } @@ -684,4 +749,132 @@ mod tests { assert_eq!(original.0.ptr_usize(), cloned.0.ptr_usize()); }); } + #[mockalloc::test] + fn unique_promotes_and_demotes_without_moving_bytes() { + assert_eq!(std::mem::size_of::
(), 4); + assert_eq!(std::mem::size_of::(), 8); + let first = IString::intern("Redis is very fast. "); // 20 bytes + let borrowed = first.as_str(); + let address = borrowed.as_ptr(); + assert_eq!(first.mem_allocated(), 24); + assert!(get_cache_guard() + .get_val(borrowed) + .unwrap() + .shared() + .is_none()); + + let second = IString::intern(borrowed); + let third = first.clone(); + assert_eq!(second.as_ptr(), address); + assert_eq!(third.as_ptr(), address); + assert_eq!( + first.mem_allocated(), + 24 + std::mem::size_of::() + ); + assert_eq!( + get_cache_guard() + .get_val(borrowed) + .unwrap() + .shared() + .unwrap() + .rc + .get(), + 3 + ); + reinit_cache(); + drop(first); + assert_eq!(second.as_str(), "Redis is very fast. "); + drop(third); + assert!(get_cache_guard() + .get_val(second.as_str()) + .unwrap() + .shared() + .is_none()); + assert_eq!(second.mem_allocated(), 24); + let promoted_again = second.clone(); + assert_eq!(promoted_again.as_ptr(), address); + drop(second); + assert_eq!(promoted_again.as_str(), "Redis is very fast. "); + drop(promoted_again); + assert!(get_cache_guard().get_val("Redis is very fast. ").is_none()); + } + + #[mockalloc::test] + fn defrag_moves_unique_but_preserves_shared_pointers() { + struct MovingAllocator(usize); + impl DefragAllocator for MovingAllocator { + unsafe fn realloc_ptr(&mut self, ptr: *mut T, layout: Layout) -> *mut T { + self.0 += 1; + // SAFETY: the caller supplies a live allocation with this layout. + unsafe { + let new = self.alloc(layout); + copy_nonoverlapping(ptr.cast(), new, layout.size()); + self.free(ptr, layout); + new.cast() + } + } + unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 { + // SAFETY: layout comes from the live string allocation. + unsafe { + NonNull::new(alloc(layout)) + .unwrap_or_else(|| std::alloc::handle_alloc_error(layout)) + .as_ptr() + } + } + unsafe fn free(&mut self, ptr: *mut T, layout: Layout) { + // SAFETY: the caller has finished using this allocation. + unsafe { + dealloc(ptr.cast(), layout); + } + } + } + let mut allocator = MovingAllocator(0); + for text in ["", "short", "שלום עולם 🌍", "long\0string\nwith controls"] { + let original = IString::intern(text); + let relocated = original.defrag(&mut allocator); + assert_eq!(relocated.as_str(), text); + let shared = relocated.clone(); + let address = shared.as_ptr(); + let moves = allocator.0; + let relocated = relocated.defrag(&mut allocator); + assert_eq!(allocator.0, moves); + if text.len() > 7 { + assert_eq!(relocated.as_ptr(), address); + } + assert_eq!(shared.as_str(), text); + } + assert_eq!(allocator.0, 2); + } + + // Run in its own test process because cache initialization is process-wide. + #[test] + #[ignore = "run separately with --ignored --exact"] + fn concurrent_promotion() { + init_cache(true).unwrap(); + let keep = IString::intern("concurrently shared string"); + let address = keep.as_ptr() as usize; + let threads: Vec<_> = (0..4) + .map(|_| { + std::thread::spawn(move || { + for _ in 0..1000 { + let value = IString::intern("concurrently shared string"); + let cloned = value.clone(); + assert_eq!(cloned.as_ptr() as usize, address); + drop(value); + assert_eq!(cloned.as_str(), "concurrently shared string"); + } + }) + }) + .collect(); + for thread in threads { + thread.join().unwrap(); + } + assert!(get_cache_guard() + .get_val(keep.as_str()) + .unwrap() + .shared() + .is_none()); + drop(keep); + assert!(get_cache_guard().check_if_empty()); + } } diff --git a/src/value.rs b/src/value.rs index 8b99edac..c9cbabac 100644 --- a/src/value.rs +++ b/src/value.rs @@ -1236,7 +1236,10 @@ mod tests { assert!(matches!(x.clone().destructure(), Destructured::String(u) if u == s)); assert!(matches!(x.clone().destructure_ref(), DestructuredRef::String(u) if *u == s)); assert!(matches!(x.clone().destructure_mut(), DestructuredMut::String(u) if *u == s)); - assert_eq!(x.mem_allocated(), 24); + assert_eq!( + x.mem_allocated(), + if cfg!(feature = "thread_safe") { 24 } else { 16 } + ); } #[mockalloc::test] From 7f430fc5f3a916dc18ead41f4f95ce8667a17954 Mon Sep 17 00:00:00 2001 From: TalBarYakar Date: Thu, 24 Sep 2026 11:08:57 +0300 Subject: [PATCH 2/3] fmt fix --- src/value.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/value.rs b/src/value.rs index c9cbabac..fe0da239 100644 --- a/src/value.rs +++ b/src/value.rs @@ -1238,7 +1238,11 @@ mod tests { assert!(matches!(x.clone().destructure_mut(), DestructuredMut::String(u) if *u == s)); assert_eq!( x.mem_allocated(), - if cfg!(feature = "thread_safe") { 24 } else { 16 } + if cfg!(feature = "thread_safe") { + 24 + } else { + 16 + } ); } From 347962df2859be0e6622cd5bf823853ca1d20926 Mon Sep 17 00:00:00 2001 From: TalBarYakar Date: Thu, 24 Sep 2026 13:59:24 +0300 Subject: [PATCH 3/3] avoid double lookup for unexisting strings --- src/unsafe_string.rs | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/src/unsafe_string.rs b/src/unsafe_string.rs index 625afe39..f0429ccf 100644 --- a/src/unsafe_string.rs +++ b/src/unsafe_string.rs @@ -95,11 +95,7 @@ enum CacheGuard { } impl CacheGuard { - fn get_or_insert<'a>( - &mut self, - value: &str, - f: Box WeakIString + 'a>, - ) -> &WeakIString { + fn get_or_insert(&mut self, value: &str, f: impl FnOnce(&str) -> WeakIString) -> &WeakIString { match self { CacheGuard::ThreadSafe(c_g) => c_g.get_or_insert_with(value, |val| f(val)), CacheGuard::ThreadUnsafe(c_g) => c_g.get_or_insert_with(value, |val| f(val)), @@ -327,20 +323,21 @@ impl IString { } let mut cache = get_cache_guard(); - if let Some(existing) = cache.get_val(s) { - return existing.upgrade(); - } - - let k = cache.get_or_insert( - s, - Box::new(|s| WeakIString { + let mut inserted = false; + let entry = cache.get_or_insert(s, |s| { + inserted = true; + WeakIString { ptr: Cell::new( NonNull::new(Self::alloc(s, allocator).cast()) .expect("string allocation failed"), ), - }), - ); - k.value() + } + }); + if inserted { + entry.value() + } else { + entry.upgrade() + } } /// Create an inline string by storing bytes in upper bits