diff --git a/Cargo.toml b/Cargo.toml index 537d5c5..3d19ce0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,13 +1,24 @@ -[workspace] -members = [ - "lockmap-core", - "lockmap", - "lockmap-lru", -] -resolver = "2" - -[workspace.dependencies] +[package] +name = "lockmap" +version = "0.2.0" +edition = "2021" + +authors = ["SF-Zhou "] +homepage = "https://github.com/SF-Zhou/lockmap" +repository = "https://github.com/SF-Zhou/lockmap" +description = "A high-performance, thread-safe HashMap and LRU cache with fine-grained per-key locking." +license = "MIT OR Apache-2.0" + +[dependencies] aliasable = "0.1" -atomic-wait = "1.1" foldhash = "0.2" hashbrown = "0.15" +parking_lot = "0.12" + +[dev-dependencies] +criterion = "0" +rand = "0" + +[[bench]] +name = "bench_lockmap" +harness = false diff --git a/README.md b/README.md index 04894fb..ed0c873 100644 --- a/README.md +++ b/README.md @@ -6,17 +6,16 @@ [![Documentation](https://docs.rs/lockmap/badge.svg)](https://docs.rs/lockmap) [![FOSSA Status](https://app.fossa.com/api/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap.svg?type=shield)](https://app.fossa.com/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap?ref=badge_shield) -A family of high-performance, thread-safe concurrent map crates for Rust with **fine-grained per-key locking**. +A high-performance, thread-safe HashMap and LRU cache for Rust with **fine-grained per-key locking**. -## Workspace Crates +## Data Structures -| Crate | Description | -|-------|-------------| -| [`lockmap-core`](lockmap-core/) | Shared infrastructure: fast futex-based mutex and sharded map primitives | -| [`lockmap`](lockmap/) | Thread-safe HashMap with per-key level locking | -| [`lockmap-lru`](lockmap-lru/) | Thread-safe LRU cache with per-key locking and automatic capacity-based eviction | +| Type | Description | +|------|-------------| +| [`LockMap`](https://docs.rs/lockmap/latest/lockmap/struct.LockMap.html) | Thread-safe HashMap with per-key level locking | +| [`LruLockMap`](https://docs.rs/lockmap/latest/lockmap/struct.LruLockMap.html) | Thread-safe LRU cache with per-key locking and automatic capacity-based eviction | -## lockmap +## LockMap **LockMap** is a high-performance, thread-safe HashMap implementation that provides **fine-grained locking at the key level**. @@ -27,8 +26,9 @@ Unlike standard concurrent maps that might lock the entire map or large buckets, * **Key-Level Locking**: Acquire exclusive locks for specific keys. Operations on different keys run in parallel. * **Sharding Architecture**: Internal sharding reduces contention on the map structure itself during insertions and removals. * **Deadlock Prevention**: Provides `batch_lock` to safely acquire locks on multiple keys simultaneously using a deterministic order. -* **Efficient Waiting**: Uses a hybrid spin-then-park Futex implementation for low-overhead locking. -* **Entry API**: Ergonomic RAII guards (`EntryByVal`, `EntryByRef`) for managing locks. +* **Single Hash Computation**: Each key is hashed once; the pre-computed hash is stored alongside the key and reused for shard selection, table probing, and rehashing. +* **No Key Duplication**: Uses `hashbrown::HashTable` so each key is stored only once, inside the entry state. +* **Entry API**: Ergonomic unified RAII guard (`Entry`) for managing locks. ### Usage @@ -69,19 +69,19 @@ keys.insert("key2".to_string()); // `locked_entries` holds all the locks let mut locked_entries = map.batch_lock::>(keys); -if let Some(mut entry) = locked_entries.get_mut("key1") { +if let Some(entry) = locked_entries.get_mut("key1") { entry.insert("updated_in_batch".into()); } // All locks released when `locked_entries` is dropped ``` -## lockmap-lru +## LruLockMap **LruLockMap** extends the per-key locking design with **LRU (Least Recently Used) eviction**. Each internal shard maintains its own LRU ordering via an intrusive doubly-linked list, ensuring that eviction decisions are local and lock-free from other shards. ### Features -* **Per-Key Locking**: Same fine-grained locking as `lockmap`. +* **Per-Key Locking**: Same fine-grained locking as `LockMap`. * **Per-Shard LRU Eviction**: Each shard independently tracks access order and evicts least recently used entries when capacity is exceeded. * **Non-Blocking Eviction**: In-use entries are skipped during eviction; traversal continues to the next candidate, ensuring progress even when the tail is held. * **Intrusive Linked List**: LRU bookkeeping uses pointers embedded directly in each entry, avoiding extra allocations. @@ -91,7 +91,7 @@ if let Some(mut entry) = locked_entries.get_mut("key1") { ### Usage ```rust -use lockmap_lru::LruLockMap; +use lockmap::LruLockMap; // Create a cache with capacity 1000 let cache = LruLockMap::::new(1000); @@ -125,6 +125,31 @@ Unlike `std::sync::Mutex`, **this library does not implement lock poisoning**. I The `map.get(key)` method clones the value while holding an internal shard lock. > **Note**: If your value type `V` is expensive to clone (e.g., deep copy of large structures), or if `clone()` acquires other locks, use `map.entry(key).get()` instead. This moves the clone operation outside the internal map lock, preventing blocking of other threads accessing the same shard. +## Changelog + +### 0.2.0 + +This is a major restructuring release. The previous workspace of three crates (`lockmap`, `lockmap-lru`, `lockmap-core`) has been consolidated into a single `lockmap` crate. + +#### Breaking Changes + +* **Unified crate**: `lockmap-lru` and `lockmap-core` no longer exist as separate crates. Import both `LockMap` and `LruLockMap` from `lockmap`: + ```rust + use lockmap::{LockMap, LruLockMap}; + ``` +* **Unified `Entry` type**: The separate `EntryByVal` and `EntryByRef` types in `LockMap` have been replaced by a single [`Entry`](https://docs.rs/lockmap/latest/lockmap/struct.Entry.html) type. The key is now stored inside the entry state and accessed via `entry.key()`. +* **`parking_lot` mutex**: The custom futex-based mutex and `atomic-wait` dependency have been replaced by `parking_lot::RawMutex`. +* **`HashTable` storage for `LockMap`**: `LockMap` now uses `hashbrown::HashTable` (like `LruLockMap`) with the key and pre-computed hash stored in the entry state. Each operation hashes the key only once. + +#### Migration Guide + +| 0.1.x | 0.2.0 | +|-------|-------| +| `use lockmap::EntryByVal` | `use lockmap::Entry` | +| `use lockmap::EntryByRef` | `use lockmap::Entry` | +| `use lockmap_lru::LruLockMap` | `use lockmap::LruLockMap` | +| `use lockmap_lru::LruEntry` | `use lockmap::LruEntry` | + ## License [![FOSSA Status](https://app.fossa.com/api/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap.svg?type=large)](https://app.fossa.com/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap?ref=badge_large) diff --git a/lockmap/benches/bench_lockmap.rs b/benches/bench_lockmap.rs similarity index 100% rename from lockmap/benches/bench_lockmap.rs rename to benches/bench_lockmap.rs diff --git a/lockmap-core/Cargo.toml b/lockmap-core/Cargo.toml deleted file mode 100644 index 8c33958..0000000 --- a/lockmap-core/Cargo.toml +++ /dev/null @@ -1,14 +0,0 @@ -[package] -name = "lockmap-core" -version = "0.1.0" -edition = "2021" - -authors = ["SF-Zhou "] -homepage = "https://github.com/SF-Zhou/lockmap" -repository = "https://github.com/SF-Zhou/lockmap" -description = "Core utilities for the lockmap family of crates, including a fast futex-based mutex and sharded map infrastructure." -license = "MIT OR Apache-2.0" - -[dependencies] -atomic-wait.workspace = true -foldhash.workspace = true diff --git a/lockmap-core/src/futex.rs b/lockmap-core/src/futex.rs deleted file mode 100644 index c15254c..0000000 --- a/lockmap-core/src/futex.rs +++ /dev/null @@ -1,264 +0,0 @@ -// Modified from https://github.com/rust-lang/rust/blob/master/library/std/src/sys/sync/mutex/futex.rs -// Copyright (c) The Rust Project Contributors - -//! Fast user-space mutex implementation using futex operations. -//! -//! This module provides a lightweight mutex implementation that optimizes for -//! the common case of low contention while gracefully handling high contention -//! scenarios through futex-based blocking. -//! -//! The implementation uses a three-state atomic variable to track lock state: -//! - `UNLOCKED`: No thread holds the lock -//! - `LOCKED`: One thread holds the lock, no others waiting -//! - `CONTENDED`: One thread holds the lock, others are waiting -//! -//! This design minimizes atomic operations in the fast path while ensuring -//! fair wakeup behavior when contention occurs. - -use std::sync::atomic::{ - AtomicU32, - Ordering::{Acquire, Relaxed, Release}, -}; - -/// A fast user-space mutex implementation using futex operations. -/// -/// This mutex is optimized for low contention scenarios by first attempting -/// to acquire the lock using atomic operations before falling back to -/// kernel-level futex operations when contention occurs. -/// -/// # Performance Characteristics -/// - Very fast in uncontended cases (single atomic operation) -/// - Efficiently handles contended cases using futex syscalls -/// - Minimal memory overhead (single 32-bit atomic) -/// -/// # Safety -/// This mutex does not provide poisoning support unlike `std::sync::Mutex`. -/// If a thread panics while holding the lock, the lock may remain in a locked -/// state permanently. -pub struct Mutex { - futex: AtomicU32, -} - -/// Represents an unlocked mutex state. -const UNLOCKED: u32 = 0; - -/// Represents a locked mutex with no waiting threads. -const LOCKED: u32 = 1; - -/// Represents a locked mutex with threads waiting for acquisition. -const CONTENDED: u32 = 2; - -impl Mutex { - /// Creates a new mutex in the unlocked state. - /// - /// # Returns - /// - /// A new `Mutex` instance ready for use. - #[inline] - pub const fn new() -> Self { - Self { - futex: AtomicU32::new(UNLOCKED), - } - } - - /// Attempts to acquire the lock without blocking. - /// - /// # Returns - /// - /// * `true` if the lock was successfully acquired - /// * `false` if the lock is currently held by another thread - #[inline] - pub fn try_lock(&self) -> bool { - self.futex - .compare_exchange(UNLOCKED, LOCKED, Acquire, Relaxed) - .is_ok() - } - - /// Acquires the lock, blocking the current thread until it becomes available. - /// - /// This function will not return until the lock has been acquired. - /// - /// # Panics - /// - /// This function may panic if the current thread already holds the lock. - /// However, this is not guaranteed and should not be relied upon for correctness. - #[inline] - pub fn lock(&self) { - if !self.try_lock() { - self.lock_contended(); - } - } - - /// Handles lock acquisition when contention is detected. - /// - /// This is the slow path for lock acquisition, used when `try_lock` fails. - /// It first spins briefly to see if the lock becomes available quickly, - /// then falls back to using futex operations to block the thread. - #[cold] - fn lock_contended(&self) { - // Spin first to speed things up if the lock is released quickly. - let mut state = self.spin(); - - // If it's unlocked now, attempt to take the lock - // without marking it as contended. - if state == UNLOCKED { - match self - .futex - .compare_exchange(UNLOCKED, LOCKED, Acquire, Relaxed) - { - Ok(_) => return, // Locked! - Err(s) => state = s, - } - } - - loop { - // Put the lock in contended state. - // We avoid an unnecessary write if it is already set to CONTENDED, - // to be friendlier for the caches. - if state != CONTENDED && self.futex.swap(CONTENDED, Acquire) == UNLOCKED { - // We changed it from UNLOCKED to CONTENDED, so we just successfully locked it. - return; - } - - // Wait for the futex to change state, assuming it is still CONTENDED. - atomic_wait::wait(&self.futex, CONTENDED); - - // Spin again after waking up. - state = self.spin(); - } - } - - /// Performs a brief spin loop waiting for the lock to become available. - /// - /// This reduces the overhead of immediately falling back to futex operations - /// in cases where the lock is held only briefly. - /// - /// # Returns - /// - /// The current state of the futex after spinning. - fn spin(&self) -> u32 { - let mut spin = 100; - loop { - // We only use `load` (and not `swap` or `compare_exchange`) - // while spinning, to be easier on the caches. - let state = self.futex.load(Relaxed); - - // We stop spinning when the mutex is UNLOCKED, - // but also when it's CONTENDED. - if state != LOCKED || spin == 0 { - return state; - } - - std::hint::spin_loop(); - spin -= 1; - } - } - - /// Releases the lock. - /// - /// # Safety - /// - /// This function should only be called by the thread that currently holds the lock. - /// Calling this function when the lock is not held or from a different thread - /// may lead to undefined behavior. - /// - /// # Panics - /// - /// This function may panic if called when the lock is not held, but this - /// is not guaranteed and should not be relied upon for correctness. - #[inline] - pub fn unlock(&self) { - if self.futex.swap(UNLOCKED, Release) == CONTENDED { - // We only wake up one thread. When that thread locks the mutex, it - // will mark the mutex as CONTENDED (see lock_contended above), - // which makes sure that any other waiting threads will also be - // woken up eventually. - self.wake(); - } - } - - /// Wakes up one waiting thread. - /// - /// This is called when releasing a contended lock to notify waiting threads - /// that the lock is now available. - #[cold] - fn wake(&self) { - atomic_wait::wake_one(&self.futex); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Arc; - - #[test] - fn test_futex() { - let lock = Arc::new(Mutex::new()); - let current = Arc::new(AtomicU32::new(0)); - const N: usize = 8; - const M: usize = 1 << 20; - - let mut tasks = vec![]; - for _ in 0..N { - let lock = lock.clone(); - let current = current.clone(); - tasks.push(std::thread::spawn(move || { - for _ in 0..M { - lock.lock(); - assert_eq!(current.fetch_add(1, Acquire), 0); - current.fetch_sub(1, Acquire); - lock.unlock(); - } - })); - } - for task in tasks { - task.join().unwrap(); - } - } - - #[test] - fn test_concurrent() { - let lock = Arc::new(Mutex::new()); - let counter = Arc::new(AtomicU32::new(0)); - const THREAD_COUNT: usize = 4; - const ITERATIONS: usize = 10000; - - let mut handles = vec![]; - - // Spawn multiple threads that increment and decrement a shared counter - for _ in 0..THREAD_COUNT { - let lock = Arc::clone(&lock); - let counter = Arc::clone(&counter); - - handles.push(std::thread::spawn(move || { - for _ in 0..ITERATIONS { - // Lock and modify shared state - lock.lock(); - let value = counter.load(Relaxed); - std::thread::yield_now(); // Force a context switch to increase contention - counter.store(value + 1, Relaxed); - lock.unlock(); - - // Do some work without the lock - std::thread::yield_now(); - - // Lock and modify shared state again - lock.lock(); - let value = counter.load(Relaxed); - std::thread::yield_now(); // Force a context switch to increase contention - counter.store(value - 1, Relaxed); - lock.unlock(); - } - })); - } - - // Wait for all threads to complete - for handle in handles { - handle.join().unwrap(); - } - - // Verify the final counter value is 0 - assert_eq!(counter.load(Relaxed), 0); - } -} diff --git a/lockmap-core/src/lib.rs b/lockmap-core/src/lib.rs deleted file mode 100644 index 6f753e0..0000000 --- a/lockmap-core/src/lib.rs +++ /dev/null @@ -1,12 +0,0 @@ -//! Core utilities for the lockmap family of crates. -//! -//! This crate provides: -//! - A fast, futex-based [`Mutex`] for low-overhead synchronization -//! - [`ShardsMap`] and [`ShardMap`] for sharded concurrent map infrastructure -//! - [`SimpleAction`] and [`UpdateAction`] enums for map update operations - -mod futex; -mod shards_map; - -pub use futex::*; -pub use shards_map::*; diff --git a/lockmap-core/src/shards_map.rs b/lockmap-core/src/shards_map.rs deleted file mode 100644 index 97a78fd..0000000 --- a/lockmap-core/src/shards_map.rs +++ /dev/null @@ -1,454 +0,0 @@ -use foldhash::fast::RandomState; -use std::borrow::Borrow; -use std::collections::HashMap; -use std::hash::{BuildHasher, Hash}; -use std::sync::Mutex; - -/// Represents the action to be taken on a value during a simple update operation. -/// -/// This enum is used with `simple_update` methods that can only keep existing values -/// unchanged or remove them entirely, without the ability to replace them with new values. -pub enum SimpleAction { - /// Keep the current value unchanged. - Keep, - /// Remove the value from the map. - Remove, -} - -/// Represents the action to be taken on a value during an update operation. -/// -/// This enum provides full control over what happens to a value during an update, -/// allowing the value to be kept unchanged or replaced with a new value. -pub enum UpdateAction { - /// Keep the current value unchanged. - Keep, - /// Replace the current value with the provided new value. - Replace(V), -} - -/// A thread-safe hashmap shard. -/// -/// This struct wraps a `HashMap` protected by a `Mutex` to ensure thread safety. -#[derive(Debug)] -pub struct ShardMap { - /// The underlying hashmap protected by a `Mutex`. - map: Mutex>, -} - -impl ShardMap -where - K: Eq + Hash, -{ - /// Creates a new `ShardMap` with the specified initial capacity. - /// - /// # Arguments - /// - /// * `capacity` - The initial capacity of the hashmap. - /// - /// # Returns - /// - /// A new `ShardMap` instance. - pub fn with_capacity(capacity: usize) -> Self { - Self { - map: Mutex::new(HashMap::with_capacity_and_hasher( - capacity, - RandomState::default(), - )), - } - } - - /// Returns the number of elements in the shard. - /// - /// # Returns - /// - /// The number of key-value pairs currently stored in this shard. - pub fn len(&self) -> usize { - self.map.lock().unwrap().len() - } - - /// Returns `true` if the shard contains no elements. - /// - /// # Returns - /// - /// `true` if this shard is empty, `false` otherwise. - pub fn is_empty(&self) -> bool { - self.map.lock().unwrap().is_empty() - } - - /// Updates the value associated with a key using a function that can only keep or remove. - /// - /// This is a simpler version of `update` that doesn't support replacing values, - /// only keeping them unchanged or removing them entirely. - /// - /// # Arguments - /// - /// * `key` - The key to update - /// * `func` - A function that takes an `Option<&mut V>` and returns a tuple containing - /// the action to take (`SimpleAction::Keep` or `SimpleAction::Remove`) and a result value - /// - /// # Returns - /// - /// The result value returned by the provided function. - pub fn simple_update(&self, key: &Q, func: F) -> R - where - K: Borrow, - Q: Eq + Hash + ?Sized, - F: FnOnce(Option<&mut V>) -> (SimpleAction, R), - { - let mut map = self.map.lock().unwrap(); - let value = map.get_mut(key); - let has_value = value.is_some(); - let (action, ret) = func(value); - if has_value && matches!(action, SimpleAction::Remove) { - let _ = map.remove_entry(key); - } - ret - } - - /// Updates the value associated with the given key using the provided function. - /// - /// # Arguments - /// - /// * `key` - The key to update. - /// * `func` - A function that takes an `Option<&mut V>` and returns a tuple containing the action to take and the result. - /// - /// # Returns - /// - /// The result returned by the provided function. - pub fn update(&self, key: K, func: F) -> R - where - F: FnOnce(Option<&mut V>) -> (UpdateAction, R), - { - let mut map = self.map.lock().unwrap(); - match map.get_mut(&key) { - Some(value) => { - let (action, ret) = func(Some(value)); - match action { - UpdateAction::Keep => {} - UpdateAction::Replace(v) => { - *value = v; - } - } - ret - } - None => { - let (action, ret) = func(None); - match action { - UpdateAction::Keep => {} - UpdateAction::Replace(value) => { - map.insert(key, value); - } - } - ret - } - } - } - - /// Updates the value associated with a key using a reference to the key. - /// - /// This method is similar to `update` but takes a reference to the key instead of - /// owning it. If the key doesn't exist and needs to be inserted, it will be - /// created from the reference using the `From` trait. - /// - /// # Arguments - /// - /// * `key` - A reference to the key to update - /// * `func` - A function that takes an `Option<&mut V>` and returns a tuple containing - /// the action to take and a result value - /// - /// # Returns - /// - /// The result value returned by the provided function. - pub fn update_by_ref(&self, key: &Q, func: F) -> R - where - K: Borrow + for<'c> From<&'c Q>, - Q: Eq + Hash + ?Sized, - F: FnOnce(Option<&mut V>) -> (UpdateAction, R), - { - let mut map = self.map.lock().unwrap(); - match map.get_mut(key) { - Some(value) => { - let (action, ret) = func(Some(value)); - match action { - UpdateAction::Keep => {} - UpdateAction::Replace(v) => { - *value = v; - } - } - ret - } - None => { - let (action, ret) = func(None); - match action { - UpdateAction::Keep => {} - UpdateAction::Replace(value) => { - map.insert(key.into(), value); - } - } - ret - } - } - } -} - -/// A collection of `ShardMap` instances, providing sharded access to a hashmap. -pub struct ShardsMap { - /// The vector of `ShardMap` instances. - shards: Vec>, - hasher: RandomState, -} - -impl ShardsMap -where - K: Eq + Hash, -{ - /// Creates a new `ShardsMap` with the specified capacity and number of shards. - /// - /// # Arguments - /// - /// * `capacity` - The total initial capacity of the hashmap. - /// * `shard_amount` - The number of shards to create. - /// - /// # Returns - /// - /// A new `ShardsMap` instance. - pub fn with_capacity_and_shard_amount(capacity: usize, shard_amount: usize) -> Self { - let shard_capacity = capacity / shard_amount; - Self { - shards: (0..shard_amount) - .map(|_| ShardMap::with_capacity(shard_capacity)) - .collect::>(), - hasher: RandomState::default(), - } - } - - /// Returns the total number of elements across all shards. - /// - /// # Returns - /// - /// The sum of elements in all shards. - pub fn len(&self) -> usize { - self.shards.iter().map(|s| s.len()).sum() - } - - /// Returns `true` if all shards contain no elements. - /// - /// # Returns - /// - /// `true` if all shards are empty, `false` otherwise. - pub fn is_empty(&self) -> bool { - self.shards.iter().all(|s| s.is_empty()) - } - - /// Updates the value associated with the given key using the provided function. - /// - /// # Arguments - /// - /// * `key` - The key to update. - /// * `func` - A function that takes an `Option<&mut V>` and returns a tuple containing the action to take and the result. - /// - /// # Returns - /// - /// The result returned by the provided function. - pub fn simple_update(&self, key: &Q, func: F) -> R - where - K: Borrow, - Q: Eq + Hash + ?Sized, - F: FnOnce(Option<&mut V>) -> (SimpleAction, R), - { - self.shard(key).simple_update(key, func) - } - - /// Updates the value associated with the given key using the provided function. - /// - /// # Arguments - /// - /// * `key` - The key to update. - /// * `func` - A function that takes an `Option<&mut V>` and returns a tuple containing the action to take and the result. - /// - /// # Returns - /// - /// The result returned by the provided function. - pub fn update(&self, key: K, func: F) -> R - where - F: FnOnce(Option<&mut V>) -> (UpdateAction, R), - { - self.shard(&key).update(key, func) - } - - /// Updates the value associated with the given key using the provided function. - /// - /// # Arguments - /// - /// * `key` - The key to update. - /// * `func` - A function that takes an `Option<&mut V>` and returns a tuple containing the action to take and the result. - /// - /// # Returns - /// - /// The result returned by the provided function. - pub fn update_by_ref(&self, key: &Q, func: F) -> R - where - K: Borrow + for<'c> From<&'c Q>, - Q: Eq + Hash + ?Sized, - F: FnOnce(Option<&mut V>) -> (UpdateAction, R), - { - self.shard(key).update_by_ref(key, func) - } - - /// Gets the appropriate shard for a given key. - /// - /// Uses consistent hashing to distribute keys across shards, ensuring - /// that the same key always maps to the same shard. - /// - /// # Arguments - /// - /// * `key` - The key to hash and map to a shard - /// - /// # Returns - /// - /// A reference to the `ShardMap` that should contain this key. - #[inline(always)] - fn shard(&self, key: &Q) -> &ShardMap - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - let idx = self.hasher.hash_one(key) as usize % self.shards.len(); - &self.shards[idx] - } -} - -#[cfg(test)] -mod tests { - use std::sync::{ - atomic::{AtomicU32, Ordering}, - Arc, - }; - - use super::*; - - #[test] - fn test_shards_map() { - let shards_map = ShardsMap::::with_capacity_and_shard_amount(256, 16); - assert!(shards_map.is_empty()); - assert_eq!(shards_map.len(), 0); - shards_map.update(1, |v| { - assert_eq!(v, None); - (UpdateAction::Replace(1), ()) - }); - assert!(!shards_map.is_empty()); - assert_eq!(shards_map.len(), 1); - shards_map.update(2, |v| { - assert_eq!(v, None); - (UpdateAction::Keep, ()) - }); - shards_map.simple_update(&3, |v| { - assert_eq!(v, None); - (SimpleAction::Remove, ()) - }); - assert!(!shards_map.is_empty()); - assert_eq!(shards_map.len(), 1); - shards_map.update(1, |v| { - assert_eq!(v.cloned(), Some(1)); - (UpdateAction::Replace(2), ()) - }); - shards_map.update(1, |v| { - assert_eq!(v.cloned(), Some(2)); - (UpdateAction::Keep, ()) - }); - shards_map.simple_update(&1, |v| { - assert_eq!(v.cloned(), Some(2)); - (SimpleAction::Remove, ()) - }); - assert!(shards_map.is_empty()); - assert_eq!(shards_map.len(), 0); - shards_map.simple_update(&1, |v| { - assert_eq!(v, None); - (SimpleAction::Remove, ()) - }); - assert!(shards_map.is_empty()); - assert_eq!(shards_map.len(), 0); - } - - #[test] - fn test_shards_map_2() { - let shards_map = ShardsMap::::with_capacity_and_shard_amount(256, 16); - shards_map.update_by_ref("hello", |v| { - assert_eq!(v, None); - (UpdateAction::Replace("world".to_string()), ()) - }); - shards_map.update_by_ref("hello", |v| { - assert_eq!(v.unwrap(), "world"); - (UpdateAction::Replace("lockmap".to_string()), ()) - }); - shards_map.simple_update("hello", |v| { - assert_eq!(v, Some(&mut "lockmap".to_string())); - (SimpleAction::Remove, ()) - }); - shards_map.simple_update("hello", |v| { - assert_eq!(v, None); - (SimpleAction::Remove, ()) - }); - shards_map.update_by_ref("hello", |v| { - assert_eq!(v, None); - (UpdateAction::Replace("lockmap".to_string()), ()) - }); - shards_map.simple_update("hello", |v| { - assert_eq!(v.unwrap(), "lockmap"); - (SimpleAction::Remove, ()) - }); - shards_map.simple_update("hello", |v| { - assert_eq!(v, None); - (SimpleAction::Remove, ()) - }); - shards_map.update_by_ref("hello", |v| { - assert_eq!(v, None); - (UpdateAction::Keep, ()) - }); - shards_map.update_by_ref(&"hello".to_owned(), |v| { - assert_eq!(v, None); - (UpdateAction::Keep, ()) - }); - shards_map.simple_update("hello", |v| { - assert_eq!(v, None); - (SimpleAction::Keep, ()) - }); - } - - #[test] - fn test_shards_map_concurrent() { - let lock_map = Arc::new(ShardsMap::::with_capacity_and_shard_amount( - 256, 16, - )); - let current = Arc::new(AtomicU32::default()); - const N: usize = 1 << 12; - const M: usize = 8; - - lock_map.update(1, |_| (UpdateAction::Replace(0), ())); - - let threads = (0..M) - .map(|_| { - let lock_map = lock_map.clone(); - let current = current.clone(); - std::thread::spawn(move || { - for _ in 0..N { - lock_map.update(1, |v| { - let now = current.fetch_add(1, Ordering::AcqRel); - assert_eq!(now, 0); - *v.unwrap() += 1; - let now = current.fetch_sub(1, Ordering::AcqRel); - assert_eq!(now, 1); - (UpdateAction::Keep, ()) - }); - } - }) - }) - .collect::>(); - threads.into_iter().for_each(|t| t.join().unwrap()); - - assert_eq!( - lock_map.update(1, |v| (UpdateAction::Replace(0), *v.unwrap())), - N * M - ); - } -} diff --git a/lockmap-lru/Cargo.toml b/lockmap-lru/Cargo.toml deleted file mode 100644 index 4fd9e38..0000000 --- a/lockmap-lru/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[package] -name = "lockmap-lru" -version = "0.1.1" -edition = "2021" - -authors = ["SF-Zhou "] -homepage = "https://github.com/SF-Zhou/lockmap" -repository = "https://github.com/SF-Zhou/lockmap" -description = "A high-performance, thread-safe LRU cache built on lockmap's per-key locking architecture." -license = "MIT OR Apache-2.0" - -[dependencies] -aliasable.workspace = true -foldhash.workspace = true -hashbrown.workspace = true -lockmap-core = { version = "0.1.0", path = "../lockmap-core" } - -[dev-dependencies] -rand = "0" diff --git a/lockmap-lru/README.md b/lockmap-lru/README.md deleted file mode 100644 index bc27328..0000000 --- a/lockmap-lru/README.md +++ /dev/null @@ -1,55 +0,0 @@ -# lockmap-lru - -[![Crates.io](https://img.shields.io/crates/v/lockmap-lru.svg)](https://crates.io/crates/lockmap-lru) -[![Documentation](https://docs.rs/lockmap-lru/badge.svg)](https://docs.rs/lockmap-lru) - -**lockmap-lru** is a high-performance, thread-safe LRU cache for Rust, built on top of the [lockmap](https://crates.io/crates/lockmap) architecture. - -It provides **fine-grained per-key locking** combined with **automatic capacity-based eviction**. Each internal shard maintains its own LRU ordering via an intrusive doubly-linked list, ensuring that eviction decisions are local and lock-free from other shards. - -## Features - -* **Per-Key Locking**: Acquire exclusive locks for specific keys. Operations on different keys run in parallel. -* **Per-Shard LRU Eviction**: Each shard independently tracks access order and evicts the least recently used entries when capacity is exceeded. -* **Non-Blocking Eviction**: In-use entries are skipped during eviction; traversal continues to the next candidate, ensuring eviction always makes progress. -* **Intrusive Linked List**: LRU bookkeeping uses pointers embedded directly in each entry, avoiding extra allocations. -* **No Key Duplication**: Uses `hashbrown::HashTable` so each key is stored only once (inside the entry state), saving memory and avoiding extra clones. -* **Single Hash**: One `RandomState` hasher shared across all shards; each operation hashes the key once. Shard selection uses upper bits; the full hash is stored in each entry and reused for eviction and cleanup. -* **Single-Probe Lookups**: Uses `HashTable::entry` / `HashTable::find_entry` for find-or-insert / find-or-remove in a single probe, avoiding double lookups. -* **Entry API**: Ergonomic RAII guard (`LruEntry`) for managing locks. The key is obtained directly from the entry's internal state, eliminating redundant copies. - -## Usage - -```rust -use lockmap_lru::LruLockMap; - -// Create a cache with capacity 1000 -let cache = LruLockMap::::new(1000); - -// 1. Basic Insert -cache.insert_by_ref("key", "value".into()); - -// 2. Get a value (promotes it in the LRU list) -assert_eq!(cache.get("key"), Some("value".into())); - -// 3. Entry API: Exclusive access -{ - let mut entry = cache.entry_by_ref("key"); - assert_eq!(entry.get().as_deref(), Some("value")); - entry.insert("new_value".to_string()); -} // Lock released here - -// 4. Remove -assert_eq!(cache.remove("key"), Some("new_value".into())); -``` - -## LRU Eviction Details - -- The total capacity is divided evenly among shards. -- On every access (get, insert, entry, remove, contains_key), the accessed entry is promoted to the head of its shard's LRU list. -- When a shard's entry count exceeds its capacity (after an insert or entry creation), the least recently used entries are evicted from the tail. -- Entries currently held by an `LruEntry` guard are **skipped** during eviction. Traversal continues from tail towards head, evicting any eligible entries, so eviction always makes progress even when the tail entry is held by another thread. - -## License - -Licensed under either of [Apache License, Version 2.0](../LICENSE-APACHE) or [MIT License](../LICENSE-MIT) at your option. diff --git a/lockmap-lru/src/lib.rs b/lockmap-lru/src/lib.rs deleted file mode 100644 index ba7bda5..0000000 --- a/lockmap-lru/src/lib.rs +++ /dev/null @@ -1,58 +0,0 @@ -//! A high-performance, thread-safe LRU cache with per-key locking. -//! -//! # Overview -//! -//! `lockmap-lru` provides a concurrent LRU (Least Recently Used) cache built on -//! the [`lockmap`](https://crates.io/crates/lockmap) architecture. It combines -//! fine-grained per-key locking with automatic capacity-based eviction. -//! -//! Each shard maintains its own LRU list using an intrusive doubly-linked list -//! embedded in the entry state. The underlying storage uses -//! [`hashbrown::HashTable`] to avoid key duplication — the key and its -//! pre-computed hash live only inside the `State` node. A single `RandomState` -//! hasher is shared across all shards: high bits of the hash select the shard, -//! and the full hash is passed to the `HashTable`, eliminating redundant hashing. -//! Lookups use `HashTable::entry` / `HashTable::find_entry` for single-probe -//! find-or-insert / find-or-remove, avoiding double lookups. -//! -//! On every access, the accessed entry is promoted to the head of the list. -//! When a shard exceeds its capacity, the least recently used entries are evicted -//! from the tail. In-use entries (held by an [`LruEntry`] guard) are skipped and -//! eviction continues to the next candidate, ensuring progress even when the -//! tail entry is held by another thread. -//! -//! # Features -//! -//! - **Per-key locking**: Same fine-grained locking as `lockmap` -//! - **Per-shard LRU eviction**: Each shard independently manages its own LRU list -//! - **Non-blocking eviction**: Entries currently in use are skipped; eviction walks past them to evict other candidates -//! - **Intrusive linked list**: Zero-allocation LRU bookkeeping via pointers embedded in each entry -//! - **No key duplication**: Uses `hashbrown::HashTable` so the key is stored only once, inside the entry state -//! - **Single hash**: One `RandomState` hasher for the entire map; hash computed once per operation -//! -//! # Examples -//! -//! ``` -//! use lockmap_lru::LruLockMap; -//! -//! // Create an LRU cache with capacity 100 -//! let cache = LruLockMap::::new(100); -//! -//! // Basic operations -//! cache.insert("key1".to_string(), 42); -//! assert_eq!(cache.get("key1"), Some(42)); -//! -//! // Entry API for exclusive access -//! { -//! let mut entry = cache.entry("key2".to_string()); -//! entry.insert(123); -//! } -//! -//! // Remove a value -//! assert_eq!(cache.remove("key1"), Some(42)); -//! assert_eq!(cache.get("key1"), None); -//! ``` -#[doc = include_str!("../README.md")] -mod lru_lockmap; - -pub use lru_lockmap::*; diff --git a/lockmap/Cargo.toml b/lockmap/Cargo.toml deleted file mode 100644 index 588f77f..0000000 --- a/lockmap/Cargo.toml +++ /dev/null @@ -1,22 +0,0 @@ -[package] -name = "lockmap" -version = "0.1.17" -edition = "2021" - -authors = ["SF-Zhou "] -homepage = "https://github.com/SF-Zhou/lockmap" -repository = "https://github.com/SF-Zhou/lockmap" -description = "A high-performance, thread-safe HashMap implementation for Rust that provides fine-grained locking at the key level." -license = "MIT OR Apache-2.0" - -[dependencies] -aliasable.workspace = true -lockmap-core = { version = "0.1.0", path = "../lockmap-core" } - -[dev-dependencies] -criterion = "0" -rand = "0" - -[[bench]] -name = "bench_lockmap" -harness = false diff --git a/lockmap/README.md b/lockmap/README.md deleted file mode 100644 index 04b3442..0000000 --- a/lockmap/README.md +++ /dev/null @@ -1,80 +0,0 @@ -# lockmap - -[![Rust](https://github.com/SF-Zhou/lockmap/actions/workflows/rust.yml/badge.svg)](https://github.com/SF-Zhou/lockmap/actions/workflows/rust.yml) -[![codecov](https://codecov.io/gh/SF-Zhou/lockmap/graph/badge.svg?token=7U9JFC64U4)](https://codecov.io/gh/SF-Zhou/lockmap) -[![Crates.io](https://img.shields.io/crates/v/lockmap.svg)](https://crates.io/crates/lockmap) -[![Documentation](https://docs.rs/lockmap/badge.svg)](https://docs.rs/lockmap) -[![FOSSA Status](https://app.fossa.com/api/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap.svg?type=shield)](https://app.fossa.com/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap?ref=badge_shield) - -**LockMap** is a high-performance, thread-safe HashMap implementation for Rust that provides **fine-grained locking at the key level**. - -Unlike standard concurrent maps that might lock the entire map or large buckets, `LockMap` allows you to hold an exclusive lock on a specific key (including non-existent ones) for complex atomic operations, minimizing contention across different keys. - -## Features - -* **Key-Level Locking**: Acquire exclusive locks for specific keys. Operations on different keys run in parallel. -* **Sharding Architecture**: Internal sharding reduces contention on the map structure itself during insertions and removals. -* **Deadlock Prevention**: Provides `batch_lock` to safely acquire locks on multiple keys simultaneously using a deterministic order. -* **Efficient Waiting**: Uses a hybrid spin-then-park Futex implementation for low-overhead locking. -* **Entry API**: Ergonomic RAII guards (`EntryByVal`, `EntryByRef`) for managing locks. - -## Important Caveats - -### 1. No Lock Poisoning - -Unlike `std::sync::Mutex`, **this library does not implement lock poisoning**. If a thread panics while holding an `Entry`, the lock is released immediately (via Drop) to avoid deadlocks, but the data is **not** marked as poisoned. -> **Warning**: Users must ensure exception safety. If a panic occurs during a partial update, the data associated with that key may be left in an inconsistent state for subsequent readers. - -### 2. `get()` Performance - -The `map.get(key)` method clones the value while holding an internal shard lock. -> **Note**: If your value type `V` is expensive to clone (e.g., deep copy of large structures), or if `clone()` acquires other locks, use `map.entry(key).get()` instead. This moves the clone operation outside the internal map lock, preventing blocking of other threads accessing the same shard. - -## Usage - -```rust -use lockmap::LockMap; -use std::collections::BTreeSet; - -// Create a new lock map -let map = LockMap::::new(); - -// 1. Basic Insert -map.insert_by_ref("key", "value".into()); - -// 2. Get a value (Clones the value) -assert_eq!(map.get("key"), Some("value".into())); - -// 3. Entry API: Exclusive access (Read/Write) -// This locks ONLY "key", other threads can access "other_key" concurrently. -{ - let mut entry = map.entry_by_ref("key"); - - // Check value - assert_eq!(entry.get().as_deref(), Some("value")); - - // Update value atomically - entry.insert("new value".to_string()); -} // Lock is automatically released here - -// 4. Remove a value -assert_eq!(map.remove("key"), Some("new value".into())); - -// 5. Batch Locking (Deadlock safe) -// Acquires locks for multiple keys in a deterministic order. -let mut keys = BTreeSet::new(); -keys.insert("key1".to_string()); -keys.insert("key2".to_string()); - -// `locked_entries` holds all the locks -let mut locked_entries = map.batch_lock::>(keys); - -if let Some(mut entry) = locked_entries.get_mut("key1") { - entry.insert("updated_in_batch".into()); -} -// All locks released when `locked_entries` is dropped -``` - -## License - -[![FOSSA Status](https://app.fossa.com/api/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap.svg?type=large)](https://app.fossa.com/projects/git%2Bgithub.com%2FSF-Zhou%2Flockmap?ref=badge_large) diff --git a/lockmap/src/lib.rs b/lockmap/src/lib.rs deleted file mode 100644 index 6b4ab1c..0000000 --- a/lockmap/src/lib.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! A thread-safe hashmap implementation providing per-key level locking and atomic operations. -//! -//! # Overview -//! `lockmap` provides a concurrent hashmap implementation that allows fine-grained locking at the key level. -//! It uses internal sharding for better performance under high concurrency. -//! -//! # Features -//! - Thread-safe access with per-key locking -//! - Entry API for exclusive access to values -//! - Efficient concurrent operations through sharding -//! - Safe atomic updates -//! - No poisoning, the lock is released normally on panic -//! -//! # Examples -//! ``` -//! use lockmap::LockMap; -//! -//! let map = LockMap::::new(); -//! -//! // Basic operations -//! map.insert("key1".into(), 42); -//! assert_eq!(map.get("key1"), Some(42)); -//! -//! // Entry API for exclusive access -//! { -//! let mut entry = map.entry("key2".into()); -//! entry.get_mut().replace(123); -//! } -//! -//! // Remove a value -//! assert_eq!(map.remove("key1"), Some(42)); -//! assert_eq!(map.get("key1"), None); -//! ``` -#[doc = include_str!("../README.md")] -mod lockmap; - -pub use lockmap::*; diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..69aa2be --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,67 @@ +//! A high-performance, thread-safe HashMap and LRU cache with fine-grained per-key locking. +//! +//! # Overview +//! +//! This crate provides two concurrent map implementations: +//! +//! - [`LockMap`]: A thread-safe HashMap with per-key level locking +//! - [`LruLockMap`]: A thread-safe LRU cache with per-key locking and automatic eviction +//! +//! Both data structures use internal sharding for high concurrency and allow you to +//! hold an exclusive lock on a specific key for complex atomic operations, minimizing +//! contention across different keys. +//! +//! # Features +//! +//! - **Per-key locking**: Acquire exclusive locks for specific keys; operations on different keys run in parallel +//! - **Sharding architecture**: Internal sharding reduces contention on the map structure itself +//! - **Single hash computation**: Key and pre-computed hash stored together; each operation hashes once +//! - **No key duplication**: Uses [`hashbrown::HashTable`] so each key is stored only once +//! - **Deadlock prevention**: `LockMap` provides [`batch_lock`](LockMap::batch_lock) for safe multi-key locking +//! - **LRU eviction**: `LruLockMap` automatically evicts least recently used entries when capacity is exceeded +//! - **Non-blocking eviction**: In-use entries are skipped during eviction; traversal continues to the next candidate +//! +//! # Examples +//! +//! ## LockMap +//! +//! ``` +//! use lockmap::LockMap; +//! +//! let map = LockMap::::new(); +//! +//! map.insert("key".to_string(), 42); +//! assert_eq!(map.get("key"), Some(42)); +//! +//! { +//! let mut entry = map.entry("key2".to_string()); +//! entry.insert(123); +//! } +//! +//! assert_eq!(map.remove("key"), Some(42)); +//! ``` +//! +//! ## LruLockMap +//! +//! ``` +//! use lockmap::LruLockMap; +//! +//! let cache = LruLockMap::::new(1000); +//! +//! cache.insert("key".to_string(), 42); +//! assert_eq!(cache.get("key"), Some(42)); +//! +//! { +//! let mut entry = cache.entry("key2".to_string()); +//! entry.insert(123); +//! } +//! +//! assert_eq!(cache.remove("key"), Some(42)); +//! ``` + +#[doc = include_str!("../README.md")] +mod lockmap; +mod lru_lockmap; + +pub use lockmap::{Entry, LockMap}; +pub use lru_lockmap::{LruEntry, LruLockMap}; diff --git a/lockmap/src/lockmap.rs b/src/lockmap.rs similarity index 53% rename from lockmap/src/lockmap.rs rename to src/lockmap.rs index 0713415..6442717 100644 --- a/lockmap/src/lockmap.rs +++ b/src/lockmap.rs @@ -1,16 +1,20 @@ use aliasable::boxed::AliasableBox; -use lockmap_core::{Mutex, ShardsMap, SimpleAction, UpdateAction}; +use foldhash::fast::RandomState; +use hashbrown::hash_table::Entry as TableEntry; +use hashbrown::HashTable; +use parking_lot::lock_api::RawMutex as _; +use parking_lot::RawMutex; use std::borrow::Borrow; use std::cell::UnsafeCell; use std::collections::BTreeSet; -use std::hash::Hash; +use std::hash::{BuildHasher, Hash}; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::OnceLock; -/// Internal flags for `State`. -/// -/// This struct wraps an `u32` to store both the reference count and a "has value" flag. -/// The highest bit is used for the flag, and the remaining bits for the reference count. +// --------------------------------------------------------------------------- +// StateFlags +// --------------------------------------------------------------------------- + struct StateFlags(u32); impl StateFlags { @@ -34,25 +38,29 @@ impl StateFlags { } fn pending_cleanup(&self) -> bool { - self.0 == 0 // equal to `self.refcnt() == 0 && !self.has_value()` + self.0 == 0 } } -/// Internal state for a key-value pair in the `LockMap`. -/// -/// This type manages the stored value, the per-key lock, and a reference count -/// used for both synchronization optimization and memory management. -struct State { +// --------------------------------------------------------------------------- +// State – per-key state with key and pre-computed hash +// --------------------------------------------------------------------------- + +struct State { + key: K, + hash: u64, flags: AtomicU32, - mutex: Mutex, + mutex: RawMutex, value: UnsafeCell>, } -impl State { - fn new(value: Option, refcnt: u32) -> AliasableBox { +impl State { + fn new(key: K, value: Option, refcnt: u32, hash: u64) -> AliasableBox { AliasableBox::from_unique(Box::new(Self { + key, + hash, flags: AtomicU32::new(StateFlags::new(refcnt, value.is_some()).0), - mutex: Mutex::new(), + mutex: RawMutex::INIT, value: UnsafeCell::new(value), })) } @@ -101,18 +109,89 @@ impl State { } } -/// A thread-safe hashmap that supports locking entries at the key level. -pub struct LockMap { - map: ShardsMap>>, +// --------------------------------------------------------------------------- +// ShardInner – HashTable storage +// --------------------------------------------------------------------------- + +struct ShardInner { + table: HashTable>>, } -impl Default for LockMap { - fn default() -> Self { - Self::new() +impl ShardInner { + fn with_capacity(capacity: usize) -> Self { + Self { + table: HashTable::with_capacity(capacity), + } + } +} + +// --------------------------------------------------------------------------- +// ShardMap +// --------------------------------------------------------------------------- + +struct ShardMap { + inner: std::sync::Mutex>, +} + +impl ShardMap { + fn with_capacity(capacity: usize) -> Self { + Self { + inner: std::sync::Mutex::new(ShardInner::with_capacity(capacity)), + } + } + + fn len(&self) -> usize { + self.inner.lock().unwrap().table.len() + } + + fn is_empty(&self) -> bool { + self.inner.lock().unwrap().table.is_empty() } } -/// Returns the default number of shards to use for the `LockMap`. +// --------------------------------------------------------------------------- +// LockMap +// --------------------------------------------------------------------------- + +/// A thread-safe hashmap that supports locking entries at the key level. +/// +/// `LockMap` provides a concurrent hashmap with fine-grained per-key locking. +/// Each key can be independently locked, so operations on different keys can +/// proceed in parallel. The map is internally sharded to reduce contention on +/// the map structure itself. +/// +/// # Storage Design +/// +/// The key and its pre-computed hash are stored together in the internal entry +/// state, so each operation hashes the key only once. The full hash is also +/// reused for shard selection and table probing. +/// +/// # Examples +/// +/// ``` +/// use lockmap::LockMap; +/// +/// let map = LockMap::::new(); +/// +/// // Basic operations +/// map.insert("key1".to_string(), 42); +/// assert_eq!(map.get("key1"), Some(42)); +/// +/// // Entry API for exclusive access +/// { +/// let mut entry = map.entry("key2".to_string()); +/// entry.insert(123); +/// } +/// +/// // Remove a value +/// assert_eq!(map.remove("key1"), Some(42)); +/// assert_eq!(map.get("key1"), None); +/// ``` +pub struct LockMap { + shards: Vec>, + hasher: RandomState, +} + fn default_shard_amount() -> usize { static DEFAULT_SHARD_AMOUNT: OnceLock = OnceLock::new(); *DEFAULT_SHARD_AMOUNT.get_or_init(|| { @@ -120,139 +199,152 @@ fn default_shard_amount() -> usize { }) } -/// The main thread-safe map type providing per-key level locking. +impl Default for LockMap { + fn default() -> Self { + Self::new() + } +} + impl LockMap { /// Creates a new `LockMap` with the default number of shards. - /// - /// # Returns - /// - /// A new `LockMap` instance. pub fn new() -> Self { - Self { - map: ShardsMap::with_capacity_and_shard_amount(0, default_shard_amount()), - } + Self::with_capacity_and_shard_amount(0, default_shard_amount()) } - /// Creates a new `LockMap` with the specified initial capacity and the default number of shards. - /// - /// # Arguments - /// - /// * `capacity` - The initial capacity of the hashmap. - /// - /// # Returns - /// - /// A new `LockMap` instance. + /// Creates a new `LockMap` with the specified initial capacity. pub fn with_capacity(capacity: usize) -> Self { - Self { - map: ShardsMap::with_capacity_and_shard_amount(capacity, default_shard_amount()), - } + Self::with_capacity_and_shard_amount(capacity, default_shard_amount()) } /// Creates a new `LockMap` with the specified initial capacity and number of shards. - /// - /// # Arguments - /// - /// * `capacity` - The initial capacity of the hashmap. - /// * `shard_amount` - The number of shards to create. - /// - /// # Returns - /// - /// A new `LockMap` instance. pub fn with_capacity_and_shard_amount(capacity: usize, shard_amount: usize) -> Self { + assert!(shard_amount > 0, "shard_amount must be greater than 0"); + let per_shard_capacity = capacity.div_ceil(shard_amount); Self { - map: ShardsMap::with_capacity_and_shard_amount(capacity, shard_amount), + shards: (0..shard_amount) + .map(|_| ShardMap::with_capacity(per_shard_capacity)) + .collect(), + hasher: RandomState::default(), } } /// Returns the number of elements in the map. pub fn len(&self) -> usize { - self.map.len() + self.shards.iter().map(|s| s.len()).sum() } /// Returns `true` if the map contains no elements. pub fn is_empty(&self) -> bool { - self.map.is_empty() + self.shards.iter().all(|s| s.is_empty()) + } + + // --- shard routing --- + + /// Compute the shard index from the full hash value. + /// + /// Uses the upper 32 bits of the hash for shard selection. The internal + /// `HashTable` uses the lower bits for bucket selection, so using the upper + /// bits avoids correlation between shard assignment and bucket placement. + #[inline(always)] + fn shard_index(&self, hash: u64) -> usize { + ((hash >> 32) as usize) % self.shards.len() } + /// Rehash closure for `HashTable` growth — reuses the stored hash from + /// each `State`. This avoids calling the hasher entirely during rehash. + #[inline(always)] + fn state_hasher() -> impl Fn(&AliasableBox>) -> u64 { + |s| s.hash + } + + // ------------------------------------------------------------------ + // Public API + // ------------------------------------------------------------------ + /// Gets exclusive access to an entry in the map. /// - /// The returned `EntryByVal` provides exclusive access to the key and its associated value - /// until it is dropped. + /// The returned [`Entry`] provides exclusive access to the key and its + /// associated value until it is dropped. /// /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// # use lockmap::LockMap; /// let map = LockMap::::new(); /// { /// let mut entry = map.entry("key".to_string()); /// entry.insert(42); - /// // let _ = map.get("key".to_string()); // DEADLOCK! - /// // map.insert("key".to_string(), 21); // DEADLOCK! - /// // map.remove("key".to_string()); // DEADLOCK! - /// // let mut entry2 = map.entry("key".to_string()); // DEADLOCK! /// } /// ``` - pub fn entry(&self, key: K) -> EntryByVal<'_, K, V> - where - K: Clone, - { - let ptr: *mut State = self.map.update(key.clone(), |s| match s { - Some(state) => { - state.inc_ref(); - let ptr = &**state as *const State as *mut State; - (UpdateAction::Keep, ptr) - } - None => { - let state = State::new(None, 1); - let ptr = &*state as *const State as *mut State; - (UpdateAction::Replace(state), ptr) + pub fn entry(&self, key: K) -> Entry<'_, K, V> { + let hash = self.hasher.hash_one(&key); + let shard = &self.shards[self.shard_index(hash)]; + let ptr: *mut State = { + let mut inner = shard.inner.lock().unwrap(); + match inner + .table + .entry(hash, |s| s.key.borrow() == &key, Self::state_hasher()) + { + TableEntry::Occupied(occupied) => { + let ptr = &**occupied.get() as *const State as *mut State; + unsafe { &*ptr }.inc_ref(); + ptr + } + TableEntry::Vacant(vacant) => { + let state = State::new(key, None, 1, hash); + let ptr = &*state as *const State as *mut State; + vacant.insert(state); + ptr + } } - }); - - self.guard_by_val(ptr, key) + }; + self.guard(ptr) } - /// Gets exclusive access to an entry in the map. - /// - /// The returned `EntryByVal` provides exclusive access to the key and its associated value - /// until it is dropped. + /// Gets exclusive access to an entry by reference. /// /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// # use lockmap::LockMap; /// let map = LockMap::::new(); /// { /// let mut entry = map.entry_by_ref("key"); /// entry.insert(42); - /// // let _ = map.get("key"); // DEADLOCK! - /// // map.insert_by_ref("key", 21); // DEADLOCK! - /// // map.remove("key"); // DEADLOCK! - /// // let mut entry2 = map.entry_by_ref("key"); // DEADLOCK! /// } /// ``` - pub fn entry_by_ref<'a, 'b, Q>(&'a self, key: &'b Q) -> EntryByRef<'a, 'b, K, Q, V> + pub fn entry_by_ref(&self, key: &Q) -> Entry<'_, K, V> where K: Borrow + for<'c> From<&'c Q>, Q: Eq + Hash + ?Sized, { - let ptr: *mut State = self.map.update_by_ref(key, |s| match s { - Some(state) => { - state.inc_ref(); - let ptr = &**state as *const State as *mut State; - (UpdateAction::Keep, ptr) - } - None => { - let state = State::new(None, 1); - let ptr = &*state as *const State as *mut State; - (UpdateAction::Replace(state), ptr) + let hash = self.hasher.hash_one(key); + let shard = &self.shards[self.shard_index(hash)]; + let ptr: *mut State = { + let mut inner = shard.inner.lock().unwrap(); + match inner + .table + .entry(hash, |s| s.key.borrow() == key, Self::state_hasher()) + { + TableEntry::Occupied(occupied) => { + let ptr = &**occupied.get() as *const State as *mut State; + unsafe { &*ptr }.inc_ref(); + ptr + } + TableEntry::Vacant(vacant) => { + let owned_key: K = key.into(); + let state = State::new(owned_key, None, 1, hash); + let ptr = &*state as *const State as *mut State; + vacant.insert(state); + ptr + } } - }); - - self.guard_by_ref(ptr, key) + }; + self.guard(ptr) } /// Gets the value associated with the given key. @@ -260,13 +352,6 @@ impl LockMap { /// If other threads are currently accessing the key, this will wait /// until exclusive access is available before returning. /// - /// # Arguments - /// * `key` - The key to look up - /// - /// # Returns - /// * `Some(V)` if the key exists - /// * `None` if the key doesn't exist - /// /// # Performance Note /// /// When no other thread holds an entry for this key, the `clone()` operation @@ -277,11 +362,12 @@ impl LockMap { /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// use lockmap::LockMap; /// /// let map = LockMap::::new(); - /// map.insert_by_ref("key", 42); + /// map.insert("key".to_string(), 42); /// assert_eq!(map.get("key"), Some(42)); /// assert_eq!(map.get("missing"), None); /// ``` @@ -291,160 +377,151 @@ impl LockMap { V: Clone, Q: Eq + Hash + ?Sized, { - let mut ptr: *mut State = std::ptr::null_mut(); - let value = self.map.simple_update(key, |s| match s { - Some(state) => { + let hash = self.hasher.hash_one(key); + let shard = &self.shards[self.shard_index(hash)]; + let mut ptr: *mut State = std::ptr::null_mut(); + + let value = { + let inner = shard.inner.lock().unwrap(); + let p = inner + .table + .find(hash, |s| s.key.borrow() == key) + .map(|s| &**s as *const State as *mut State) + .unwrap_or(std::ptr::null_mut()); + if !p.is_null() { + let state = unsafe { &*p }; if state.flags().refcnt() == 0 { - // SAFETY: We are inside the map's shard lock, and the reference count is 0, - // meaning no other thread can be holding an `Entry` for this key. - let value = unsafe { state.value_ref() }.clone(); - (SimpleAction::Keep, value) + // SAFETY: refcnt == 0 means no Entry guard exists. + unsafe { state.value_ref() }.clone() } else { state.inc_ref(); - ptr = &**state as *const State as *mut State; - (SimpleAction::Keep, None) + ptr = p; + None } + } else { + None } - None => (SimpleAction::Keep, None), - }); + }; if ptr.is_null() { return value; } - self.guard_by_ref(ptr, key).get().clone() + self.guard(ptr).get().clone() } - /// Sets a value in the map. - /// - /// If other threads are currently accessing the key, this will wait - /// until exclusive access is available before updating. - /// - /// # Arguments - /// * `key` - The key to update - /// * `value` - The value to set + /// Sets a value in the map, returning the previous value if any. /// /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// use lockmap::LockMap; /// /// let map = LockMap::::new(); - /// - /// // Set a value /// assert_eq!(map.insert("key".to_string(), 42), None); - /// - /// // Update existing value /// assert_eq!(map.insert("key".to_string(), 123), Some(42)); /// ``` - pub fn insert(&self, key: K, value: V) -> Option - where - K: Clone, - { - let (ptr, value) = self.map.update(key.clone(), move |s| match s { - Some(state) => { - let flags = state.flags(); - if flags.refcnt() == 0 { - // SAFETY: We are inside the map's shard lock, and the reference count is 0, - // meaning no other thread can be holding an `Entry` for this key. - let value = unsafe { state.value_mut() }.replace(value); - if !flags.has_value() { + pub fn insert(&self, key: K, value: V) -> Option { + let hash = self.hasher.hash_one(&key); + let shard = &self.shards[self.shard_index(hash)]; + let (ptr, old) = { + let mut inner = shard.inner.lock().unwrap(); + match inner + .table + .entry(hash, |s| s.key.borrow() == &key, Self::state_hasher()) + { + TableEntry::Occupied(occupied) => { + let p = &**occupied.get() as *const State as *mut State; + let state = unsafe { &*p }; + let flags = state.flags(); + if flags.refcnt() == 0 { + // SAFETY: refcnt == 0 → exclusive. + let old = unsafe { state.value_mut() }.replace(value); state.set_value_state(true); + (std::ptr::null_mut(), old) + } else { + state.inc_ref(); + (p, Some(value)) } - (UpdateAction::Keep, (std::ptr::null_mut(), value)) - } else { - state.inc_ref(); - let ptr: *mut State = &**state as *const State as *mut State; - (UpdateAction::Keep, (ptr, Some(value))) + } + TableEntry::Vacant(vacant) => { + let state = State::new(key, Some(value), 0, hash); + vacant.insert(state); + (std::ptr::null_mut(), None) } } - None => { - let state = State::new(Some(value), 0); - (UpdateAction::Replace(state), (std::ptr::null_mut(), None)) - } - }); + }; if ptr.is_null() { - return value; + return old; } - self.guard_by_val(ptr, key).swap(value) + self.guard(ptr).swap(old) } - /// Sets a value in the map. - /// - /// If other threads are currently accessing the key, this will wait - /// until exclusive access is available before updating. - /// - /// # Arguments - /// * `key` - The key to update - /// * `value` - The value to set + /// Sets a value in the map by reference key. /// /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// use lockmap::LockMap; /// /// let map = LockMap::::new(); - /// - /// // Set a value /// map.insert_by_ref("key", 42); - /// - /// // Update existing value - /// map.insert_by_ref("key", 123); + /// assert_eq!(map.get("key"), Some(42)); /// ``` pub fn insert_by_ref(&self, key: &Q, value: V) -> Option where K: Borrow + for<'c> From<&'c Q>, Q: Eq + Hash + ?Sized, { - let (ptr, value) = self.map.update_by_ref(key, move |s| match s { - Some(state) => { - let flags = state.flags(); - if flags.refcnt() == 0 { - // SAFETY: We are inside the map's shard lock, and the reference count is 0, - // meaning no other thread can be holding an `Entry` for this key. - let value = unsafe { state.value_mut() }.replace(value); - if !flags.has_value() { + let hash = self.hasher.hash_one(key); + let shard = &self.shards[self.shard_index(hash)]; + let (ptr, old) = { + let mut inner = shard.inner.lock().unwrap(); + match inner + .table + .entry(hash, |s| s.key.borrow() == key, Self::state_hasher()) + { + TableEntry::Occupied(occupied) => { + let p = &**occupied.get() as *const State as *mut State; + let state = unsafe { &*p }; + let flags = state.flags(); + if flags.refcnt() == 0 { + let old = unsafe { state.value_mut() }.replace(value); state.set_value_state(true); + (std::ptr::null_mut(), old) + } else { + state.inc_ref(); + (p, Some(value)) } - (UpdateAction::Keep, (std::ptr::null_mut(), value)) - } else { - state.inc_ref(); - let ptr: *mut State = &**state as *const State as *mut State; - (UpdateAction::Keep, (ptr, Some(value))) + } + TableEntry::Vacant(vacant) => { + let owned_key: K = key.into(); + let state = State::new(owned_key, Some(value), 0, hash); + vacant.insert(state); + (std::ptr::null_mut(), None) } } - None => { - let state = State::new(Some(value), 0); - (UpdateAction::Replace(state), (std::ptr::null_mut(), None)) - } - }); + }; if ptr.is_null() { - return value; + return old; } - self.guard_by_ref(ptr, key).swap(value) + self.guard(ptr).swap(old) } /// Checks if the map contains a key. /// - /// If other threads are currently accessing the key, this will wait - /// until exclusive access is available before checking. - /// - /// # Arguments - /// * `key` - The key to check - /// - /// # Returns - /// * `true` if the key exists - /// * `false` if the key doesn't exist - /// /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// use lockmap::LockMap; /// @@ -458,49 +535,49 @@ impl LockMap { K: Borrow, Q: Eq + Hash + ?Sized, { - let mut ptr: *mut State = std::ptr::null_mut(); - let value = self.map.simple_update(key, |s| match s { - Some(state) => { + let hash = self.hasher.hash_one(key); + let shard = &self.shards[self.shard_index(hash)]; + let mut ptr: *mut State = std::ptr::null_mut(); + + let found = { + let inner = shard.inner.lock().unwrap(); + let p = inner + .table + .find(hash, |s| s.key.borrow() == key) + .map(|s| &**s as *const State as *mut State) + .unwrap_or(std::ptr::null_mut()); + if !p.is_null() { + let state = unsafe { &*p }; if state.flags().refcnt() == 0 { - // SAFETY: We are inside the map's shard lock, and the reference count is 0, - // meaning no other thread can be holding an `Entry` for this key. - (SimpleAction::Keep, unsafe { state.value_ref() }.is_some()) + unsafe { state.value_ref() }.is_some() } else { state.inc_ref(); - ptr = &**state as *const State as *mut State; - (SimpleAction::Keep, false) + ptr = p; + false } + } else { + false } - None => (SimpleAction::Keep, false), - }); + }; if ptr.is_null() { - return value; + return found; } - self.guard_by_ref(ptr, key).get().is_some() + self.guard(ptr).get().is_some() } /// Removes a key from the map. /// - /// If other threads are currently accessing the key, this will wait - /// until exclusive access is available before removing. - /// - /// # Arguments - /// * `key` - The key to remove - /// - /// # Returns - /// * `Some(V)` if the key exists - /// * `None` if the key doesn't exist - /// /// **Locking behaviour:** Deadlock if called when holding the same entry. /// /// # Examples + /// /// ``` /// use lockmap::LockMap; /// /// let map = LockMap::::new(); - /// map.insert_by_ref("key", 42); + /// map.insert("key".to_string(), 42); /// assert_eq!(map.remove("key"), Some(42)); /// assert_eq!(map.get("key"), None); /// ``` @@ -509,52 +586,36 @@ impl LockMap { K: Borrow, Q: Eq + Hash + ?Sized, { - let mut ptr: *mut State = std::ptr::null_mut(); - let value = self.map.simple_update(key, |s| match s { - Some(state) => { - if state.flags().refcnt() == 0 { - // SAFETY: We are inside the map's shard lock, and the reference count is 0, - // meaning no other thread can be holding an `Entry` for this key. - let value = unsafe { state.value_mut() }.take(); - (SimpleAction::Remove, value) - } else { + let hash = self.hasher.hash_one(key); + let shard = &self.shards[self.shard_index(hash)]; + + let ptr = { + let mut inner = shard.inner.lock().unwrap(); + match inner.table.find_entry(hash, |s| s.key.borrow() == key) { + Ok(occupied) => { + let p = &**occupied.get() as *const State as *mut State; + let state = unsafe { &*p }; + if state.flags().refcnt() == 0 { + // SAFETY: refcnt == 0 → exclusive. + let value = unsafe { state.value_mut() }.take(); + let _ = occupied.remove(); + return value; + } state.inc_ref(); - ptr = &**state as *const State as *mut State; - (SimpleAction::Keep, None) + p } + Err(_) => return None, } - None => (SimpleAction::Keep, None), - }); - - if ptr.is_null() { - return value; - } + }; - self.guard_by_ref(ptr, key).remove() + self.guard(ptr).remove() } /// Acquires exclusive locks for a batch of keys in a deadlock-safe manner. /// - /// This function is designed to prevent deadlocks that can occur when multiple threads - /// try to acquire locks on the same set of keys in different orders. It achieves this - /// by taking a `BTreeSet` of keys, which ensures the keys are processed and locked + /// Takes a `BTreeSet` of keys, ensuring they are processed and locked /// in a consistent, sorted order across all threads. /// - /// The function iterates through the sorted keys, acquiring an exclusive lock for each - /// key and its associated value. The returned `Vec` contains RAII guards, which - /// automatically release the locks when they are dropped. - /// - /// # Arguments - /// - /// * `keys` - The `BTreeSet` of keys to lock. The use of `BTreeSet` is crucial as it - /// enforces a global, canonical locking order, thus preventing deadlocks. - /// - /// # Returns - /// - /// A `Vec>` containing the RAII guards for each locked key. - /// - /// **Locking behaviour:** Deadlock if called when holding the same entry. - /// /// # Examples /// /// ```rust @@ -566,22 +627,17 @@ impl LockMap { /// map.insert(2, 200); /// map.insert(3, 300); /// - /// // Create a set of keys to lock. Note that the order in the set doesn't matter - /// // since BTreeSet will sort them. /// let mut keys = BTreeSet::new(); /// keys.insert(3); /// keys.insert(1); /// keys.insert(2); /// - /// // Acquire locks for all keys in a deadlock-safe manner /// let mut locked_entries = map.batch_lock::>(keys); /// - /// // The locks are held as long as `locked_entries` is in scope /// locked_entries.get_mut(&1).and_then(|entry| entry.insert(101)); /// locked_entries.get_mut(&2).and_then(|entry| entry.insert(201)); /// locked_entries.get_mut(&3).and_then(|entry| entry.insert(301)); /// - /// // When `locked_entries` is dropped, all locks are released /// drop(locked_entries); /// /// assert_eq!(map.get(&1), Some(101)); @@ -591,68 +647,18 @@ impl LockMap { pub fn batch_lock<'a, M>(&'a self, keys: BTreeSet) -> M where K: Clone, - M: FromIterator<(K, EntryByVal<'a, K, V>)>, + M: FromIterator<(K, Entry<'a, K, V>)>, { keys.into_iter() .map(|key| (key.clone(), self.entry(key))) .collect() } - /// Attempts to remove an entry from the map if it's no longer needed. - /// - /// An entry is considered no longer needed if its reference count is 0 - /// and it contains no value. - fn try_remove_entry(&self, key: &Q) - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - self.map.simple_update(key, |value| match value { - Some(state) => { - // SAFETY: We are inside the map's shard lock. If the reference count is 0 here, - // then no `Entry` is currently held for this key, and no other thread - // can increment the reference count without first acquiring this same shard lock. - // Therefore, if the stored value is also `None`, it is safe to remove - // the entry from the map. - if state.flags().pending_cleanup() { - (SimpleAction::Remove, ()) - } else { - (SimpleAction::Keep, ()) - } - } - // The key might have been removed by another thread (e.g., via `remove`) - // between the reference count decrement and this call. - None => (SimpleAction::Keep, ()), - }); - } - - fn guard_by_val(&self, ptr: *mut State, key: K) -> EntryByVal<'_, K, V> { - // SAFETY: The pointer `ptr` is valid because it was just retrieved from the map - // and its reference count was incremented, ensuring it won't be dropped. - // The `AliasableBox` in the map ensures the `State` remains at a stable - // memory location. - unsafe { (*ptr).mutex.lock() }; - EntryByVal { - map: self, - key, - state: ptr, - } - } - - fn guard_by_ref<'a, 'b, Q>( - &'a self, - ptr: *mut State, - key: &'b Q, - ) -> EntryByRef<'a, 'b, K, Q, V> - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - // SAFETY: Same as `guard_by_val`. + fn guard(&self, ptr: *mut State) -> Entry<'_, K, V> { + // SAFETY: ptr is valid (ref-counted) and stable (AliasableBox). unsafe { (*ptr).mutex.lock() }; - EntryByRef { + Entry { map: self, - key, state: ptr, } } @@ -664,280 +670,144 @@ impl std::fmt::Debug for LockMap { } } -/// An RAII guard providing exclusive access to a key-value pair in the `LockMap`. +// --------------------------------------------------------------------------- +// Entry +// --------------------------------------------------------------------------- + +/// An RAII guard providing exclusive access to a key-value pair in the [`LockMap`]. /// -/// When dropped, this type automatically unlocks the entry allowing other threads to access it. +/// The key is obtained directly from the internal `State`, so there is no need +/// for separate by-value / by-reference entry types. /// -/// # Type Parameters -/// * `'a` - Lifetime of the `LockMap` -/// * `K` - Key type that must implement `Eq + Hash` -/// * `V` - Value type +/// When dropped, this type automatically unlocks the entry and may trigger +/// cleanup of empty entries. /// /// # Examples +/// /// ``` /// use lockmap::LockMap; /// /// let map = LockMap::new(); /// { -/// // Get exclusive access to an entry /// let mut entry = map.entry("key"); -/// -/// // Modify the value /// entry.insert(42); -/// -/// // EntryByVal is automatically unlocked when dropped /// } /// ``` -pub struct EntryByVal<'a, K: Eq + Hash, V> { +pub struct Entry<'a, K: Eq + Hash, V> { map: &'a LockMap, - key: K, - state: *mut State, + state: *mut State, } -impl EntryByVal<'_, K, V> { +// SAFETY: The guard holds a per-key mutex lock and a valid, ref-counted pointer. +unsafe impl Send for Entry<'_, K, V> {} +unsafe impl Sync for Entry<'_, K, V> {} + +impl Entry<'_, K, V> { /// Returns a reference to the entry's key. - /// - /// # Returns - /// - /// A reference to the key associated with this entry. pub fn key(&self) -> &K { - &self.key + // SAFETY: The state pointer is valid for the lifetime of this guard + // and the key is immutable. + unsafe { &(*self.state).key } } /// Returns a reference to the entry's value. - /// - /// # Returns - /// - /// A reference to `Some(V)` if the entry has a value, or `None` if the entry is vacant. pub fn get(&self) -> &Option { - // SAFETY: The entry holds the lock on the `State`, so it is safe to access the value. unsafe { (*self.state).value_ref() } } /// Returns a mutable reference to the entry's value. - /// - /// # Returns - /// - /// A mutable reference to `Some(V)` if the entry has a value, or `None` if the entry is vacant. pub fn get_mut(&mut self) -> &mut Option { - // SAFETY: The entry holds the lock on the `State`, so it is safe to access the value. unsafe { (*self.state).value_mut() } } - /// Sets the value of the entry, returning the old value if it existed. - /// - /// # Arguments - /// - /// * `value` - The new value to insert - /// - /// # Returns - /// - /// The previous value if the entry was occupied, or `None` if it was vacant. + /// Sets the value, returning the old value if any. pub fn insert(&mut self, value: V) -> Option { self.get_mut().replace(value) } - /// Swaps the value of the entry with the provided value. - /// - /// # Arguments - /// - /// * `value` - The new value to swap in (wrapped in `Option`) - /// - /// # Returns - /// - /// The previous value of the entry. + /// Swaps the value with the provided one, returning the old value. pub fn swap(&mut self, mut value: Option) -> Option { std::mem::swap(self.get_mut(), &mut value); value } - /// Removes the value from the entry, returning it if it existed. - /// - /// # Returns - /// - /// The value that was stored in the entry, or `None` if the entry was vacant. + /// Removes the value, returning it if it existed. pub fn remove(&mut self) -> Option { - // SAFETY: The entry holds the lock on the `State`, so it is safe to access the value. self.get_mut().take() } } -impl std::fmt::Debug for EntryByVal<'_, K, V> { +impl std::fmt::Debug for Entry<'_, K, V> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("EntryByVal") - .field("key", &self.key) + f.debug_struct("Entry") + .field("key", self.key()) .field("value", self.get()) .finish() } } -impl Drop for EntryByVal<'_, K, V> { +impl Drop for Entry<'_, K, V> { fn drop(&mut self) { - // Update flags based on current value state - unsafe { &*self.state }.set_value_state(self.get().is_some()); - - // SAFETY: The entry holds the lock on the `State`, so it is safe to unlock it. - unsafe { &*self.state }.mutex.unlock(); - - // SAFETY: The pointer `self.state` remains valid here because the `EntryByVal` - // incremented the `State`'s reference count when it was created. While `self` is - // alive in this `drop` call, the reference count is therefore at least 1, and this - // `fetch_sub(1, ...)` is decrementing that last reference held by the entry. The - // `State` is only deallocated once its reference count reaches zero, which can only - // occur after this `fetch_sub` completes. Thus, dereferencing `self.state` to access - // the reference count is safe at this point. - let flags = unsafe { (*self.state).dec_ref() }; - if flags.pending_cleanup() { - self.map.try_remove_entry(&self.key); - } - } -} - -/// An RAII guard providing exclusive access to a key-value pair in the `LockMap`. -/// -/// When dropped, this type automatically unlocks the entry allowing other threads to access it. -/// -/// # Type Parameters -/// * `'a` - Lifetime of the `LockMap` -/// * `'b` - Lifetime of the key reference -/// * `K` - Key type that must implement `Eq + Hash` -/// * `Q` - Query type that can be borrowed from `K` -/// * `V` - Value type -/// -/// # Examples -/// ``` -/// use lockmap::LockMap; -/// -/// let map = LockMap::::new(); -/// { -/// // Get exclusive access to an entry -/// let mut entry = map.entry_by_ref("key"); -/// -/// // Modify the value -/// entry.insert(42); -/// -/// // EntryByRef is automatically unlocked when dropped -/// } -/// ``` -pub struct EntryByRef<'a, 'b, K: Eq + Hash + Borrow, Q: Eq + Hash + ?Sized, V> { - map: &'a LockMap, - key: &'b Q, - state: *mut State, -} - -impl, Q: Eq + Hash + ?Sized, V> EntryByRef<'_, '_, K, Q, V> { - /// Returns a reference to the entry's key. - /// - /// # Returns - /// - /// A reference to the key associated with this entry. - pub fn key(&self) -> &Q { - self.key - } - - /// Returns a reference to the entry's value. - /// - /// # Returns - /// - /// A reference to `Some(V)` if the entry has a value, or `None` if the entry is vacant. - pub fn get(&self) -> &Option { - // SAFETY: The entry holds the lock on the `State`, so it is safe to access the value. - unsafe { (*self.state).value_ref() } - } - - /// Returns a mutable reference to the entry's value. - /// - /// # Returns - /// - /// A mutable reference to `Some(V)` if the entry has a value, or `None` if the entry is vacant. - pub fn get_mut(&mut self) -> &mut Option { - // SAFETY: The entry holds the lock on the `State`, so it is safe to access the value. - unsafe { (*self.state).value_mut() } - } - - /// Sets the value of the entry, returning the old value if it existed. - /// - /// # Arguments - /// - /// * `value` - The new value to insert - /// - /// # Returns - /// - /// The previous value if the entry was occupied, or `None` if it was vacant. - pub fn insert(&mut self, value: V) -> Option { - self.get_mut().replace(value) - } - - /// Swaps the value of the entry with the provided value. - /// - /// # Arguments - /// - /// * `value` - The new value to swap in (wrapped in `Option`) - /// - /// # Returns - /// - /// The previous value of the entry. - pub fn swap(&mut self, mut value: Option) -> Option { - std::mem::swap(self.get_mut(), &mut value); - value - } - - /// Removes the value from the entry, returning it if it existed. - /// - /// # Returns - /// - /// The value that was stored in the entry, or `None` if the entry was vacant. - pub fn remove(&mut self) -> Option { - // SAFETY: The entry holds the lock on the `State`, so it is safe to access the value. - self.get_mut().take() - } -} + // 1. Update value state flag + let has_value = self.get().is_some(); + let state_ref = unsafe { &*self.state }; + state_ref.set_value_state(has_value); + + // 2. Unlock the entry's mutex + // SAFETY: We hold the lock (acquired in guard()). + unsafe { state_ref.mutex.unlock() }; + + // 3. CAS loop to decrement reference count + let mut current = state_ref.flags.load(Ordering::Acquire); + loop { + let flags = StateFlags(current); + + // If this is the last guard and no value, proceed to cleanup + if flags.refcnt() == 1 && !flags.has_value() { + break; + } -impl std::fmt::Debug for EntryByRef<'_, '_, K, Q, V> -where - K: Eq + Hash + Borrow + std::fmt::Debug, - Q: Eq + Hash + ?Sized + std::fmt::Debug, - V: std::fmt::Debug, -{ - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("EntryByRef") - .field("key", &self.key) - .field("value", self.get()) - .finish() - } -} + let new_flags = StateFlags::new(flags.refcnt() - 1, flags.has_value()); + match state_ref.flags.compare_exchange_weak( + current, + new_flags.0, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return, // Not last guard or still has value, return early + Err(actual) => current = actual, + } + } -impl, Q: Eq + Hash + ?Sized, V> Drop for EntryByRef<'_, '_, K, Q, V> { - fn drop(&mut self) { - // Update flags based on current value state - unsafe { &*self.state }.set_value_state(self.get().is_some()); - - // SAFETY: The entry holds the lock on the `State`, so it is safe to unlock it. - unsafe { &*self.state }.mutex.unlock(); - - // SAFETY: The pointer `self.state` remains valid here because the `EntryByRef` - // incremented the `State`'s reference count when it was created. While `self` is - // alive in this `drop` call, the reference count is therefore at least 1, and this - // `fetch_sub(1, ...)` is decrementing that last reference held by the entry. The - // `State` is only deallocated once its reference count reaches zero, which can only - // occur after this `fetch_sub` completes. Thus, dereferencing `self.state` to access - // the reference count is safe at this point. - let flags = unsafe { (*self.state).dec_ref() }; - if flags.pending_cleanup() { - self.map.try_remove_entry(self.key); + // 4. Acquire shard lock using the stored hash (no re-hashing). + let shard_idx = self.map.shard_index(state_ref.hash); + let shard = &self.map.shards[shard_idx]; + let mut inner = shard.inner.lock().unwrap(); + + // 5. Decrement reference count again; cleanup if needed + let final_flags = state_ref.dec_ref(); + if final_flags.pending_cleanup() { + let state_ptr = self.state as *const State; + if let Ok(entry) = inner + .table + .find_entry(state_ref.hash, |s| std::ptr::eq(&**s, state_ptr)) + { + let _ = entry.remove(); + } } } } +// =========================================================================== +// Tests +// =========================================================================== + #[cfg(test)] mod tests { use super::*; - use std::{ - collections::HashMap, - sync::{ - atomic::{AtomicU32, Ordering}, - Arc, - }, + use std::sync::{ + atomic::{AtomicU32, Ordering}, + Arc, }; #[test] @@ -1126,7 +996,8 @@ mod tests { std::thread::spawn(move || { for _ in 0..N { let keys = (0..3).map(|_| rand::random::() % 32).collect(); - let mut entries: HashMap<_, _> = lock_map.batch_lock(keys); + let mut entries: std::collections::HashMap<_, _> = + lock_map.batch_lock(keys); for entry in entries.values_mut() { assert!(entry.get().is_none()); entry.insert(1); @@ -1203,9 +1074,7 @@ mod tests { #[test] fn test_lockmap_get_set_by_ref() { - let lock_map = Arc::new(LockMap::::with_capacity_and_shard_amount( - 256, 16, - )); + let lock_map = Arc::new(LockMap::::with_capacity_and_shard_amount(256, 16)); #[cfg(not(miri))] const N: usize = 1 << 18; #[cfg(miri)] @@ -1313,4 +1182,21 @@ mod tests { THREADS as u32 * OPS_PER_THREAD as u32 ); } + + #[test] + fn test_lockmap_grow() { + let lock_map = Arc::new(LockMap::::with_capacity(4)); + #[cfg(not(miri))] + const N: usize = 1 << 12; + #[cfg(miri)] + const N: usize = 1 << 6; + + for i in 0..N { + lock_map.insert(i as u32, i as u32); + } + + for i in 0..N { + assert_eq!(lock_map.get(&(i as u32)), Some(i as u32)); + } + } } diff --git a/lockmap-lru/src/lru_lockmap.rs b/src/lru_lockmap.rs similarity index 90% rename from lockmap-lru/src/lru_lockmap.rs rename to src/lru_lockmap.rs index 06d20ad..a95a7c4 100644 --- a/lockmap-lru/src/lru_lockmap.rs +++ b/src/lru_lockmap.rs @@ -2,7 +2,8 @@ use aliasable::boxed::AliasableBox; use foldhash::fast::RandomState; use hashbrown::hash_table::Entry; use hashbrown::HashTable; -use lockmap_core::Mutex; +use parking_lot::lock_api::RawMutex as _; +use parking_lot::RawMutex; use std::borrow::Borrow; use std::cell::UnsafeCell; use std::hash::{BuildHasher, Hash}; @@ -48,7 +49,7 @@ struct State { key: K, hash: u64, flags: AtomicU32, - mutex: Mutex, + mutex: RawMutex, value: UnsafeCell>, // Intrusive doubly-linked list pointers for LRU ordering. // These are only accessed while the shard's `std::sync::Mutex` is held. @@ -62,7 +63,7 @@ impl State { key, hash, flags: AtomicU32::new(StateFlags::new(refcnt, value.is_some()).0), - mutex: Mutex::new(), + mutex: RawMutex::INIT, value: UnsafeCell::new(value), prev: UnsafeCell::new(std::ptr::null_mut()), next: UnsafeCell::new(std::ptr::null_mut()), @@ -201,33 +202,23 @@ impl LruShardInner { /// /// Walks from the tail (LRU end) towards the head (MRU end). Entries that /// are currently in use (refcnt > 0) are **skipped** and traversal continues - /// to the next candidate. This ensures eviction still makes progress even - /// when the tail entry is held by another thread. + /// to the next candidate. /// /// `current` is the entry that was just accessed or inserted — it must not /// be evicted even though it is at the head of the list. fn try_evict(&mut self, current: *mut State) { let mut cursor = self.tail; while self.table.len() > self.max_size && !cursor.is_null() && cursor != current { - // SAFETY: `cursor` is guaranteed non-null by the while condition. - // We read `prev` before a potential detach so we can advance the - // cursor even after `cursor` is detached and freed. `prev` may be - // null (when cursor is the head), which is fine because the while - // condition will catch it on the next iteration. let prev = unsafe { *(*cursor).prev.get() }; let state = unsafe { &*cursor }; if state.flags().refcnt() > 0 { - // This entry is currently in use — skip it and try the next one. cursor = prev; continue; } - // SAFETY: `cursor` is valid and in the list. unsafe { self.detach(cursor) }; - // Remove from the HashTable using the stored hash and pointer - // equality. This avoids re-hashing the key. let hash = state.hash; if let Ok(entry) = self.table.find_entry(hash, |s| std::ptr::eq(&**s, cursor)) { let _ = entry.remove(); @@ -276,9 +267,9 @@ impl LruShardMap { /// A thread-safe LRU cache that supports locking entries at the key level. /// -/// `LruLockMap` extends the per-key locking design of `LockMap` with LRU -/// (Least Recently Used) eviction. The total capacity is divided evenly among -/// internal shards, and each shard independently evicts its least-recently-used +/// `LruLockMap` extends the per-key locking design of [`LockMap`](crate::LockMap) +/// with LRU (Least Recently Used) eviction. The total capacity is divided evenly +/// among internal shards, and each shard independently evicts its least-recently-used /// entries when it exceeds its share of the capacity. /// /// # Eviction Policy @@ -294,7 +285,7 @@ impl LruShardMap { /// # Examples /// /// ``` -/// use lockmap_lru::LruLockMap; +/// use lockmap::LruLockMap; /// /// let cache = LruLockMap::::new(1024); /// @@ -362,16 +353,10 @@ impl LruLockMap { /// `max_size.div_ceil(shard_amount)`, so the effective total capacity may be /// rounded up compared to the value originally passed to [`with_options`]. /// - /// Note that this is not a strict upper bound on [`len`](Self::len). In - /// particular, eviction may skip entries that are currently in use (e.g. - /// with a positive reference count), and implementation details such as a - /// minimum of one entry per shard can also cause the actual number of - /// stored entries to temporarily exceed this configured target. - /// /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::with_options(100, 100, 10); /// assert_eq!(cache.max_size(), 100); /// ``` @@ -390,15 +375,11 @@ impl LruLockMap { /// of excess entries. Eviction will occur lazily on subsequent insertions /// and other operations that may increase the cache size. /// - /// # Arguments - /// - /// * `max_size` - The new maximum number of entries across all shards - /// /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; - /// let mut cache = LruLockMap::::with_options(100, 100, 10); + /// # use lockmap::LruLockMap; + /// let cache = LruLockMap::::with_options(100, 100, 10); /// cache.set_max_size(200); /// assert_eq!(cache.max_size(), 200); /// ``` @@ -411,25 +392,18 @@ impl LruLockMap { // --- shard routing --- - /// Compute the shard index from the full hash value. - /// - /// Uses the upper 32 bits of the hash for shard selection. The internal - /// `HashTable` uses the lower bits for bucket selection, so using the upper - /// bits avoids correlation between shard assignment and bucket placement. #[inline(always)] fn shard_index(&self, hash: u64) -> usize { ((hash >> 32) as usize) % self.shards.len() } - /// Rehash closure for `HashTable` growth — reuses the stored hash from - /// each `State`. This avoids calling the hasher entirely during rehash. #[inline(always)] fn state_hasher() -> impl Fn(&AliasableBox>) -> u64 { |s| s.hash } // ------------------------------------------------------------------ - // Public API – mirrors lockmap::LockMap + // Public API // ------------------------------------------------------------------ /// Gets exclusive access to an entry in the cache. @@ -443,7 +417,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// { /// let mut entry = cache.entry("key".to_string()); @@ -455,7 +429,6 @@ impl LruLockMap { let shard = &self.shards[self.shard_index(hash)]; let ptr: *mut State = { let mut inner = shard.inner.lock().unwrap(); - // Single lookup: find existing entry or prepare insert slot. let ptr = match inner .table @@ -488,7 +461,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// { /// let mut entry = cache.entry_by_ref("key"); @@ -538,7 +511,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// cache.insert("key".to_string(), 42); /// assert_eq!(cache.get("key"), Some(42)); @@ -592,7 +565,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// assert_eq!(cache.insert("key".to_string(), 42), None); /// assert_eq!(cache.insert("key".to_string(), 123), Some(42)); @@ -646,7 +619,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// cache.insert_by_ref("key", 42); /// assert_eq!(cache.get("key"), Some(42)); @@ -706,7 +679,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// cache.insert("key".to_string(), 42); /// assert!(cache.contains_key("key")); @@ -758,7 +731,7 @@ impl LruLockMap { /// # Examples /// /// ``` - /// # use lockmap_lru::LruLockMap; + /// # use lockmap::LruLockMap; /// let cache = LruLockMap::::new(100); /// cache.insert("key".to_string(), 42); /// assert_eq!(cache.remove("key"), Some(42)); @@ -774,7 +747,6 @@ impl LruLockMap { let ptr = { let mut inner = shard.inner.lock().unwrap(); - // Single lookup via find_entry. let p = match inner.table.find_entry(hash, |s| s.key.borrow() == key) { Ok(occupied) => { let p = &**occupied.get() as *const State as *mut State; @@ -782,9 +754,6 @@ impl LruLockMap { if state.flags().refcnt() == 0 { // SAFETY: refcnt == 0 → exclusive. let value = unsafe { state.value_mut() }.take(); - // Remove from table first (consuming the OccupiedEntry - // to release the borrow), then detach from linked list - // while the returned box keeps the State memory alive. let (state_box, _) = occupied.remove(); unsafe { inner.detach(p) }; drop(state_box); @@ -846,8 +815,6 @@ unsafe impl Sync for LruEntry<'_, K, impl LruEntry<'_, K, V> { /// Returns a reference to the entry's key. pub fn key(&self) -> &K { - // SAFETY: The state pointer is valid for the lifetime of this guard - // and the key is immutable. unsafe { &(*self.state).key } } @@ -888,21 +855,6 @@ impl std::fmt::Debug for Lru } impl Drop for LruEntry<'_, K, V> { - /// Drop implementation for `LruEntry`. - /// - /// 1. Update the value state flag to ensure `flags` reflects whether the value exists. - /// 2. Unlock the entry's mutex, allowing other threads to access this entry. - /// 3. Atomically decrement the reference count (`refcnt`). - /// - If this is the last guard and the entry has no value, break and proceed to cleanup. - /// - If not the last guard or the entry still has a value, return early. - /// - Uses a CAS loop to ensure thread-safe decrement of the reference count. - /// 4. Acquire the corresponding shard lock to safely access and modify the LRU list and table. - /// 5. Decrement the reference count again; if `refcnt == 0` and no value, cleanup: - /// - Remove the entry from the LRU list. - /// - Remove the entry from the table, freeing memory. - /// - /// # Safety - /// - Mutex unlocking and linked list operations are performed while holding the shard lock, ensuring thread safety. fn drop(&mut self) { // 1. Update value state flag let has_value = self.get().is_some(); @@ -910,7 +862,8 @@ impl Drop for LruEntry<'_, K, V> { state_ref.set_value_state(has_value); // 2. Unlock the entry's mutex - state_ref.mutex.unlock(); + // SAFETY: We hold the lock (acquired in guard()). + unsafe { state_ref.mutex.unlock() }; // 3. CAS loop to decrement reference count let mut current = state_ref.flags.load(Ordering::Acquire); @@ -942,8 +895,6 @@ impl Drop for LruEntry<'_, K, V> { // 5. Decrement reference count again; cleanup if needed let final_flags = state_ref.dec_ref(); if final_flags.pending_cleanup() { - // SAFETY: The entry is in the linked list and the shard lock is held. - // Detach from linked list first, then remove from table. unsafe { inner.detach(self.state) }; let state_ptr = self.state as *const State; if let Ok(entry) = inner @@ -1065,7 +1016,7 @@ mod tests { #[test] fn test_set_max_size() { - let mut cache = LruLockMap::::with_options(3, 3, 4); + let cache = LruLockMap::::with_options(3, 3, 4); assert_eq!(cache.max_size(), 4); cache.set_max_size(6); assert_eq!(cache.max_size(), 8); @@ -1143,7 +1094,6 @@ mod tests { let _entry = cache.entry(1); // Insert key=4 — should try to evict key=1 but it's in use. - // With the improved eviction, it skips key=1 and evicts key=2 instead. let cache2 = cache.clone(); let t = std::thread::spawn(move || { cache2.insert(4, 40); @@ -1167,29 +1117,20 @@ mod tests { #[test] fn test_lru_evict_skips_multiple_in_use() { - // Verify that eviction walks past multiple in-use entries and still - // evicts other eligible entries. let cache = LruLockMap::::with_options(3, 3, 1); cache.insert(1, 10); cache.insert(2, 20); cache.insert(3, 30); - // Hold the two LRU-most entries (key=1 is tail, key=2 is next) let _entry1 = cache.entry(1); let _entry2 = cache.entry(2); - // Insert key=4 from a separate thread (since entry() would deadlock - // on same thread for key=1 or key=2, but key=4 is new). - // The eviction should skip key=1 and key=2 (both in use), evict key=3. cache.insert(4, 40); - // key=1, key=2 still present (in use) assert_eq!(*_entry1.get(), Some(10)); assert_eq!(*_entry2.get(), Some(20)); - // key=3 evicted, key=4 present - // Note: we use contains_key via a thread since entry() would interact with LRU assert_eq!(cache.get(&3), None); assert_eq!(cache.get(&4), Some(40)); @@ -1205,7 +1146,6 @@ mod tests { cache.insert(2, 20); cache.insert(3, 30); - // Overwriting an existing key should NOT increase count cache.insert(2, 200); assert_eq!(cache.len(), 3); assert_eq!(cache.get(&1), Some(10)); @@ -1224,7 +1164,6 @@ mod tests { cache.remove(&2); assert_eq!(cache.len(), 2); - // Should be able to add more without evicting cache.insert(4, 40); assert_eq!(cache.len(), 3); assert_eq!(cache.get(&1), Some(10)); @@ -1451,7 +1390,6 @@ mod tests { #[test] fn test_concurrent_with_eviction() { - // Small capacity to force frequent evictions under contention let cache = Arc::new(LruLockMap::::with_options(32, 4, 1)); #[cfg(not(miri))] const N: usize = 1 << 14; @@ -1491,7 +1429,6 @@ mod tests { t.join().unwrap(); } - // The cache should not have grown unbounded assert!(cache.len() <= 64); }