From d4d1ab5cc32ab23bfedc256e4772164f99f4b028 Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Thu, 27 Aug 2026 01:04:40 +0800 Subject: [PATCH 1/2] perf: sharded `OnceTable` --- asyncband/src/internal/once_table.rs | 103 ++++++++++++++++--------- asyncband/src/once/once_map/mod.rs | 62 +++++---------- asyncband/src/once/once_map/tests.rs | 10 +-- asyncband/src/singleflight/mod.rs | 29 ++----- asyncband/src/singleflight/tests.rs | 12 +-- tests-integration/tests/traits_test.rs | 22 ++++++ 6 files changed, 127 insertions(+), 111 deletions(-) diff --git a/asyncband/src/internal/once_table.rs b/asyncband/src/internal/once_table.rs index 27865ed3..bb01e409 100644 --- a/asyncband/src/internal/once_table.rs +++ b/asyncband/src/internal/once_table.rs @@ -20,11 +20,17 @@ use std::fmt; use std::hash::BuildHasher; use std::hash::Hash; use std::sync::Arc; +use std::sync::MutexGuard; use hashbrown::HashTable; +use crate::internal::mutex::Mutex; use crate::once::OnceCell; +const SHARD_COUNT: usize = 64; + +type Entries = HashTable>>; + pub struct OnceTableEntry { hash: u64, key: K, @@ -57,7 +63,7 @@ impl OnceTableEntry { /// Shared keyed storage that lets once primitives clean up an exact entry without cloning its key. pub struct OnceTable { - entries: HashTable>>, + shards: Box<[Mutex>]>, hasher: S, } @@ -67,35 +73,40 @@ where V: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_map() - .entries(self.entries.iter().map(|entry| (&entry.key, &entry.cell))) - .finish() + let mut debug_map = f.debug_map(); + for shard in &self.shards { + let entries = shard.lock(); + debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); + } + debug_map.finish() } } impl OnceTable { pub fn with_hasher(hasher: S) -> Self { - Self { - entries: HashTable::new(), - hasher, - } + Self::with_capacity_and_hasher(0, hasher) } pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { - Self { - entries: HashTable::with_capacity(capacity), - hasher, - } + let shard_capacity = capacity.div_ceil(SHARD_COUNT); + let shards = (0..SHARD_COUNT) + .map(|_| Mutex::new(HashTable::with_capacity(shard_capacity))) + .collect(); + Self { shards, hasher } + } + + fn shard(&self, hash: u64) -> MutexGuard<'_, Entries> { + self.shards[hash as usize & (SHARD_COUNT - 1)].lock() } #[cfg(test)] pub fn len(&self) -> usize { - self.entries.len() + self.shards.iter().map(|shard| shard.lock().len()).sum() } #[cfg(test)] pub fn is_empty(&self) -> bool { - self.entries.is_empty() + self.shards.iter().all(|shard| shard.lock().is_empty()) } } @@ -104,37 +115,42 @@ where K: Eq + Hash, S: BuildHasher, { - pub fn get_or_insert(&mut self, key: K) -> &Arc> { + pub fn get_or_insert(&self, key: K) -> Arc> { let hash = self.hasher.hash_one(&key); - self.entries - .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) - .or_insert_with(|| { - Arc::new(OnceTableEntry { - hash, - key, - cell: OnceCell::new(), + let mut shard = self.shard(hash); + Arc::clone( + shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(OnceTableEntry { + hash, + key, + cell: OnceCell::new(), + }) }) - }) - .into_mut() + .into_mut(), + ) } - pub fn get(&self, key: &Q) -> Option<&Arc>> + pub fn get(&self, key: &Q) -> Option>> where K: Borrow, Q: Eq + Hash + ?Sized, { let hash = self.hasher.hash_one(key); - self.entries.find(hash, |entry| entry.key.borrow() == key) + self.shard(hash) + .find(hash, |entry| entry.key.borrow() == key) + .map(Arc::clone) } - pub fn remove(&mut self, key: &Q) -> Option>> + pub fn remove(&self, key: &Q) -> Option>> where K: Borrow, Q: Eq + Hash + ?Sized, { let hash = self.hasher.hash_one(key); - let entry = self - .entries + let mut shard = self.shard(hash); + let entry = shard .find_entry(hash, |entry| entry.key.borrow() == key) .ok()?; let (entry, _) = entry.remove(); @@ -142,18 +158,32 @@ where } /// Removes the entry if the table still contains the same allocation. - pub fn remove_entry(&mut self, entry: &Arc>) { - let Ok(occupied) = self - .entries - .find_entry(entry.hash, |existing| Arc::ptr_eq(existing, entry)) - else { + pub fn remove_entry(&self, entry: &Arc>) { + let mut shard = self.shard(entry.hash); + let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) else { return; }; drop(occupied.remove()); } - pub fn insert(&mut self, key: K, value: V) { + pub fn cleanup_abandoned_entry(&self, entry: Arc>) { + let mut shard = self.shard(entry.hash); + // If the table still owns this entry, a count of two means the current call is its only + // owner outside the table. remove_entry rejects an entry that was detached or replaced. + if Arc::strong_count(&entry) == 2 && !entry.initialized() { + if let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, &entry)) + { + drop(occupied.remove()); + } + } + + // Drop this call's reference before unlocking so a waiting cleanup observes the updated + // reference count. + drop(entry); + } + + pub fn insert(&self, key: K, value: V) { self.remove(&key); let hash = self.hasher.hash_one(&key); @@ -162,6 +192,7 @@ where key, cell: OnceCell::from_value(value), }); - self.entries.insert_unique(hash, entry, |entry| entry.hash); + self.shard(hash) + .insert_unique(hash, entry, |entry| entry.hash); } } diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index d704ac6a..f5426f38 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -21,7 +21,6 @@ use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; -use crate::internal::mutex::Mutex; use crate::internal::once_table::OnceTable; use crate::internal::once_table::OnceTableEntry; @@ -34,7 +33,7 @@ mod tests; /// to wrap the `V` in an `Arc` to make cloning cheap. #[derive(Debug)] pub struct OnceMap { - map: Mutex>, + map: OnceTable, } // Holds one call's entry so Drop can clean it up if the computation is abandoned. @@ -78,15 +77,7 @@ where return; }; - let mut table = self.once_map.map.lock(); - // If the table still owns this entry, a count of two means the current call is its only - // owner outside the table. remove_entry rejects an entry that was detached or replaced. - if Arc::strong_count(&entry) == 2 && !entry.initialized() { - table.remove_entry(&entry); - } - // Drop this call's reference before unlocking so a waiting cleanup observes the updated - // reference count. - drop(entry); + self.once_map.map.cleanup_abandoned_entry(entry); } } @@ -109,17 +100,14 @@ where /// Creates a new OnceMap with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(RandomState::new())), + map: OnceTable::with_hasher(RandomState::new()), } } /// Creates a new OnceMap with the default hasher and the specified capacity. pub fn with_capacity(capacity: usize) -> Self { Self { - map: Mutex::new(OnceTable::with_capacity_and_hasher( - capacity, - RandomState::new(), - )), + map: OnceTable::with_capacity_and_hasher(capacity, RandomState::new()), } } } @@ -133,14 +121,14 @@ where /// Creates a new OnceMap with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(hasher)), + map: OnceTable::with_hasher(hasher), } } /// Create a OnceMap with the specified capacity and hasher. pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { Self { - map: Mutex::new(OnceTable::with_capacity_and_hasher(capacity, hasher)), + map: OnceTable::with_capacity_and_hasher(capacity, hasher), } } @@ -155,14 +143,10 @@ where where F: AsyncFnOnce() -> V, { - let entry = { - let mut map = self.map.lock(); - let entry = map.get_or_insert(key); - if let Some(value) = entry.get() { - return value.clone(); - } - Arc::clone(entry) - }; + let entry = self.map.get_or_insert(key); + if let Some(value) = entry.get() { + return value.clone(); + } let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_init(func).await.clone(); @@ -181,14 +165,10 @@ where where F: AsyncFnOnce() -> Result, { - let entry = { - let mut map = self.map.lock(); - let entry = map.get_or_insert(key); - if let Some(value) = entry.get() { - return Ok(value.clone()); - } - Arc::clone(entry) - }; + let entry = self.map.get_or_insert(key); + if let Some(value) = entry.get() { + return Ok(value.clone()); + } let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_try_init(func).await?.clone(); @@ -202,8 +182,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let map = self.map.lock(); - let entry = map.get(key)?; + let entry = self.map.get(key)?; entry.get().cloned() } @@ -217,8 +196,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); - map.remove(key); + self.map.remove(key); } /// Remove the given key from the map and return a *clone* of the value if exists. @@ -232,7 +210,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.lock().remove(key)?; + let entry = self.map.remove(key)?; entry.get().cloned() } } @@ -244,13 +222,11 @@ where S: Default + BuildHasher, { fn from_iter>(iter: T) -> Self { - let mut map = OnceTable::with_hasher(S::default()); + let map = OnceTable::with_hasher(S::default()); for (key, value) in iter { map.insert(key, value); } - Self { - map: Mutex::new(map), - } + Self { map } } } diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index 1dafc702..b3cb88d1 100644 --- a/asyncband/src/once/once_map/tests.rs +++ b/asyncband/src/once/once_map/tests.rs @@ -29,7 +29,7 @@ async fn failed_compute_removes_empty_entry() { let result: Result = map.try_compute("key", async || Err("fail")).await; assert_eq!(result, Err("fail")); - assert!(map.map.lock().is_empty()); + assert!(map.map.is_empty()); } #[tokio::test] @@ -46,7 +46,7 @@ async fn panicked_compute_removes_empty_entry() { }); assert!(task.await.unwrap_err().is_panic()); - assert!(map.map.lock().is_empty()); + assert!(map.map.is_empty()); } #[tokio::test] @@ -65,11 +65,11 @@ async fn cancelled_compute_removes_empty_entry() { }); started_rx.await.unwrap(); - assert_eq!(map.map.lock().len(), 1); + assert_eq!(map.map.len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(map.map.lock().is_empty()); + assert!(map.map.is_empty()); } #[tokio::test] @@ -91,7 +91,7 @@ async fn failed_compute_preserves_entry_for_waiter_retry() { release_tx.send(()).unwrap(); assert_eq!(first.await, Err("fail")); - assert_eq!(map.map.lock().len(), 1); + assert_eq!(map.map.len(), 1); assert_eq!(retry.await, Ok(1)); assert_eq!(map.get("key"), Some(1)); } diff --git a/asyncband/src/singleflight/mod.rs b/asyncband/src/singleflight/mod.rs index cb117304..cbb37528 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -23,7 +23,6 @@ use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; -use crate::internal::mutex::Mutex; use crate::internal::once_table::OnceTable; use crate::internal::once_table::OnceTableEntry; @@ -34,7 +33,7 @@ mod tests; /// units of work can be executed with duplicate suppression. #[derive(Debug)] pub struct Group { - map: Mutex>, + map: OnceTable, } // Holds one call's entry so Drop can clean it up if the work is abandoned. @@ -53,10 +52,7 @@ where S: BuildHasher, { fn new(group: &'a Group, key: K) -> Self { - let entry = { - let mut map = group.map.lock(); - Arc::clone(map.get_or_insert(key)) - }; + let entry = group.map.get_or_insert(key); Self { group, @@ -83,15 +79,7 @@ where return; }; - let mut table = self.group.map.lock(); - // If the table still owns this entry, a count of two means the current call is its only - // owner outside the table. remove_entry rejects an entry that was detached or replaced. - if Arc::strong_count(&entry) == 2 && !entry.initialized() { - table.remove_entry(&entry); - } - // Drop this call's reference before unlocking so a waiting cleanup observes the updated - // reference count. - drop(entry); + self.group.map.cleanup_abandoned_entry(entry); } } @@ -114,7 +102,7 @@ where /// Creates a new Group with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(RandomState::new())), + map: OnceTable::with_hasher(RandomState::new()), } } } @@ -128,7 +116,7 @@ where /// Creates a new Group with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(hasher)), + map: OnceTable::with_hasher(hasher), } } @@ -193,7 +181,7 @@ where let result = entry .get_or_init(async || { let result = func().await; - self.map.lock().remove_entry(entry); + self.map.remove_entry(entry); result }) .await @@ -257,7 +245,7 @@ where let result = entry .get_or_try_init(async || { let result = func().await?; - self.map.lock().remove_entry(entry); + self.map.remove_entry(entry); Ok(result) }) .await? @@ -275,7 +263,6 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); - map.remove(key); + self.map.remove(key); } } diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index 9f3dbb15..61eb3b1c 100644 --- a/asyncband/src/singleflight/tests.rs +++ b/asyncband/src/singleflight/tests.rs @@ -36,7 +36,7 @@ async fn panicked_work_removes_empty_entry() { }); assert!(task.await.unwrap_err().is_panic()); - assert!(group.map.lock().is_empty()); + assert!(group.map.is_empty()); let result = group.work("key", || async { "success".to_owned() }).await; assert_eq!(result, "success"); @@ -58,11 +58,11 @@ async fn cancelled_work_removes_empty_entry() { }); started_rx.await.unwrap(); - assert_eq!(group.map.lock().len(), 1); + assert_eq!(group.map.len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(group.map.lock().is_empty()); + assert!(group.map.is_empty()); } #[tokio::test] @@ -73,7 +73,7 @@ async fn failed_try_work_removes_empty_entry() { .try_work("key", || async { Err::<&str, &str>("error") }) .await; assert_eq!(result, Err("error")); - assert!(group.map.lock().is_empty()); + assert!(group.map.is_empty()); let retry = group .try_work("key", || async { Ok::<&str, ()>("success") }) @@ -100,7 +100,7 @@ async fn failed_try_work_preserves_entry_for_waiter_retry() { release_tx.send(()).unwrap(); assert_eq!(first.await, Err("fail")); - assert_eq!(group.map.lock().len(), 1); + assert_eq!(group.map.len(), 1); assert_eq!(retry.await, Ok("success")); - assert!(group.map.lock().is_empty()); + assert!(group.map.is_empty()); } diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index db52c620..750c5132 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -16,6 +16,8 @@ // under the License. use std::cell::Cell; +use std::collections::hash_map::DefaultHasher; +use std::hash::BuildHasher; use asyncband::barrier::Barrier; use asyncband::condvar::Condvar; @@ -45,6 +47,17 @@ use asyncband::waitgroup::WaitGroup; struct PoolManager; +struct LocalBuildHasher(Cell); + +impl BuildHasher for LocalBuildHasher { + type Hasher = DefaultHasher; + + fn build_hasher(&self) -> Self::Hasher { + self.0.set(self.0.get().wrapping_add(1)); + DefaultHasher::new() + } +} + impl ManageObject for PoolManager { type Object = i64; type Error = std::convert::Infallible; @@ -109,6 +122,15 @@ fn movable_public_types_are_send() { assert_send::>>(); } +#[tokio::test] +async fn keyed_types_accept_non_sync_hashers_for_local_use() { + let map = OnceMap::<&str, u32, _>::with_hasher(LocalBuildHasher(Cell::new(0))); + assert_eq!(map.compute("key", async || 1).await, 1); + + let group = singleflight::Group::<&str, u32, _>::with_hasher(LocalBuildHasher(Cell::new(0))); + assert_eq!(group.work("key", || async { 1 }).await, 1); +} + #[test] fn public_types_are_unpin() { fn assert_unpin() {} From b816ff67c26c02e06f66904dbd573991a6a6e60f Mon Sep 17 00:00:00 2001 From: An Long Date: Thu, 27 Aug 2026 20:59:38 +0900 Subject: [PATCH 2/2] perf: shard the once table behind read-write locks Combines the sharded table from the previous commit with the shared reader path from the whole-table RwLock approach: hits and duplicate waiters probe their shard under the read lock, initialized compute hits resolve to the value without touching the entry reference count, and only an absent key takes the shard write lock to insert. Co-authored-by: Huliiiiii Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 2 + asyncband/src/internal/mod.rs | 3 + asyncband/src/internal/once_table.rs | 104 +++++++++++++++++++++---- asyncband/src/internal/rwlock.rs | 62 +++++++++++++++ asyncband/src/once/once_map/mod.rs | 20 ++--- asyncband/src/once/once_map/tests.rs | 67 ++++++++++++++++ asyncband/src/singleflight/tests.rs | 52 +++++++++++++ tests-integration/tests/traits_test.rs | 15 ++++ 8 files changed, 299 insertions(+), 26 deletions(-) create mode 100644 asyncband/src/internal/rwlock.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 78bc8acf..ddc9391d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ All notable changes to this project will be documented in this file. * Replace `Semaphore::forget` with `Semaphore::drain_permits` and `Semaphore::forget_exact` with `Semaphore::reduce_permits`; permit-level `forget` methods are unchanged. * Rename `ShutdownSend` and `ShutdownRecv` to `Shutdown` and `ShutdownGuard`; rename `shutdown::new_pair` to `shutdown::new`; make `Shutdown` awaitable for requesting shutdown and awaiting completion; and rename the remaining operations to `request_shutdown`, `watch`, `into_watch`, `is_shutdown_requested`, `shutdown_requested`, and `shutdown_requested_owned`. * Raise the minimum supported Rust version from 1.85.0 to 1.86.0. +* Require the hasher to be `Sync` for `OnceMap` and `singleflight::Group` to be `Sync`; lookups now hash keys and pick shards outside the exclusive lock, so concurrent readers share the hasher. The default `RandomState` and other common hashers are unaffected. ### Bug fixes @@ -31,3 +32,4 @@ All notable changes to this project will be documented in this file. ### Improvements * Remove the `slab` dependency in favor of a focused internal waiter arena. +* Shard the keyed table behind `OnceMap` and `singleflight::Group` and serve initialized hits and duplicate waiters under a shard read lock, so concurrent lookups no longer serialize and operations on different keys proceed in parallel. diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index cdd33e34..8b0142e9 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -44,6 +44,9 @@ pub(crate) mod countdown; #[allow(dead_code)] pub(crate) mod once_table; +#[cfg(any(feature = "once-map", feature = "singleflight"))] +pub(crate) mod rwlock; + #[cfg(any(feature = "lazy-cell", feature = "once-cell"))] // `LazyCell` and `OnceCell` use different subsets of `ValueCell`, so single-feature builds leave // some operations in the shared implementation unused. diff --git a/asyncband/src/internal/once_table.rs b/asyncband/src/internal/once_table.rs index bb01e409..fdbd3b7a 100644 --- a/asyncband/src/internal/once_table.rs +++ b/asyncband/src/internal/once_table.rs @@ -20,11 +20,12 @@ use std::fmt; use std::hash::BuildHasher; use std::hash::Hash; use std::sync::Arc; -use std::sync::MutexGuard; +use std::sync::RwLockReadGuard; +use std::sync::RwLockWriteGuard; use hashbrown::HashTable; -use crate::internal::mutex::Mutex; +use crate::internal::rwlock::RwLock; use crate::once::OnceCell; const SHARD_COUNT: usize = 64; @@ -61,9 +62,16 @@ impl OnceTableEntry { } } +/// Outcome of looking a key up for compute: an initialized entry resolves to its value while the +/// shard lock is still held, so contended hits never touch the entry's shared reference count. +pub enum OnceTableLookup { + Hit(V), + Pending(Arc>), +} + /// Shared keyed storage that lets once primitives clean up an exact entry without cloning its key. pub struct OnceTable { - shards: Box<[Mutex>]>, + shards: Box<[RwLock>]>, hasher: S, } @@ -75,7 +83,7 @@ where fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut debug_map = f.debug_map(); for shard in &self.shards { - let entries = shard.lock(); + let entries = shard.read(); debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); } debug_map.finish() @@ -90,23 +98,27 @@ impl OnceTable { pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { let shard_capacity = capacity.div_ceil(SHARD_COUNT); let shards = (0..SHARD_COUNT) - .map(|_| Mutex::new(HashTable::with_capacity(shard_capacity))) + .map(|_| RwLock::new(HashTable::with_capacity(shard_capacity))) .collect(); Self { shards, hasher } } - fn shard(&self, hash: u64) -> MutexGuard<'_, Entries> { - self.shards[hash as usize & (SHARD_COUNT - 1)].lock() + fn shard_read(&self, hash: u64) -> RwLockReadGuard<'_, Entries> { + self.shards[hash as usize & (SHARD_COUNT - 1)].read() + } + + fn shard_write(&self, hash: u64) -> RwLockWriteGuard<'_, Entries> { + self.shards[hash as usize & (SHARD_COUNT - 1)].write() } #[cfg(test)] pub fn len(&self) -> usize { - self.shards.iter().map(|shard| shard.lock().len()).sum() + self.shards.iter().map(|shard| shard.read().len()).sum() } #[cfg(test)] pub fn is_empty(&self) -> bool { - self.shards.iter().all(|shard| shard.lock().is_empty()) + self.shards.iter().all(|shard| shard.read().is_empty()) } } @@ -117,7 +129,14 @@ where { pub fn get_or_insert(&self, key: K) -> Arc> { let hash = self.hasher.hash_one(&key); - let mut shard = self.shard(hash); + { + let shard = self.shard_read(hash); + if let Some(entry) = shard.find(hash, |entry| entry.key.eq(&key)) { + return Arc::clone(entry); + } + } + + let mut shard = self.shard_write(hash); Arc::clone( shard .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) @@ -138,18 +157,68 @@ where Q: Eq + Hash + ?Sized, { let hash = self.hasher.hash_one(key); - self.shard(hash) + self.shard_read(hash) .find(hash, |entry| entry.key.borrow() == key) .map(Arc::clone) } + /// Clones the value of an initialized entry under the shard read lock, so hits never touch + /// the entry's shared reference count. + pub fn get_value(&self, key: &Q) -> Option + where + K: Borrow, + Q: Eq + Hash + ?Sized, + V: Clone, + { + let hash = self.hasher.hash_one(key); + self.shard_read(hash) + .find(hash, |entry| entry.key.borrow() == key)? + .get() + .cloned() + } + + /// Looks the key up under the shard read lock: an initialized hit resolves to its value + /// without cloning the entry, a pending entry is returned to wait on, and only an absent key + /// takes the shard write lock to insert. + pub fn lookup_or_insert(&self, key: K) -> OnceTableLookup + where + V: Clone, + { + let hash = self.hasher.hash_one(&key); + { + let shard = self.shard_read(hash); + if let Some(entry) = shard.find(hash, |entry| entry.key.eq(&key)) { + if let Some(value) = entry.get() { + return OnceTableLookup::Hit(value.clone()); + } + return OnceTableLookup::Pending(Arc::clone(entry)); + } + } + + let mut shard = self.shard_write(hash); + let entry = shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(OnceTableEntry { + hash, + key, + cell: OnceCell::new(), + }) + }) + .into_mut(); + if let Some(value) = entry.get() { + return OnceTableLookup::Hit(value.clone()); + } + OnceTableLookup::Pending(Arc::clone(entry)) + } + pub fn remove(&self, key: &Q) -> Option>> where K: Borrow, Q: Eq + Hash + ?Sized, { let hash = self.hasher.hash_one(key); - let mut shard = self.shard(hash); + let mut shard = self.shard_write(hash); let entry = shard .find_entry(hash, |entry| entry.key.borrow() == key) .ok()?; @@ -159,7 +228,7 @@ where /// Removes the entry if the table still contains the same allocation. pub fn remove_entry(&self, entry: &Arc>) { - let mut shard = self.shard(entry.hash); + let mut shard = self.shard_write(entry.hash); let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) else { return; }; @@ -168,9 +237,12 @@ where } pub fn cleanup_abandoned_entry(&self, entry: Arc>) { - let mut shard = self.shard(entry.hash); + let mut shard = self.shard_write(entry.hash); // If the table still owns this entry, a count of two means the current call is its only - // owner outside the table. remove_entry rejects an entry that was detached or replaced. + // owner outside the table: entries are only cloned out of their shard while holding the + // shard lock, so the write lock excludes new owners while the count is checked, and + // owners that release outside the lock do so only after the cell is initialized or the + // entry was detached. The ptr_eq probe rejects an entry that was detached or replaced. if Arc::strong_count(&entry) == 2 && !entry.initialized() { if let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, &entry)) { @@ -192,7 +264,7 @@ where key, cell: OnceCell::from_value(value), }); - self.shard(hash) + self.shard_write(hash) .insert_unique(hash, entry, |entry| entry.hash); } } diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs new file mode 100644 index 00000000..bedf51bc --- /dev/null +++ b/asyncband/src/internal/rwlock.rs @@ -0,0 +1,62 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt; +use std::sync::PoisonError; + +pub struct RwLock(std::sync::RwLock); + +impl fmt::Debug for RwLock { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl RwLock { + pub const fn new(t: T) -> Self { + Self(std::sync::RwLock::new(t)) + } + + pub fn read(&self) -> std::sync::RwLockReadGuard<'_, T> { + self.0.read().unwrap_or_else(PoisonError::into_inner) + } + + pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, T> { + self.0.write().unwrap_or_else(PoisonError::into_inner) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::internal::rwlock::RwLock; + + #[test] + fn test_poison_rwlock() { + let rwlock = Arc::new(RwLock::new(42)); + let r = rwlock.clone(); + let handle = std::thread::spawn(move || { + let _guard = r.write(); + panic!("poison"); + }); + let _ = handle.join(); + assert_eq!(*rwlock.read(), 42); + let guard = rwlock.write(); + assert_eq!(*guard, 42); + } +} diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index f5426f38..02d46c42 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use crate::internal::once_table::OnceTable; use crate::internal::once_table::OnceTableEntry; +use crate::internal::once_table::OnceTableLookup; #[cfg(test)] mod tests; @@ -143,10 +144,10 @@ where where F: AsyncFnOnce() -> V, { - let entry = self.map.get_or_insert(key); - if let Some(value) = entry.get() { - return value.clone(); - } + let entry = match self.map.lookup_or_insert(key) { + OnceTableLookup::Hit(value) => return value, + OnceTableLookup::Pending(entry) => entry, + }; let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_init(func).await.clone(); @@ -165,10 +166,10 @@ where where F: AsyncFnOnce() -> Result, { - let entry = self.map.get_or_insert(key); - if let Some(value) = entry.get() { - return Ok(value.clone()); - } + let entry = match self.map.lookup_or_insert(key) { + OnceTableLookup::Hit(value) => return Ok(value), + OnceTableLookup::Pending(entry) => entry, + }; let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_try_init(func).await?.clone(); @@ -182,8 +183,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.get(key)?; - entry.get().cloned() + self.map.get_value(key) } /// Remove the given key from the map. diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index b3cb88d1..ca4341b5 100644 --- a/asyncband/src/once/once_map/tests.rs +++ b/asyncband/src/once/once_map/tests.rs @@ -95,3 +95,70 @@ async fn failed_compute_preserves_entry_for_waiter_retry() { assert_eq!(retry.await, Ok(1)); assert_eq!(map.get("key"), Some(1)); } + +#[tokio::test] +async fn cancelled_waiter_preserves_pending_entry() { + let map = OnceMap::<&str, i32>::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let leader = map.compute("key", async move || { + release_rx.await.unwrap(); + 1 + }); + tokio::pin!(leader); + assert!(poll_once(leader.as_mut()).is_pending()); + + let mut waiter = Box::pin(map.compute("key", async || unreachable!())); + assert!(poll_once(waiter.as_mut()).is_pending()); + drop(waiter); + + assert_eq!(map.map.len(), 1); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + assert_eq!(map.get("key"), Some(1)); +} + +#[tokio::test] +async fn cancelled_waiter_preserves_initialized_entry() { + let map = OnceMap::<&str, i32>::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let leader = map.compute("key", async move || { + release_rx.await.unwrap(); + 1 + }); + tokio::pin!(leader); + assert!(poll_once(leader.as_mut()).is_pending()); + + let mut waiter = Box::pin(map.compute("key", async || unreachable!())); + assert!(poll_once(waiter.as_mut()).is_pending()); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + + drop(waiter); + + assert_eq!(map.map.len(), 1); + assert_eq!(map.get("key"), Some(1)); +} + +#[tokio::test] +async fn cancelled_compute_preserves_replacement_entry() { + let map = OnceMap::<&str, i32>::new(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + + let mut first = Box::pin(map.compute("key", async move || { + release_rx.await.unwrap(); + 1 + })); + assert!(poll_once(first.as_mut()).is_pending()); + + map.discard("key"); + assert_eq!(map.compute("key", async || 2).await, 2); + + drop(first); + + assert_eq!(map.map.len(), 1); + assert_eq!(map.get("key"), Some(2)); +} diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index 61eb3b1c..0139e3fd 100644 --- a/asyncband/src/singleflight/tests.rs +++ b/asyncband/src/singleflight/tests.rs @@ -104,3 +104,55 @@ async fn failed_try_work_preserves_entry_for_waiter_retry() { assert_eq!(retry.await, Ok("success")); assert!(group.map.is_empty()); } + +#[tokio::test] +async fn cancelled_waiter_preserves_inflight_entry() { + let group = Group::<&str, i32>::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let leader = group.work("key", || async move { + release_rx.await.unwrap(); + 1 + }); + tokio::pin!(leader); + assert!(poll_once(leader.as_mut()).is_pending()); + + let mut waiter = Box::pin(group.work("key", || async { unreachable!() })); + assert!(poll_once(waiter.as_mut()).is_pending()); + drop(waiter); + + assert_eq!(group.map.len(), 1); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + assert!(group.map.is_empty()); +} + +#[tokio::test] +async fn cancelled_work_preserves_replacement_entry() { + let group = Group::<&str, i32>::new(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + let (replacement_tx, replacement_rx) = tokio::sync::oneshot::channel(); + + let mut first = Box::pin(group.work("key", || async move { + release_rx.await.unwrap(); + 1 + })); + assert!(poll_once(first.as_mut()).is_pending()); + + group.forget("key"); + + let second = group.work("key", || async move { + replacement_rx.await.unwrap(); + 2 + }); + tokio::pin!(second); + assert!(poll_once(second.as_mut()).is_pending()); + + drop(first); + assert_eq!(group.map.len(), 1); + + replacement_tx.send(()).unwrap(); + assert_eq!(second.await, 2); + assert!(group.map.is_empty()); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 750c5132..20a6ba5c 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -58,6 +58,19 @@ impl BuildHasher for LocalBuildHasher { } } +// A hasher that is Send but not Sync. OnceMap and Group hash keys and pick shards outside the +// exclusive lock, so being Sync requires the hasher to be Sync, while being Send does not. +#[allow(dead_code)] +struct SendOnlyState(std::marker::PhantomData>); + +impl BuildHasher for SendOnlyState { + type Hasher = DefaultHasher; + + fn build_hasher(&self) -> Self::Hasher { + DefaultHasher::new() + } +} + impl ManageObject for PoolManager { type Object = i64; type Error = std::convert::Infallible; @@ -117,6 +130,8 @@ fn movable_public_types_are_send() { fn assert_send() {} assert_send::>>(); + assert_send::>(); + assert_send::>(); assert_send::>(); assert_send::>(); assert_send::>>();