From 8317a740f206cbd929496ed950166c359b7adc2c Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Fri, 28 Aug 2026 03:11:55 +0800 Subject: [PATCH 1/9] benchmark: refine OnceMap and singleflight multithreaded benchmarks --- benchmarks/asyncband/once_map/compute.rs | 191 ++++++++++++---------- benchmarks/asyncband/once_map/get.rs | 16 +- benchmarks/asyncband/once_map/support.rs | 13 +- benchmarks/asyncband/singleflight/work.rs | 39 +---- benchmarks/asyncband/support.rs | 15 -- 5 files changed, 131 insertions(+), 143 deletions(-) diff --git a/benchmarks/asyncband/once_map/compute.rs b/benchmarks/asyncband/once_map/compute.rs index e5733fc..71ffd2c 100644 --- a/benchmarks/asyncband/once_map/compute.rs +++ b/benchmarks/asyncband/once_map/compute.rs @@ -17,11 +17,12 @@ use std::cell::Cell; -use asyncband::once::OnceMap; use divan::Bencher; use divan::black_box; +use super::support::BenchMap; use super::support::CONTENDED_ENTRY_COUNTS; +use super::support::CONTENDED_THREAD_SLOTS; use super::support::THREAD_COUNTS; use super::support::preloaded_map; use crate::support::bench_context; @@ -32,18 +33,28 @@ use crate::support::poll_ready; use crate::support::spin_poll_ready; use crate::support::thread_slot_ticket; use crate::support::wait_until_open; -use crate::support::yield_polls; const CACHED_ENTRY_COUNTS: &[usize] = &[0, 64, 1024]; -const WAITER_COUNTS: &[usize] = &[1, 8, 32]; -const MISS_KEY_SPAN: usize = 1 << 16; -const COALESCED_LEADER_POLLS: usize = 32; +const COMPUTATION_COUNTS: &[usize] = &[2, 9, 33]; +const MISS_KEYSPACE_SIZE: usize = 1 << 16; + +enum MixedInput { + Hit(usize), + Miss(usize), +} + +fn miss_key(cached_entries: usize, slot: usize, ticket: usize) -> usize { + // Workers start at different phases but traverse the same keyspace. + let thread_offset = + slot % CONTENDED_THREAD_SLOTS * (MISS_KEYSPACE_SIZE / CONTENDED_THREAD_SLOTS); + cached_entries + (thread_offset + ticket) % MISS_KEYSPACE_SIZE +} #[divan::bench] fn compute_vacant(bencher: Bencher) { let mut context = bench_context(); bencher - .with_inputs(OnceMap::::new) + .with_inputs(BenchMap::default) .bench_local_values(|map| { let result = black_box(poll_ready( map.compute(black_box(0), || async { black_box(1) }), @@ -57,7 +68,7 @@ fn compute_vacant(bencher: Bencher) { fn compute_occupied(bencher: Bencher) { let mut context = bench_context(); bencher - .with_inputs(|| [(0, 1)].into_iter().collect::>()) + .with_inputs(|| [(0, 1)].into_iter().collect::()) .bench_local_values(|map| { let result = black_box(poll_ready( map.compute(black_box(0), || async { black_box(2) }), @@ -74,7 +85,7 @@ fn try_compute_error(bencher: Bencher, cached_entries: usize) { .with_inputs(|| { (0..cached_entries) .map(|key| (key, key)) - .collect::>() + .collect::() }) .bench_local_values(|map| { let result = black_box(poll_ready( @@ -85,38 +96,61 @@ fn try_compute_error(bencher: Bencher, cached_entries: usize) { }); } -#[divan::bench(args = WAITER_COUNTS)] -fn coalesced_compute_batch(bencher: Bencher, waiter_count: usize) { +#[divan::bench(args = COMPUTATION_COUNTS)] +fn coalesced_compute_batch(bencher: Bencher, computation_count: usize) { let mut context = bench_context(); bencher.bench_local(|| { - let map = OnceMap::::new(); + let map = BenchMap::default(); let gate = Cell::new(false); - let mut leader = Box::pin(map.compute(0, || async { - wait_until_open(&gate).await; - black_box(1usize) - })); - poll_pending(leader.as_mut(), &mut context); - - let mut waiters = (0..waiter_count) - .map(|_| Box::pin(map.compute(0, || async { unreachable!() }))) + let mut computations = (0..computation_count) + .map(|_| { + Box::pin(map.compute(0, || async { + wait_until_open(&gate).await; + black_box(1usize) + })) + }) .collect::>(); - for waiter in &mut waiters { - poll_pending(waiter.as_mut(), &mut context); + for computation in &mut computations { + poll_pending(computation.as_mut(), &mut context); } gate.set(true); - black_box(poll_pinned_ready(leader.as_mut(), &mut context)); - drop(leader); - for mut waiter in waiters { - black_box(poll_pinned_ready(waiter.as_mut(), &mut context)); + for mut computation in computations { + black_box(poll_pinned_ready(computation.as_mut(), &mut context)); + } + }); +} + +#[divan::bench(args = COMPUTATION_COUNTS)] +fn independent_compute_batch(bencher: Bencher, computation_count: usize) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let map = BenchMap::default(); + let gate = Cell::new(false); + let mut computations = (0..computation_count) + .map(|key| { + Box::pin(map.compute(key, || async { + wait_until_open(&gate).await; + black_box(1usize) + })) + }) + .collect::>(); + for computation in &mut computations { + poll_pending(computation.as_mut(), &mut context); + } + + gate.set(true); + for mut computation in computations { + black_box(poll_pinned_ready(computation.as_mut(), &mut context)); } }); } #[divan::bench(threads = THREAD_COUNTS)] fn contended_compute_hit_same_key(bencher: Bencher) { - let map = [(0, 1)].into_iter().collect::>(); + let map = [(0, 1)].into_iter().collect::(); bencher.bench(|| { let mut context = bench_context(); @@ -131,75 +165,68 @@ fn contended_compute_hit_same_key(bencher: Bencher) { fn contended_compute_hit_disjoint(bencher: Bencher, cached_entries: usize) { let map = preloaded_map(cached_entries); - bencher.bench(|| { - let mut context = bench_context(); - let (slot, ticket) = thread_slot_ticket(); - let key = (slot + ticket) % cached_entries; - black_box(spin_poll_ready( - map.compute(black_box(key), || async { unreachable!() }), - &mut context, - )) - }); + bencher + .with_inputs(|| { + let (slot, ticket) = thread_slot_ticket(); + (slot + ticket * CONTENDED_THREAD_SLOTS) % cached_entries + }) + .bench_values(|key| { + let mut context = bench_context(); + black_box(spin_poll_ready( + map.compute(black_box(key), || async { unreachable!() }), + &mut context, + )) + }); } #[divan::bench(threads = THREAD_COUNTS, args = CONTENDED_ENTRY_COUNTS)] fn contended_compute_miss_churn(bencher: Bencher, cached_entries: usize) { let map = preloaded_map(cached_entries); - bencher.bench(|| { - let mut context = bench_context(); - let (slot, ticket) = thread_slot_ticket(); - let key = cached_entries + slot * MISS_KEY_SPAN + ticket % MISS_KEY_SPAN; - let value = spin_poll_ready( - map.compute(black_box(key), || async move { key }), - &mut context, - ); - map.discard(&key); - black_box(value) - }); -} - -#[divan::bench(threads = THREAD_COUNTS, args = CONTENDED_ENTRY_COUNTS)] -fn contended_compute_mixed(bencher: Bencher, cached_entries: usize) { - let map = preloaded_map(cached_entries); - - bencher.bench(|| { - let mut context = bench_context(); - let (slot, ticket) = thread_slot_ticket(); - if ticket % 2 == 0 { - let key = (slot + ticket / 2) % cached_entries; - black_box(spin_poll_ready( - map.compute(black_box(key), || async { unreachable!() }), - &mut context, - )) - } else { - let key = cached_entries + slot * MISS_KEY_SPAN + (ticket / 2) % MISS_KEY_SPAN; + bencher + .with_inputs(|| { + let (slot, ticket) = thread_slot_ticket(); + miss_key(cached_entries, slot, ticket) + }) + .bench_values(|key| { + let mut context = bench_context(); let value = spin_poll_ready( map.compute(black_box(key), || async move { key }), &mut context, ); map.discard(&key); black_box(value) - } - }); + }); } -// The leader stays in flight for several polls so calls on other threads coalesce as duplicate -// waiters, and discards the key while in flight so every cycle re-runs the vacant-leader path -// instead of settling into steady-state hits. -#[divan::bench(threads = THREAD_COUNTS)] -fn contended_compute_coalesced(bencher: Bencher) { - let map = OnceMap::::new(); +#[divan::bench(threads = THREAD_COUNTS, args = CONTENDED_ENTRY_COUNTS)] +fn contended_compute_mixed(bencher: Bencher, cached_entries: usize) { + let map = preloaded_map(cached_entries); - bencher.bench(|| { - let mut context = bench_context(); - black_box(spin_poll_ready( - map.compute(black_box(0), || async { - yield_polls(COALESCED_LEADER_POLLS).await; - map.discard(&0); - black_box(1) - }), - &mut context, - )) - }); + bencher + .with_inputs(|| { + let (slot, ticket) = thread_slot_ticket(); + if (slot + ticket) % 2 == 0 { + MixedInput::Hit((slot + ticket / 2 * CONTENDED_THREAD_SLOTS) % cached_entries) + } else { + MixedInput::Miss(miss_key(cached_entries, slot, ticket / 2)) + } + }) + .bench_values(|input| { + let mut context = bench_context(); + match input { + MixedInput::Hit(key) => black_box(spin_poll_ready( + map.compute(black_box(key), || async { unreachable!() }), + &mut context, + )), + MixedInput::Miss(key) => { + let value = spin_poll_ready( + map.compute(black_box(key), || async move { key }), + &mut context, + ); + map.discard(&key); + black_box(value) + } + } + }); } diff --git a/benchmarks/asyncband/once_map/get.rs b/benchmarks/asyncband/once_map/get.rs index cef31a0..287eb14 100644 --- a/benchmarks/asyncband/once_map/get.rs +++ b/benchmarks/asyncband/once_map/get.rs @@ -15,18 +15,19 @@ // specific language governing permissions and limitations // under the License. -use asyncband::once::OnceMap; use divan::Bencher; use divan::black_box; +use super::support::BenchMap; use super::support::CONTENDED_ENTRY_COUNTS; +use super::support::CONTENDED_THREAD_SLOTS; use super::support::THREAD_COUNTS; use super::support::preloaded_map; use crate::support::thread_slot_ticket; #[divan::bench(threads = THREAD_COUNTS)] fn contended_get_hit_same_key(bencher: Bencher) { - let map = [(0, 1)].into_iter().collect::>(); + let map = [(0, 1)].into_iter().collect::(); bencher.bench(|| black_box(map.get(black_box(&0)))); } @@ -35,9 +36,10 @@ fn contended_get_hit_same_key(bencher: Bencher) { fn contended_get_hit_disjoint(bencher: Bencher, cached_entries: usize) { let map = preloaded_map(cached_entries); - bencher.bench(|| { - let (slot, ticket) = thread_slot_ticket(); - let key = (slot + ticket) % cached_entries; - black_box(map.get(black_box(&key))) - }); + bencher + .with_inputs(|| { + let (slot, ticket) = thread_slot_ticket(); + (slot + ticket * CONTENDED_THREAD_SLOTS) % cached_entries + }) + .bench_values(|key| black_box(map.get(black_box(&key)))); } diff --git a/benchmarks/asyncband/once_map/support.rs b/benchmarks/asyncband/once_map/support.rs index 44c59c8..b5b155b 100644 --- a/benchmarks/asyncband/once_map/support.rs +++ b/benchmarks/asyncband/once_map/support.rs @@ -15,14 +15,19 @@ // specific language governing permissions and limitations // under the License. +use std::collections::hash_map::DefaultHasher; +use std::hash::BuildHasherDefault; + use asyncband::once::OnceMap; pub const CONTENDED_ENTRY_COUNTS: &[usize] = &[64, 1024]; +pub const CONTENDED_THREAD_SLOTS: usize = 32; pub const THREAD_COUNTS: &[usize] = &[1, 2, 8, 32]; -// The contended get and compute benches share one map across OS threads and spread keys with -// thread_slot_ticket, so "disjoint" means threads mostly touch different keys at any moment rather -// than strict per-thread key ownership. -pub fn preloaded_map(cached_entries: usize) -> OnceMap { +type BenchHasher = BuildHasherDefault; + +pub type BenchMap = OnceMap; + +pub fn preloaded_map(cached_entries: usize) -> BenchMap { (0..cached_entries).map(|key| (key, key)).collect() } diff --git a/benchmarks/asyncband/singleflight/work.rs b/benchmarks/asyncband/singleflight/work.rs index 62ffefb..7d70600 100644 --- a/benchmarks/asyncband/singleflight/work.rs +++ b/benchmarks/asyncband/singleflight/work.rs @@ -16,6 +16,8 @@ // under the License. use std::cell::Cell; +use std::hash::BuildHasherDefault; +use std::hash::DefaultHasher; use asyncband::singleflight::Group; use divan::Bencher; @@ -29,12 +31,10 @@ use crate::support::poll_ready; use crate::support::spin_poll_ready; use crate::support::thread_slot_ticket; use crate::support::wait_until_open; -use crate::support::yield_polls; const WAITER_COUNTS: &[usize] = &[1, 8, 32]; const THREAD_COUNTS: &[usize] = &[1, 2, 8, 32]; const DISJOINT_KEY_SPAN: usize = 1 << 16; -const COALESCED_LEADER_POLLS: usize = 32; #[divan::bench] fn work_ready(bencher: Bencher) { @@ -67,9 +67,9 @@ fn try_work_error(bencher: Bencher) { #[divan::bench(args = WAITER_COUNTS)] fn coalesced_work_batch(bencher: Bencher, waiter_count: usize) { let mut context = bench_context(); + let group = Group::::new(); bencher.bench_local(|| { - let group = Group::::new(); let gate = Cell::new(false); let mut leader = Box::pin(group.work(0, || async { wait_until_open(&gate).await; @@ -93,40 +93,9 @@ fn coalesced_work_batch(bencher: Bencher, waiter_count: usize) { }); } -#[divan::bench(threads = THREAD_COUNTS)] -fn contended_work_same_key(bencher: Bencher) { - let group = Group::::new(); - - bencher.bench(|| { - let mut context = bench_context(); - black_box(spin_poll_ready( - group.work(black_box(0), || async { black_box(1) }), - &mut context, - )) - }); -} - -// The leader stays in flight for several polls so calls on other threads coalesce as duplicate -// waiters instead of leading their own cycles. -#[divan::bench(threads = THREAD_COUNTS)] -fn contended_work_coalesced(bencher: Bencher) { - let group = Group::::new(); - - bencher.bench(|| { - let mut context = bench_context(); - black_box(spin_poll_ready( - group.work(black_box(0), || async { - yield_polls(COALESCED_LEADER_POLLS).await; - black_box(1) - }), - &mut context, - )) - }); -} - #[divan::bench(threads = THREAD_COUNTS)] fn contended_work_disjoint_churn(bencher: Bencher) { - let group = Group::::new(); + let group = Group::>::default(); bencher.bench(|| { let mut context = bench_context(); diff --git a/benchmarks/asyncband/support.rs b/benchmarks/asyncband/support.rs index ce12a63..cba7cbd 100644 --- a/benchmarks/asyncband/support.rs +++ b/benchmarks/asyncband/support.rs @@ -95,21 +95,6 @@ pub(super) fn thread_slot_ticket() -> (usize, usize) { (slot, ticket) } -// Stays pending for the given number of polls without registering a waker, so it must only be -// polled through spin_poll_ready. Keeps a leader in flight long enough for calls on other threads -// to join as waiters. -pub(super) async fn yield_polls(mut polls: usize) { - poll_fn(move |_| { - if polls == 0 { - Poll::Ready(()) - } else { - polls -= 1; - Poll::Pending - } - }) - .await -} - // Move the input into the benchmark output so Divan drops it outside the timed section. #[inline] pub(super) fn defer_input_drop(input: I, output: O) -> (I, O) { From be50824e8e920860f66834e9d6c5ad3a976bf1a5 Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Thu, 27 Aug 2026 01:04:40 +0800 Subject: [PATCH 2/9] 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 ++-- 5 files changed, 105 insertions(+), 111 deletions(-) diff --git a/asyncband/src/internal/once_table.rs b/asyncband/src/internal/once_table.rs index 27865ed..bb01e40 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 d704ac6..f5426f3 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 1dafc70..b3cb88d 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 cb11730..cbb3752 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 9f3dbb1..61eb3b1 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()); } From d1c5c2b481e343815ebc4b83f55f9fff00dbe84a Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Thu, 27 Aug 2026 19:30:11 +0800 Subject: [PATCH 3/9] perf: specialize OnceMap and singleflight storage --- Cargo.lock | 26 ++ Cargo.toml | 1 + asyncband/Cargo.toml | 3 +- asyncband/src/internal/mod.rs | 16 +- asyncband/src/internal/mutex.rs | 7 + asyncband/src/internal/once_table.rs | 198 -------------- asyncband/src/once/once_map/mod.rs | 48 ++-- asyncband/src/once/once_map/table.rs | 394 +++++++++++++++++++++++++++ asyncband/src/singleflight/mod.rs | 19 +- asyncband/src/singleflight/table.rs | 161 +++++++++++ 10 files changed, 634 insertions(+), 239 deletions(-) delete mode 100644 asyncband/src/internal/once_table.rs create mode 100644 asyncband/src/once/once_map/table.rs create mode 100644 asyncband/src/singleflight/table.rs diff --git a/Cargo.lock b/Cargo.lock index 0c320cf..eddb010 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -69,6 +69,7 @@ name = "asyncband" version = "0.6.7" dependencies = [ "hashbrown", + "scc", "tokio", ] @@ -789,12 +790,37 @@ dependencies = [ "untrusted", ] +[[package]] +name = "saa" +version = "5.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f5acb362a0e75c2a963532fa7fabf13dff81626dc494df16488d30befcbea0" + +[[package]] +name = "scc" +version = "3.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8af0b99483d1c3e59471d4f0cb58b244169436a8979c889a91a3f697075ea01" +dependencies = [ + "saa", + "sdd", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sdd" +version = "4.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f95d4cc3459db1e608946153c95cf20158af0b86131ae4ae451f90f54549e7d1" +dependencies = [ + "saa", +] + [[package]] name = "semver" version = "1.0.28" diff --git a/Cargo.toml b/Cargo.toml index 872eae1..7fe5b5e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ asyncband = { path = "asyncband" } # Optional runtime dependencies hashbrown = { version = "0.17.1", default-features = false } +scc = "3.8.6" # Dev dependencies async-channel = { version = "2.5.0" } diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index 3871529..b2886c9 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -53,7 +53,7 @@ mpsc = [] mutex = [] once = ["semaphore"] once-cell = ["semaphore"] -once-map = ["dep:hashbrown", "once-cell"] +once-map = ["dep:hashbrown", "dep:scc", "once-cell"] oneshot = [] pool = ["semaphore"] rwlock = [] @@ -66,6 +66,7 @@ waitgroup = [] hashbrown = { workspace = true, default-features = false, features = [ "inline-more", ], optional = true } +scc = { workspace = true, optional = true } [dev-dependencies] tokio = { workspace = true, features = ["full"] } diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index cdd33e3..07ae169 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -38,12 +38,6 @@ mod arena; #[allow(dead_code)] pub(crate) mod countdown; -#[cfg(any(feature = "once-map", feature = "singleflight"))] -// `OnceMap` and `singleflight` use different subsets of `OnceTable`, so single-feature builds -// leave some operations in the shared implementation unused. -#[allow(dead_code)] -pub(crate) mod once_table; - #[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. @@ -91,3 +85,13 @@ pub(crate) mod waitlist; // `new`. One constructor is therefore unused in every single-primitive build. #[allow(dead_code)] pub(crate) mod waitset; + +#[cfg(any(feature = "once-map", feature = "singleflight"))] +pub fn default_shard_count() -> usize { + // Tested on a 32-core machine, the optimal shard count for `OnceMap` and `Singleflight` is 256. + // So I use 8 as the coefficient, which is 256 / 32. + // Need to test on other machines to see if this coefficient is optimal. + // Dashmap use 4. + (std::thread::available_parallelism().map_or(1, |parallelism| parallelism.get()) * 8) + .next_power_of_two() +} diff --git a/asyncband/src/internal/mutex.rs b/asyncband/src/internal/mutex.rs index 9477e26..73577fb 100644 --- a/asyncband/src/internal/mutex.rs +++ b/asyncband/src/internal/mutex.rs @@ -36,6 +36,13 @@ impl Mutex { } } +#[cfg(any(feature = "once-map", feature = "singleflight"))] +/// Alignment uses 60% more memory (64/40) but improves write performance by 25% at 32 threads. (On +/// Zen 5 CPUs) +/// Need to test on other architectures. +#[repr(align(64))] +pub struct CachePaddedMutex(pub Mutex); + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/asyncband/src/internal/once_table.rs b/asyncband/src/internal/once_table.rs deleted file mode 100644 index bb01e40..0000000 --- a/asyncband/src/internal/once_table.rs +++ /dev/null @@ -1,198 +0,0 @@ -// 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::borrow::Borrow; -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, - cell: OnceCell, -} - -impl OnceTableEntry { - pub fn initialized(&self) -> bool { - self.cell.initialized() - } - - pub fn get(&self) -> Option<&V> { - self.cell.get() - } - - pub async fn get_or_init(&self, init: F) -> &V - where - F: AsyncFnOnce() -> V, - { - self.cell.get_or_init(init).await - } - - pub async fn get_or_try_init(&self, init: F) -> Result<&V, E> - where - F: AsyncFnOnce() -> Result, - { - self.cell.get_or_try_init(init).await - } -} - -/// Shared keyed storage that lets once primitives clean up an exact entry without cloning its key. -pub struct OnceTable { - shards: Box<[Mutex>]>, - hasher: S, -} - -impl fmt::Debug for OnceTable -where - K: fmt::Debug, - V: fmt::Debug, -{ - 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(); - debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); - } - debug_map.finish() - } -} - -impl OnceTable { - pub fn with_hasher(hasher: S) -> Self { - Self::with_capacity_and_hasher(0, hasher) - } - - 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))) - .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.shards.iter().map(|shard| shard.lock().len()).sum() - } - - #[cfg(test)] - pub fn is_empty(&self) -> bool { - self.shards.iter().all(|shard| shard.lock().is_empty()) - } -} - -impl OnceTable -where - K: Eq + Hash, - S: BuildHasher, -{ - pub fn get_or_insert(&self, key: K) -> Arc> { - let hash = self.hasher.hash_one(&key); - 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(), - ) - } - - pub fn get(&self, key: &Q) -> Option>> - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - let hash = self.hasher.hash_one(key); - self.shard(hash) - .find(hash, |entry| entry.key.borrow() == key) - .map(Arc::clone) - } - - 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 entry = shard - .find_entry(hash, |entry| entry.key.borrow() == key) - .ok()?; - let (entry, _) = entry.remove(); - Some(entry) - } - - /// 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 Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) else { - return; - }; - - drop(occupied.remove()); - } - - 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); - let entry = Arc::new(OnceTableEntry { - hash, - key, - cell: OnceCell::from_value(value), - }); - 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 f5426f3..11c3848 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -21,9 +21,11 @@ use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; -use crate::internal::once_table::OnceTable; -use crate::internal::once_table::OnceTableEntry; +use table::Entry; +use table::Lookup; +use table::Table; +mod table; #[cfg(test)] mod tests; @@ -33,7 +35,7 @@ mod tests; /// to wrap the `V` in an `Arc` to make cloning cheap. #[derive(Debug)] pub struct OnceMap { - map: OnceTable, + map: Table, } // Holds one call's entry so Drop can clean it up if the computation is abandoned. @@ -43,7 +45,7 @@ where S: BuildHasher, { once_map: &'a OnceMap, - entry: Option>>, + entry: Option>>, } impl<'a, K, V, S> ComputeCleanupGuard<'a, K, V, S> @@ -51,14 +53,14 @@ where K: Eq + Hash, S: BuildHasher, { - fn new(once_map: &'a OnceMap, entry: Arc>) -> Self { + fn new(once_map: &'a OnceMap, entry: Arc>) -> Self { Self { once_map, entry: Some(entry), } } - fn entry(&self) -> &Arc> { + fn entry(&self) -> &Arc> { self.entry.as_ref().unwrap() } @@ -100,14 +102,14 @@ where /// Creates a new OnceMap with the default hasher. pub fn new() -> Self { Self { - map: OnceTable::with_hasher(RandomState::new()), + map: Table::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: OnceTable::with_capacity_and_hasher(capacity, RandomState::new()), + map: Table::with_capacity_and_hasher(capacity, RandomState::new()), } } } @@ -121,14 +123,14 @@ where /// Creates a new OnceMap with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: OnceTable::with_hasher(hasher), + map: Table::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: OnceTable::with_capacity_and_hasher(capacity, hasher), + map: Table::with_capacity_and_hasher(capacity, hasher), } } @@ -143,10 +145,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.get_or_insert(key) { + Lookup::Ready(value) => return value, + Lookup::Entry(entry) => entry, + }; let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_init(func).await.clone(); @@ -165,10 +167,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.get_or_insert(key) { + Lookup::Ready(value) => return Ok(value), + Lookup::Entry(entry) => entry, + }; let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_try_init(func).await?.clone(); @@ -182,8 +184,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.get(key)?; - entry.get().cloned() + self.map.get(key) } /// Remove the given key from the map. @@ -222,11 +223,8 @@ where S: Default + BuildHasher, { fn from_iter>(iter: T) -> Self { - let map = OnceTable::with_hasher(S::default()); - for (key, value) in iter { - map.insert(key, value); + Self { + map: iter.into_iter().collect(), } - - Self { map } } } diff --git a/asyncband/src/once/once_map/table.rs b/asyncband/src/once/once_map/table.rs new file mode 100644 index 0000000..ac7e634 --- /dev/null +++ b/asyncband/src/once/once_map/table.rs @@ -0,0 +1,394 @@ +// 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::borrow::Borrow; +use std::fmt; +use std::hash::BuildHasher; +use std::hash::BuildHasherDefault; +use std::hash::Hash; +use std::hash::Hasher; +use std::panic::UnwindSafe; +use std::sync::Arc; +use std::sync::MutexGuard; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use hashbrown::HashTable; +use scc::Equivalent; +use scc::HashIndex; +use scc::hash_index::Entry as IndexEntry; + +use crate::internal::default_shard_count; +use crate::internal::mutex::CachePaddedMutex; +use crate::internal::mutex::Mutex; +use crate::once::OnceCell; + +type Entries = HashTable>>; + +pub struct Entry { + hash: u64, + key: K, + cell: OnceCell, + was_indexed: AtomicBool, +} + +pub enum Lookup { + Ready(V), + Entry(Arc>), +} + +impl Entry { + pub fn initialized(&self) -> bool { + self.cell.initialized() + } + + pub fn get(&self) -> Option<&V> { + self.cell.get() + } + + pub async fn get_or_init(&self, init: F) -> &V + where + F: AsyncFnOnce() -> V, + { + self.cell.get_or_init(init).await + } + + pub async fn get_or_try_init(&self, init: F) -> Result<&V, E> + where + F: AsyncFnOnce() -> Result, + { + self.cell.get_or_try_init(init).await + } +} + +struct ReadyEntry(Arc>); + +impl PartialEq for ReadyEntry { + fn eq(&self, other: &Self) -> bool { + self.0.hash == other.0.hash && self.0.key == other.0.key + } +} + +impl Eq for ReadyEntry {} + +impl Hash for ReadyEntry { + fn hash(&self, state: &mut H) { + state.write_u64(self.0.hash); + } +} + +type BuildIdentityHasher = BuildHasherDefault; + +#[derive(Default)] +struct IdentityHasher(u64); + +impl Hasher for IdentityHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, bytes: &[u8]) { + for byte in bytes { + self.0 = self.0.rotate_left(8) ^ u64::from(*byte); + } + } + + fn write_u64(&mut self, value: u64) { + self.0 = value; + } +} + +pub struct Table { + shards: Box<[CachePaddedMutex>]>, + index: HashIndex, (), BuildIdentityHasher>, + hasher: S, +} + +/// `HashIndex` prevents `Table` from being automatically `UnwindSafe` unless `K` and `V` are +/// `UnwindSafe`. +/// Table operations are unwind-safe regardless, but since it was refactored from a +/// mutex-backed implementation, implement `UnwindSafe` manually to retain the same auto-trait +/// semantics. +impl UnwindSafe for Table {} + +impl fmt::Debug for Table +where + K: fmt::Debug, + V: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut debug_map = f.debug_map(); + for shard in &self.shards { + let entries = shard.0.lock(); + debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); + } + debug_map.finish() + } +} + +impl Table { + pub fn with_hasher(hasher: S) -> Self { + Self::with_capacity_and_hasher(0, hasher) + } + + pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { + let shard_count = default_shard_count(); + let shard_capacity = capacity.div_ceil(shard_count); + let shards = (0..shard_count) + .map(|_| CachePaddedMutex(Mutex::new(HashTable::with_capacity(shard_capacity)))) + .collect(); + + Self { + shards, + index: HashIndex::with_capacity_and_hasher(capacity, BuildIdentityHasher::default()), + hasher, + } + } + + fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { + self.shards[(hash as usize) & (self.shards.len() - 1)] + .0 + .lock() + } + + #[cfg(test)] + pub fn len(&self) -> usize { + self.shards.iter().map(|shard| shard.0.lock().len()).sum() + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + self.shards.iter().all(|shard| shard.0.lock().is_empty()) + } +} + +impl Table +where + K: Eq + Hash, + S: BuildHasher, +{ + fn lockup_index(&self, hash: u64, key: &Q) -> Option + where + K: Borrow, + Q: Eq + ?Sized, + V: Clone, + { + // I think we need to add a function-based peek_with to scc instead of using new type here. + // Alternatively, we could write our own HashIndex. + struct LookupKey<'a, Q: ?Sized> { + hash: u64, + key: &'a Q, + } + + impl Hash for LookupKey<'_, Q> { + fn hash(&self, state: &mut H) { + state.write_u64(self.hash); + } + } + + impl Equivalent> for LookupKey<'_, Q> + where + K: Borrow, + Q: Eq + ?Sized, + { + fn equivalent(&self, entry: &ReadyEntry) -> bool { + entry.0.key.borrow() == self.key + } + } + + self.index + .peek_with(&LookupKey { hash, key }, |entry, ()| entry.0.get().cloned()) + .flatten() + } + + pub fn get_or_insert(&self, key: K) -> Lookup + where + V: Clone, + { + let hash = self.hasher.hash_one(&key); + if let Some(value) = self.lockup_index(hash, &key) { + return Lookup::Ready(value); + } + + let entry = { + let mut shard = self.lock_shard(hash); + shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(Entry { + hash, + key, + cell: OnceCell::new(), + was_indexed: AtomicBool::new(false), + }) + }) + .into_mut() + .clone() + }; + + match self.index_if_ready(&entry) { + Some(value) => Lookup::Ready(value), + None => Lookup::Entry(entry), + } + } + + pub fn get(&self, key: &Q) -> Option + where + K: Borrow, + Q: Eq + Hash + ?Sized, + V: Clone, + { + let hash = self.hasher.hash_one(key); + + if let Some(value) = self.lockup_index(hash, key) { + Some(value) + } else { + let entry = self + .lock_shard(hash) + .find(hash, |entry| entry.key.borrow() == key) + .map(Arc::clone)?; + + self.index_if_ready(&entry) + } + } + + fn index_if_ready(&self, entry: &Arc>) -> Option + where + V: Clone, + { + let value = entry.get()?.clone(); + + { + let shard = self.lock_shard(entry.hash); + if shard + .find(entry.hash, |stored| Arc::ptr_eq(stored, entry)) + .is_some() + { + self.index(entry); + } + } + + Some(value) + } + + 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.lock_shard(hash); + + let occupied = shard + .find_entry(hash, |entry| entry.key.borrow() == key) + .ok()?; + let (entry, _) = occupied.remove(); + + self.remove_from_index(&entry); + Some(entry) + } + + pub fn cleanup_abandoned_entry(&self, entry: Arc>) { + let mut shard = self.lock_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. The pointer comparison rejects a detached or replaced entry. + 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); + } + + fn insert(&self, key: K, value: V) { + let hash = self.hasher.hash_one(&key); + let entry = Arc::new(Entry { + hash, + key, + cell: OnceCell::from_value(value), + was_indexed: AtomicBool::new(false), + }); + + let mut shard = self.lock_shard(hash); + if let Ok(occupied) = shard.find_entry(hash, |stored| stored.key.eq(&entry.key)) { + let (replaced, _) = occupied.remove(); + self.remove_from_index(&replaced); + } + shard.insert_unique(hash, Arc::clone(&entry), |entry| entry.hash); + + self.index(&entry); + } + + fn index(&self, entry: &Arc>) { + loop { + match self.index.entry_sync(ReadyEntry(Arc::clone(entry))) { + IndexEntry::Occupied(occupied) => { + if Arc::ptr_eq(&occupied.key().0, entry) { + break; + } + occupied.remove_entry(); + } + IndexEntry::Vacant(vacant) => { + vacant.insert_entry(()); + break; + } + } + } + + entry.was_indexed.store(true, Ordering::Relaxed); + } + + fn remove_from_index(&self, entry: &Arc>) { + struct EntryIdentity<'a, K, V>(&'a Arc>); + + impl Hash for EntryIdentity<'_, K, V> { + fn hash(&self, state: &mut H) { + state.write_u64(self.0.hash); + } + } + + impl Equivalent> for EntryIdentity<'_, K, V> { + fn equivalent(&self, entry: &ReadyEntry) -> bool { + Arc::ptr_eq(&entry.0, self.0) + } + } + + if entry.was_indexed.load(Ordering::Relaxed) { + self.index.remove_if_sync(&EntryIdentity(entry), |()| true); + } + } +} + +impl FromIterator<(K, V)> for Table +where + K: Eq + Hash, + S: BuildHasher + Default, +{ + fn from_iter>(iter: T) -> Self { + let iter = iter.into_iter(); + let table = Self::with_capacity_and_hasher(iter.size_hint().0, S::default()); + for (key, value) in iter { + table.insert(key, value); + } + + table + } +} diff --git a/asyncband/src/singleflight/mod.rs b/asyncband/src/singleflight/mod.rs index cbb3752..9b6c22b 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -23,9 +23,10 @@ use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; -use crate::internal::once_table::OnceTable; -use crate::internal::once_table::OnceTableEntry; +use table::Entry; +use table::Table; +mod table; #[cfg(test)] mod tests; @@ -33,7 +34,7 @@ mod tests; /// units of work can be executed with duplicate suppression. #[derive(Debug)] pub struct Group { - map: OnceTable, + map: Table, } // Holds one call's entry so Drop can clean it up if the work is abandoned. @@ -43,7 +44,7 @@ where S: BuildHasher, { group: &'a Group, - entry: Option>>, + entry: Option>>, } impl<'a, K, V, S> WorkCleanupGuard<'a, K, V, S> @@ -60,7 +61,7 @@ where } } - fn entry(&self) -> &Arc> { + fn entry(&self) -> &Arc> { self.entry.as_ref().unwrap() } @@ -102,7 +103,7 @@ where /// Creates a new Group with the default hasher. pub fn new() -> Self { Self { - map: OnceTable::with_hasher(RandomState::new()), + map: Table::with_hasher(RandomState::new()), } } } @@ -116,7 +117,7 @@ where /// Creates a new Group with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: OnceTable::with_hasher(hasher), + map: Table::with_hasher(hasher), } } @@ -181,7 +182,7 @@ where let result = entry .get_or_init(async || { let result = func().await; - self.map.remove_entry(entry); + self.map.remove_if_current(entry); result }) .await @@ -245,7 +246,7 @@ where let result = entry .get_or_try_init(async || { let result = func().await?; - self.map.remove_entry(entry); + self.map.remove_if_current(entry); Ok(result) }) .await? diff --git a/asyncband/src/singleflight/table.rs b/asyncband/src/singleflight/table.rs new file mode 100644 index 0000000..2ca0069 --- /dev/null +++ b/asyncband/src/singleflight/table.rs @@ -0,0 +1,161 @@ +// 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::borrow::Borrow; +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::default_shard_count; +use crate::internal::mutex::CachePaddedMutex; +use crate::internal::mutex::Mutex; +use crate::once::OnceCell; + +type Entries = HashTable>>; + +pub struct Entry { + hash: u64, + key: K, + cell: OnceCell, +} + +impl Entry { + fn initialized(&self) -> bool { + self.cell.initialized() + } + + pub async fn get_or_init(&self, init: F) -> &V + where + F: AsyncFnOnce() -> V, + { + self.cell.get_or_init(init).await + } + + pub async fn get_or_try_init(&self, init: F) -> Result<&V, E> + where + F: AsyncFnOnce() -> Result, + { + self.cell.get_or_try_init(init).await + } +} + +/// Storage for one in-flight call per key. +pub struct Table { + shards: Box<[CachePaddedMutex>]>, + hasher: S, +} + +impl fmt::Debug for Table +where + K: fmt::Debug, + V: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut debug_map = f.debug_map(); + for shard in &self.shards { + let entries = shard.0.lock(); + debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); + } + debug_map.finish() + } +} + +impl Table { + pub fn with_hasher(hasher: S) -> Self { + let shards = (0..default_shard_count()) + .map(|_| CachePaddedMutex(Mutex::new(HashTable::new()))) + .collect(); + Self { shards, hasher } + } + + fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { + self.shards[(hash as usize) & (self.shards.len() - 1)] + .0 + .lock() + } + + #[cfg(test)] + pub fn len(&self) -> usize { + self.shards.iter().map(|shard| shard.0.lock().len()).sum() + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + self.shards.iter().all(|shard| shard.0.lock().is_empty()) + } +} + +impl Table +where + K: Eq + Hash, + S: BuildHasher, +{ + pub fn get_or_insert(&self, key: K) -> Arc> { + let hash = self.hasher.hash_one(&key); + let mut shard = self.lock_shard(hash); + shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(Entry { + hash, + key, + cell: OnceCell::new(), + }) + }) + .into_mut() + .clone() + } + + pub fn remove(&self, key: &Q) + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let mut shard = self.lock_shard(hash); + if let Ok(occupied) = shard.find_entry(hash, |entry| entry.key.borrow() == key) { + drop(occupied.remove()); + } + } + + pub fn remove_if_current(&self, entry: &Arc>) { + let mut shard = self.lock_shard(entry.hash); + if let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) { + drop(occupied.remove()); + } + } + + pub fn cleanup_abandoned_entry(&self, entry: Arc>) { + let mut shard = self.lock_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. The pointer comparison rejects a detached or replaced entry. + 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); + } +} From 208e2d410b36d2f872d0f2af8b4c308bbf1e47c3 Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Fri, 28 Aug 2026 21:06:25 +0800 Subject: [PATCH 4/9] feat: add with_shard_amount and other related constructors --- asyncband/src/once/once_map/mod.rs | 53 ++++++++++++++++++++++++++++ asyncband/src/once/once_map/table.rs | 22 ++++++++++-- asyncband/src/singleflight/mod.rs | 22 ++++++++++++ asyncband/src/singleflight/table.rs | 11 +++++- 4 files changed, 104 insertions(+), 4 deletions(-) diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index 11c3848..7332693 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -112,6 +112,33 @@ where map: Table::with_capacity_and_hasher(capacity, RandomState::new()), } } + + /// Creates a new OnceMap with the default hasher and the specified shard amount. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_shard_amount(shard_amount: usize) -> Self { + Self { + map: Table::with_hasher_and_shard_amount(RandomState::new(), shard_amount), + } + } + + /// Creates a new OnceMap with the default hasher, the specified capacity, and the specified + /// shard amount. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_capacity_and_shard_amount(capacity: usize, shard_amount: usize) -> Self { + Self { + map: Table::with_capacity_and_hasher_and_shard_amount( + capacity, + RandomState::new(), + shard_amount, + ), + } + } } impl OnceMap @@ -134,6 +161,32 @@ where } } + /// Creates a new OnceMap with the given hasher and the specified shard amount. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { + Self { + map: Table::with_hasher_and_shard_amount(hasher, shard_amount), + } + } + + /// Creates a new OnceMap with the specified capacity, hasher, and shard amount. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_capacity_and_hasher_and_shard_amount( + capacity: usize, + hasher: S, + shard_amount: usize, + ) -> Self { + Self { + map: Table::with_capacity_and_hasher_and_shard_amount(capacity, hasher, shard_amount), + } + } + /// Compute the value for the given key if absent. /// /// If the value for the key is already being computed by another task, this task will wait for diff --git a/asyncband/src/once/once_map/table.rs b/asyncband/src/once/once_map/table.rs index ac7e634..8a61a8a 100644 --- a/asyncband/src/once/once_map/table.rs +++ b/asyncband/src/once/once_map/table.rs @@ -146,9 +146,25 @@ impl Table { } pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { - let shard_count = default_shard_count(); - let shard_capacity = capacity.div_ceil(shard_count); - let shards = (0..shard_count) + Self::with_capacity_and_hasher_and_shard_amount(capacity, hasher, default_shard_count()) + } + + pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { + Self::with_capacity_and_hasher_and_shard_amount(0, hasher, shard_amount) + } + + pub fn with_capacity_and_hasher_and_shard_amount( + capacity: usize, + hasher: S, + shard_amount: usize, + ) -> Self { + assert!( + shard_amount.is_power_of_two(), + "shard amount must be greater than zero and a power of two" + ); + + let shard_capacity = capacity.div_ceil(shard_amount); + let shards = (0..shard_amount) .map(|_| CachePaddedMutex(Mutex::new(HashTable::with_capacity(shard_capacity)))) .collect(); diff --git a/asyncband/src/singleflight/mod.rs b/asyncband/src/singleflight/mod.rs index 9b6c22b..1e602ec 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -106,6 +106,17 @@ where map: Table::with_hasher(RandomState::new()), } } + + /// Creates a new Group with the default hasher and the specified shard amount. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_shard_amount(shard_amount: usize) -> Self { + Self { + map: Table::with_hasher_and_shard_amount(RandomState::new(), shard_amount), + } + } } impl Group @@ -121,6 +132,17 @@ where } } + /// Creates a new Group with the given hasher and the specified shard amount. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { + Self { + map: Table::with_hasher_and_shard_amount(hasher, shard_amount), + } + } + /// Executes and returns the results of the given function, making sure that only one execution /// is in-flight for a given key at a time. /// diff --git a/asyncband/src/singleflight/table.rs b/asyncband/src/singleflight/table.rs index 2ca0069..637685b 100644 --- a/asyncband/src/singleflight/table.rs +++ b/asyncband/src/singleflight/table.rs @@ -80,7 +80,16 @@ where impl Table { pub fn with_hasher(hasher: S) -> Self { - let shards = (0..default_shard_count()) + Self::with_hasher_and_shard_amount(hasher, default_shard_count()) + } + + pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { + assert!( + shard_amount.is_power_of_two(), + "shard amount must be greater than zero and a power of two" + ); + + let shards = (0..shard_amount) .map(|_| CachePaddedMutex(Mutex::new(HashTable::new()))) .collect(); Self { shards, hasher } From 1bde87a4bffd60377161a0f769cdf184dc567582 Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Sat, 29 Aug 2026 01:16:23 +0800 Subject: [PATCH 5/9] refactor: use sharded RwLock hash table as ready index --- Cargo.lock | 26 ---- Cargo.toml | 1 - asyncband/Cargo.toml | 3 +- asyncband/src/internal/mod.rs | 3 + asyncband/src/internal/rwlock.rs | 41 ++++++ asyncband/src/once/once_map/table.rs | 121 ++---------------- .../src/once/once_map/table/ready_index.rs | 97 ++++++++++++++ 7 files changed, 155 insertions(+), 137 deletions(-) create mode 100644 asyncband/src/internal/rwlock.rs create mode 100644 asyncband/src/once/once_map/table/ready_index.rs diff --git a/Cargo.lock b/Cargo.lock index eddb010..0c320cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -69,7 +69,6 @@ name = "asyncband" version = "0.6.7" dependencies = [ "hashbrown", - "scc", "tokio", ] @@ -790,37 +789,12 @@ dependencies = [ "untrusted", ] -[[package]] -name = "saa" -version = "5.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f5acb362a0e75c2a963532fa7fabf13dff81626dc494df16488d30befcbea0" - -[[package]] -name = "scc" -version = "3.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8af0b99483d1c3e59471d4f0cb58b244169436a8979c889a91a3f697075ea01" -dependencies = [ - "saa", - "sdd", -] - [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "sdd" -version = "4.8.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f95d4cc3459db1e608946153c95cf20158af0b86131ae4ae451f90f54549e7d1" -dependencies = [ - "saa", -] - [[package]] name = "semver" version = "1.0.28" diff --git a/Cargo.toml b/Cargo.toml index 7fe5b5e..872eae1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,7 +33,6 @@ asyncband = { path = "asyncband" } # Optional runtime dependencies hashbrown = { version = "0.17.1", default-features = false } -scc = "3.8.6" # Dev dependencies async-channel = { version = "2.5.0" } diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index b2886c9..3871529 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -53,7 +53,7 @@ mpsc = [] mutex = [] once = ["semaphore"] once-cell = ["semaphore"] -once-map = ["dep:hashbrown", "dep:scc", "once-cell"] +once-map = ["dep:hashbrown", "once-cell"] oneshot = [] pool = ["semaphore"] rwlock = [] @@ -66,7 +66,6 @@ waitgroup = [] hashbrown = { workspace = true, default-features = false, features = [ "inline-more", ], optional = true } -scc = { workspace = true, optional = true } [dev-dependencies] tokio = { workspace = true, features = ["full"] } diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 07ae169..b759225 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -55,6 +55,9 @@ pub(crate) mod value_cell; ))] pub(crate) mod mutex; +#[cfg(feature = "once-map")] +pub(crate) mod rwlock; + #[cfg(any( feature = "mpsc", feature = "mutex", diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs new file mode 100644 index 0000000..906947f --- /dev/null +++ b/asyncband/src/internal/rwlock.rs @@ -0,0 +1,41 @@ +// 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) + } +} diff --git a/asyncband/src/once/once_map/table.rs b/asyncband/src/once/once_map/table.rs index 8a61a8a..f5c93f5 100644 --- a/asyncband/src/once/once_map/table.rs +++ b/asyncband/src/once/once_map/table.rs @@ -18,9 +18,7 @@ use std::borrow::Borrow; use std::fmt; use std::hash::BuildHasher; -use std::hash::BuildHasherDefault; use std::hash::Hash; -use std::hash::Hasher; use std::panic::UnwindSafe; use std::sync::Arc; use std::sync::MutexGuard; @@ -28,15 +26,16 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use hashbrown::HashTable; -use scc::Equivalent; -use scc::HashIndex; -use scc::hash_index::Entry as IndexEntry; use crate::internal::default_shard_count; use crate::internal::mutex::CachePaddedMutex; use crate::internal::mutex::Mutex; use crate::once::OnceCell; +mod ready_index; + +use ready_index::ReadyIndex; + type Entries = HashTable>>; pub struct Entry { @@ -75,54 +74,13 @@ impl Entry { } } -struct ReadyEntry(Arc>); - -impl PartialEq for ReadyEntry { - fn eq(&self, other: &Self) -> bool { - self.0.hash == other.0.hash && self.0.key == other.0.key - } -} - -impl Eq for ReadyEntry {} - -impl Hash for ReadyEntry { - fn hash(&self, state: &mut H) { - state.write_u64(self.0.hash); - } -} - -type BuildIdentityHasher = BuildHasherDefault; - -#[derive(Default)] -struct IdentityHasher(u64); - -impl Hasher for IdentityHasher { - fn finish(&self) -> u64 { - self.0 - } - - fn write(&mut self, bytes: &[u8]) { - for byte in bytes { - self.0 = self.0.rotate_left(8) ^ u64::from(*byte); - } - } - - fn write_u64(&mut self, value: u64) { - self.0 = value; - } -} - pub struct Table { shards: Box<[CachePaddedMutex>]>, - index: HashIndex, (), BuildIdentityHasher>, + index: ReadyIndex, hasher: S, } -/// `HashIndex` prevents `Table` from being automatically `UnwindSafe` unless `K` and `V` are -/// `UnwindSafe`. -/// Table operations are unwind-safe regardless, but since it was refactored from a -/// mutex-backed implementation, implement `UnwindSafe` manually to retain the same auto-trait -/// semantics. +/// Table operations recover poisoned ready-index locks before accessing their contents. impl UnwindSafe for Table {} impl fmt::Debug for Table @@ -170,7 +128,7 @@ impl Table { Self { shards, - index: HashIndex::with_capacity_and_hasher(capacity, BuildIdentityHasher::default()), + index: ReadyIndex::with_capacity_and_shard_amount(capacity, shard_amount), hasher, } } @@ -197,38 +155,13 @@ where K: Eq + Hash, S: BuildHasher, { - fn lockup_index(&self, hash: u64, key: &Q) -> Option + fn lookup_index(&self, hash: u64, key: &Q) -> Option where K: Borrow, Q: Eq + ?Sized, V: Clone, { - // I think we need to add a function-based peek_with to scc instead of using new type here. - // Alternatively, we could write our own HashIndex. - struct LookupKey<'a, Q: ?Sized> { - hash: u64, - key: &'a Q, - } - - impl Hash for LookupKey<'_, Q> { - fn hash(&self, state: &mut H) { - state.write_u64(self.hash); - } - } - - impl Equivalent> for LookupKey<'_, Q> - where - K: Borrow, - Q: Eq + ?Sized, - { - fn equivalent(&self, entry: &ReadyEntry) -> bool { - entry.0.key.borrow() == self.key - } - } - - self.index - .peek_with(&LookupKey { hash, key }, |entry, ()| entry.0.get().cloned()) - .flatten() + self.index.get(hash, key) } pub fn get_or_insert(&self, key: K) -> Lookup @@ -236,7 +169,7 @@ where V: Clone, { let hash = self.hasher.hash_one(&key); - if let Some(value) = self.lockup_index(hash, &key) { + if let Some(value) = self.lookup_index(hash, &key) { return Lookup::Ready(value); } @@ -270,7 +203,7 @@ where { let hash = self.hasher.hash_one(key); - if let Some(value) = self.lockup_index(hash, key) { + if let Some(value) = self.lookup_index(hash, key) { Some(value) } else { let entry = self @@ -354,41 +287,13 @@ where } fn index(&self, entry: &Arc>) { - loop { - match self.index.entry_sync(ReadyEntry(Arc::clone(entry))) { - IndexEntry::Occupied(occupied) => { - if Arc::ptr_eq(&occupied.key().0, entry) { - break; - } - occupied.remove_entry(); - } - IndexEntry::Vacant(vacant) => { - vacant.insert_entry(()); - break; - } - } - } - + self.index.insert(entry); entry.was_indexed.store(true, Ordering::Relaxed); } fn remove_from_index(&self, entry: &Arc>) { - struct EntryIdentity<'a, K, V>(&'a Arc>); - - impl Hash for EntryIdentity<'_, K, V> { - fn hash(&self, state: &mut H) { - state.write_u64(self.0.hash); - } - } - - impl Equivalent> for EntryIdentity<'_, K, V> { - fn equivalent(&self, entry: &ReadyEntry) -> bool { - Arc::ptr_eq(&entry.0, self.0) - } - } - if entry.was_indexed.load(Ordering::Relaxed) { - self.index.remove_if_sync(&EntryIdentity(entry), |()| true); + self.index.remove(entry); } } } diff --git a/asyncband/src/once/once_map/table/ready_index.rs b/asyncband/src/once/once_map/table/ready_index.rs new file mode 100644 index 0000000..ffe6925 --- /dev/null +++ b/asyncband/src/once/once_map/table/ready_index.rs @@ -0,0 +1,97 @@ +// 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::borrow::Borrow; +use std::sync::Arc; + +use hashbrown::HashTable; + +use super::Entry; +use crate::internal::rwlock::RwLock; + +const TARGET_ENTRIES_PER_BUCKET: usize = 2; + +#[repr(align(64))] +struct Bucket { + entries: RwLock>>>, +} + +impl Bucket { + fn new() -> Self { + Self { + entries: RwLock::new(HashTable::new()), + } + } +} + +pub struct ReadyIndex { + buckets: Box<[Bucket]>, +} + +impl ReadyIndex { + pub fn with_capacity_and_shard_amount(capacity: usize, shard_amount: usize) -> Self { + let bucket_count = capacity + .div_ceil(TARGET_ENTRIES_PER_BUCKET) + .max(shard_amount) + .next_power_of_two(); + + Self { + buckets: (0..bucket_count).map(|_| Bucket::new()).collect(), + } + } + + fn bucket(&self, hash: u64) -> &Bucket { + &self.buckets[(hash as usize) & (self.buckets.len() - 1)] + } + + pub fn get(&self, hash: u64, key: &Q) -> Option + where + K: Borrow, + Q: Eq + ?Sized, + V: Clone, + { + self.bucket(hash) + .entries + .read() + .find(hash, |entry| entry.key.borrow() == key) + .and_then(|entry| entry.get().cloned()) + } + + pub fn insert(&self, entry: &Arc>) + where + K: Eq, + { + let mut entries = self.bucket(entry.hash).entries.write(); + + if let Ok(occupied) = entries.find_entry(entry.hash, |stored| stored.key == entry.key) { + if Arc::ptr_eq(occupied.get(), entry) { + return; + } + drop(occupied.remove()); + } + + entries.insert_unique(entry.hash, Arc::clone(entry), |entry| entry.hash); + } + + pub fn remove(&self, entry: &Arc>) { + let mut entries = self.bucket(entry.hash).entries.write(); + + if let Ok(occupied) = entries.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) { + drop(occupied.remove()); + } + } +} From a53b14a107bbbe8e2eb279a7ea27e84e6a550d62 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 29 Aug 2026 14:56:22 +0800 Subject: [PATCH 6/9] fixup Signed-off-by: tison --- asyncband/src/mpsc/mod.rs | 20 ++++++++++---------- asyncband/src/rwlock/mod.rs | 17 +++++++++-------- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 87c7c8f..45bc740 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/mpsc/mod.rs @@ -21,13 +21,13 @@ mod bounded; mod error; mod unbounded; -pub use bounded::BoundedReceiver; -pub use bounded::BoundedSender; -pub use bounded::bounded; -pub use error::RecvError; -pub use error::SendError; -pub use error::TryRecvError; -pub use error::TrySendError; -pub use unbounded::UnboundedReceiver; -pub use unbounded::UnboundedSender; -pub use unbounded::unbounded; +pub use self::bounded::BoundedReceiver; +pub use self::bounded::BoundedSender; +pub use self::bounded::bounded; +pub use self::error::RecvError; +pub use self::error::SendError; +pub use self::error::TryRecvError; +pub use self::error::TrySendError; +pub use self::unbounded::UnboundedReceiver; +pub use self::unbounded::UnboundedSender; +pub use self::unbounded::unbounded; diff --git a/asyncband/src/rwlock/mod.rs b/asyncband/src/rwlock/mod.rs index 9f04847..36e97ca 100644 --- a/asyncband/src/rwlock/mod.rs +++ b/asyncband/src/rwlock/mod.rs @@ -76,21 +76,22 @@ use std::num::NonZeroUsize; use crate::internal::semaphore::Semaphore; mod mapped_read_guard; -pub use mapped_read_guard::MappedRwLockReadGuard; mod mapped_write_guard; -pub use mapped_write_guard::MappedRwLockWriteGuard; mod owned_mapped_read_guard; -pub use owned_mapped_read_guard::OwnedMappedRwLockReadGuard; mod owned_mapped_write_guard; -pub use owned_mapped_write_guard::OwnedMappedRwLockWriteGuard; mod owned_read_guard; -pub use owned_read_guard::OwnedRwLockReadGuard; mod owned_write_guard; -pub use owned_write_guard::OwnedRwLockWriteGuard; mod read_guard; -pub use read_guard::RwLockReadGuard; mod write_guard; -pub use write_guard::RwLockWriteGuard; + +pub use self::mapped_read_guard::MappedRwLockReadGuard; +pub use self::mapped_write_guard::MappedRwLockWriteGuard; +pub use self::owned_mapped_read_guard::OwnedMappedRwLockReadGuard; +pub use self::owned_mapped_write_guard::OwnedMappedRwLockWriteGuard; +pub use self::owned_read_guard::OwnedRwLockReadGuard; +pub use self::owned_write_guard::OwnedRwLockWriteGuard; +pub use self::read_guard::RwLockReadGuard; +pub use self::write_guard::RwLockWriteGuard; /// A reader-writer lock that allows multiple readers or a single writer at a time. /// From dab784bcce8b91e1335fe32a1517dabd447f54a8 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 29 Aug 2026 15:05:57 +0800 Subject: [PATCH 7/9] refactor: generalize cache padding wrapper Signed-off-by: tison --- asyncband/src/internal/cache_padded.rs | 48 +++++++++++++++++++ asyncband/src/internal/mod.rs | 3 ++ asyncband/src/internal/mutex.rs | 7 --- asyncband/src/once/once_map/table.rs | 17 ++++--- .../src/once/once_map/table/ready_index.rs | 23 +++------ asyncband/src/singleflight/table.rs | 17 ++++--- 6 files changed, 74 insertions(+), 41 deletions(-) create mode 100644 asyncband/src/internal/cache_padded.rs diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs new file mode 100644 index 0000000..f4400eb --- /dev/null +++ b/asyncband/src/internal/cache_padded.rs @@ -0,0 +1,48 @@ +// 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::ops::Deref; +use std::ops::DerefMut; + +/// Pads a value to reduce false sharing between adjacent values. +/// +/// On Zen 5, padding the 40-byte shard mutex to 64 bytes increased its size by 60% and improved +/// write performance by 25% at 32 threads. Other architectures still need to be measured. +#[repr(align(64))] +pub struct CachePadded { + value: T, +} + +impl CachePadded { + pub const fn new(value: T) -> Self { + Self { value } + } +} + +impl Deref for CachePadded { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.value + } +} + +impl DerefMut for CachePadded { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.value + } +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 34b2769..df5b504 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -40,6 +40,9 @@ pub(crate) mod arena; #[allow(dead_code)] pub(crate) mod countdown; +#[cfg(any(feature = "once-map", feature = "singleflight"))] +pub(crate) mod cache_padded; + #[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/mutex.rs b/asyncband/src/internal/mutex.rs index 73577fb..9477e26 100644 --- a/asyncband/src/internal/mutex.rs +++ b/asyncband/src/internal/mutex.rs @@ -36,13 +36,6 @@ impl Mutex { } } -#[cfg(any(feature = "once-map", feature = "singleflight"))] -/// Alignment uses 60% more memory (64/40) but improves write performance by 25% at 32 threads. (On -/// Zen 5 CPUs) -/// Need to test on other architectures. -#[repr(align(64))] -pub struct CachePaddedMutex(pub Mutex); - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/asyncband/src/once/once_map/table.rs b/asyncband/src/once/once_map/table.rs index f5c93f5..57f81bd 100644 --- a/asyncband/src/once/once_map/table.rs +++ b/asyncband/src/once/once_map/table.rs @@ -27,8 +27,8 @@ use std::sync::atomic::Ordering; use hashbrown::HashTable; +use crate::internal::cache_padded::CachePadded; use crate::internal::default_shard_count; -use crate::internal::mutex::CachePaddedMutex; use crate::internal::mutex::Mutex; use crate::once::OnceCell; @@ -37,6 +37,7 @@ mod ready_index; use ready_index::ReadyIndex; type Entries = HashTable>>; +type Shard = CachePadded>>; pub struct Entry { hash: u64, @@ -75,7 +76,7 @@ impl Entry { } pub struct Table { - shards: Box<[CachePaddedMutex>]>, + shards: Box<[Shard]>, index: ReadyIndex, hasher: S, } @@ -91,7 +92,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.0.lock(); + let entries = shard.lock(); debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); } debug_map.finish() @@ -123,7 +124,7 @@ impl Table { let shard_capacity = capacity.div_ceil(shard_amount); let shards = (0..shard_amount) - .map(|_| CachePaddedMutex(Mutex::new(HashTable::with_capacity(shard_capacity)))) + .map(|_| CachePadded::new(Mutex::new(HashTable::with_capacity(shard_capacity)))) .collect(); Self { @@ -134,19 +135,17 @@ impl Table { } fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { - self.shards[(hash as usize) & (self.shards.len() - 1)] - .0 - .lock() + self.shards[(hash as usize) & (self.shards.len() - 1)].lock() } #[cfg(test)] pub fn len(&self) -> usize { - self.shards.iter().map(|shard| shard.0.lock().len()).sum() + self.shards.iter().map(|shard| shard.lock().len()).sum() } #[cfg(test)] pub fn is_empty(&self) -> bool { - self.shards.iter().all(|shard| shard.0.lock().is_empty()) + self.shards.iter().all(|shard| shard.lock().is_empty()) } } diff --git a/asyncband/src/once/once_map/table/ready_index.rs b/asyncband/src/once/once_map/table/ready_index.rs index ffe6925..27bc94f 100644 --- a/asyncband/src/once/once_map/table/ready_index.rs +++ b/asyncband/src/once/once_map/table/ready_index.rs @@ -21,22 +21,12 @@ use std::sync::Arc; use hashbrown::HashTable; use super::Entry; +use crate::internal::cache_padded::CachePadded; use crate::internal::rwlock::RwLock; const TARGET_ENTRIES_PER_BUCKET: usize = 2; -#[repr(align(64))] -struct Bucket { - entries: RwLock>>>, -} - -impl Bucket { - fn new() -> Self { - Self { - entries: RwLock::new(HashTable::new()), - } - } -} +type Bucket = CachePadded>>>>; pub struct ReadyIndex { buckets: Box<[Bucket]>, @@ -50,7 +40,9 @@ impl ReadyIndex { .next_power_of_two(); Self { - buckets: (0..bucket_count).map(|_| Bucket::new()).collect(), + buckets: (0..bucket_count) + .map(|_| CachePadded::new(RwLock::new(HashTable::new()))) + .collect(), } } @@ -65,7 +57,6 @@ impl ReadyIndex { V: Clone, { self.bucket(hash) - .entries .read() .find(hash, |entry| entry.key.borrow() == key) .and_then(|entry| entry.get().cloned()) @@ -75,7 +66,7 @@ impl ReadyIndex { where K: Eq, { - let mut entries = self.bucket(entry.hash).entries.write(); + let mut entries = self.bucket(entry.hash).write(); if let Ok(occupied) = entries.find_entry(entry.hash, |stored| stored.key == entry.key) { if Arc::ptr_eq(occupied.get(), entry) { @@ -88,7 +79,7 @@ impl ReadyIndex { } pub fn remove(&self, entry: &Arc>) { - let mut entries = self.bucket(entry.hash).entries.write(); + let mut entries = self.bucket(entry.hash).write(); if let Ok(occupied) = entries.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) { drop(occupied.remove()); diff --git a/asyncband/src/singleflight/table.rs b/asyncband/src/singleflight/table.rs index 637685b..73699ef 100644 --- a/asyncband/src/singleflight/table.rs +++ b/asyncband/src/singleflight/table.rs @@ -24,12 +24,13 @@ use std::sync::MutexGuard; use hashbrown::HashTable; +use crate::internal::cache_padded::CachePadded; use crate::internal::default_shard_count; -use crate::internal::mutex::CachePaddedMutex; use crate::internal::mutex::Mutex; use crate::once::OnceCell; type Entries = HashTable>>; +type Shard = CachePadded>>; pub struct Entry { hash: u64, @@ -59,7 +60,7 @@ impl Entry { /// Storage for one in-flight call per key. pub struct Table { - shards: Box<[CachePaddedMutex>]>, + shards: Box<[Shard]>, hasher: S, } @@ -71,7 +72,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.0.lock(); + let entries = shard.lock(); debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); } debug_map.finish() @@ -90,25 +91,23 @@ impl Table { ); let shards = (0..shard_amount) - .map(|_| CachePaddedMutex(Mutex::new(HashTable::new()))) + .map(|_| CachePadded::new(Mutex::new(HashTable::new()))) .collect(); Self { shards, hasher } } fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { - self.shards[(hash as usize) & (self.shards.len() - 1)] - .0 - .lock() + self.shards[(hash as usize) & (self.shards.len() - 1)].lock() } #[cfg(test)] pub fn len(&self) -> usize { - self.shards.iter().map(|shard| shard.0.lock().len()).sum() + self.shards.iter().map(|shard| shard.lock().len()).sum() } #[cfg(test)] pub fn is_empty(&self) -> bool { - self.shards.iter().all(|shard| shard.0.lock().is_empty()) + self.shards.iter().all(|shard| shard.lock().is_empty()) } } From bee5efb84d6a8b669d835abe5e586c92ce2a2b0d Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 29 Aug 2026 15:11:39 +0800 Subject: [PATCH 8/9] fixup Signed-off-by: tison --- asyncband/src/internal/mutex.rs | 20 -------------------- asyncband/src/internal/waitset.rs | 7 ++----- 2 files changed, 2 insertions(+), 25 deletions(-) diff --git a/asyncband/src/internal/mutex.rs b/asyncband/src/internal/mutex.rs index 9477e26..ce5229f 100644 --- a/asyncband/src/internal/mutex.rs +++ b/asyncband/src/internal/mutex.rs @@ -35,23 +35,3 @@ impl Mutex { self.0.lock().unwrap_or_else(PoisonError::into_inner) } } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use crate::internal::mutex::Mutex; - - #[test] - fn test_poison_mutex() { - let mutex = Arc::new(Mutex::new(42)); - let m = mutex.clone(); - let handle = std::thread::spawn(move || { - let _guard = m.lock(); - panic!("poison"); - }); - let _ = handle.join(); - let guard = mutex.lock(); - assert_eq!(*guard, 42); - } -} diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index ddf8bf2..8f4ab1f 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -158,10 +158,7 @@ mod tests { #[test] fn waker_token_preserves_the_option_niche() { - assert_eq!( - std::mem::size_of::(), - std::mem::size_of::>() - ); + assert_eq!(size_of::(), size_of::>()); } struct TrackWake(AtomicUsize); @@ -252,7 +249,7 @@ mod tests { register(&mut waiters, &mut second, &panicking); register(&mut waiters, &mut third, &tracked); - let result = std::panic::catch_unwind(AssertUnwindSafe(|| { + let result = panic::catch_unwind(AssertUnwindSafe(|| { wake_all(waiters.take_wakers()); })); assert!(result.is_err()); From 0be37b0c4af45602bb443107d12240a5672a01ff Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 29 Aug 2026 15:33:21 +0800 Subject: [PATCH 9/9] refactor: consolidate keyed work storage Signed-off-by: tison --- asyncband/src/once/once_map/mod.rs | 416 +++++++++++++++--- asyncband/src/once/once_map/table.rs | 314 ------------- .../src/once/once_map/table/ready_index.rs | 88 ---- asyncband/src/once/once_map/tests.rs | 19 +- asyncband/src/singleflight/mod.rs | 171 ++++++- asyncband/src/singleflight/table.rs | 169 ------- asyncband/src/singleflight/tests.rs | 12 +- tests-integration/tests/once_map_test.rs | 19 + 8 files changed, 546 insertions(+), 662 deletions(-) delete mode 100644 asyncband/src/once/once_map/table.rs delete mode 100644 asyncband/src/once/once_map/table/ready_index.rs delete mode 100644 asyncband/src/singleflight/table.rs diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index 7332693..bc361e3 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -16,26 +16,359 @@ // under the License. use std::borrow::Borrow; +use std::fmt; use std::hash::BuildHasher; use std::hash::Hash; use std::hash::RandomState; +use std::panic::UnwindSafe; use std::sync::Arc; +use std::sync::MutexGuard; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; -use table::Entry; -use table::Lookup; -use table::Table; +use hashbrown::HashTable; + +use crate::internal::cache_padded::CachePadded; +use crate::internal::default_shard_count; +use crate::internal::mutex::Mutex; +use crate::internal::rwlock::RwLock; +use crate::once::OnceCell; -mod table; #[cfg(test)] mod tests; +const TARGET_ENTRIES_PER_BUCKET: usize = 2; +// Capacity may increase read parallelism, but it must not turn a reservation into an unbounded +// number of eagerly allocated locks. Four preserves the benchmarked 1,024-entry layout on hosts +// whose default rounds to 128 shards. +const MAX_READY_BUCKETS_PER_SHARD: usize = 4; + +type Entries = HashTable>>; +type Shard = CachePadded>>; +type Bucket = CachePadded>>>>; + +struct Entry { + hash: u64, + key: K, + cell: OnceCell, + // Accessed only while the corresponding primary shard is locked. Atomic interior mutability + // keeps `Entry` shareable with the read index; cross-lock ordering is not required. + was_indexed: AtomicBool, +} + +enum Lookup { + Ready(V), + Pending(Arc>), +} + +struct ReadyIndex { + buckets: Box<[Bucket]>, +} + +impl ReadyIndex { + fn new(capacity: usize, shard_amount: usize) -> Self { + let target_bucket_count = capacity.div_ceil(TARGET_ENTRIES_PER_BUCKET); + let max_bucket_count = shard_amount + .checked_mul(MAX_READY_BUCKETS_PER_SHARD) + .unwrap_or(shard_amount); + let bucket_count = target_bucket_count + .clamp(shard_amount, max_bucket_count) + .next_power_of_two(); + + Self { + buckets: (0..bucket_count) + .map(|_| CachePadded::new(RwLock::new(HashTable::new()))) + .collect(), + } + } + + fn bucket(&self, hash: u64) -> &Bucket { + &self.buckets[(hash as usize) & (self.buckets.len() - 1)] + } + + fn get(&self, hash: u64, key: &Q) -> Option + where + K: Borrow, + Q: Eq + ?Sized, + V: Clone, + { + self.bucket(hash) + .read() + .find(hash, |entry| entry.key.borrow() == key) + .and_then(|entry| entry.cell.get().cloned()) + } + + fn insert(&self, entry: &Arc>) + where + K: Eq, + { + let mut entries = self.bucket(entry.hash).write(); + let replaced = if let Ok(occupied) = + entries.find_entry(entry.hash, |stored| stored.key == entry.key) + { + if Arc::ptr_eq(occupied.get(), entry) { + return; + } + Some(occupied.remove().0) + } else { + None + }; + + entries.insert_unique(entry.hash, Arc::clone(entry), |entry| entry.hash); + drop(entries); + drop(replaced); + } + + fn remove(&self, entry: &Arc>) { + let mut entries = self.bucket(entry.hash).write(); + + let Ok(occupied) = entries.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) + else { + return; + }; + let (removed, _) = occupied.remove(); + drop(entries); + drop(removed); + } +} + /// A hash map that runs computation only once for each key and stores the result. /// /// Note that this always clones the value out of the underlying map. Because of this, it's common /// to wrap the `V` in an `Arc` to make cloning cheap. -#[derive(Debug)] pub struct OnceMap { - map: Table, + // Mutations always lock a primary shard before the corresponding ready bucket. + shards: Box<[Shard]>, + // Initialized entries are mirrored here so hit-only reads do not contend on primary shards. + ready: ReadyIndex, + hasher: S, +} + +/// Operations recover poisoned shard locks before accessing their contents. +impl UnwindSafe for OnceMap {} + +impl fmt::Debug for OnceMap +where + K: fmt::Debug, + V: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Write::write_str(f, "OnceMap ")?; + 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 OnceMap { + fn with_config(capacity: usize, hasher: S, shard_amount: usize) -> Self { + assert!( + shard_amount.is_power_of_two(), + "shard amount must be greater than zero and a power of two" + ); + + let shard_capacity = capacity / shard_amount; + let extra_capacity = capacity % shard_amount; + let shards = (0..shard_amount) + .map(|shard_index| { + let capacity = shard_capacity + usize::from(shard_index < extra_capacity); + CachePadded::new(Mutex::new(HashTable::with_capacity(capacity))) + }) + .collect(); + + Self { + shards, + ready: ReadyIndex::new(capacity, shard_amount), + hasher, + } + } + + fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { + self.shards[(hash as usize) & (self.shards.len() - 1)].lock() + } + + #[cfg(test)] + fn len(&self) -> usize { + self.shards.iter().map(|shard| shard.lock().len()).sum() + } + + #[cfg(test)] + fn is_empty(&self) -> bool { + self.shards.iter().all(|shard| shard.lock().is_empty()) + } +} + +impl OnceMap +where + K: Eq + Hash, + S: BuildHasher, +{ + fn get_or_insert(&self, key: K) -> Lookup + where + V: Clone, + { + let hash = self.hasher.hash_one(&key); + if let Some(value) = self.ready.get(hash, &key) { + return Lookup::Ready(value); + } + + let entry = { + let mut shard = self.lock_shard(hash); + shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(Entry { + hash, + key, + cell: OnceCell::new(), + was_indexed: AtomicBool::new(false), + }) + }) + .into_mut() + .clone() + }; + + match self.index_if_ready(&entry) { + Some(value) => Lookup::Ready(value), + None => Lookup::Pending(entry), + } + } + + fn get_value(&self, key: &Q) -> Option + where + K: Borrow, + Q: Eq + Hash + ?Sized, + V: Clone, + { + let hash = self.hasher.hash_one(key); + + if let Some(value) = self.ready.get(hash, key) { + Some(value) + } else { + let entry = self + .lock_shard(hash) + .find(hash, |entry| entry.key.borrow() == key) + .map(Arc::clone)?; + + self.index_if_ready(&entry) + } + } + + fn index_if_ready(&self, entry: &Arc>) -> Option + where + V: Clone, + { + let value = entry.cell.get()?.clone(); + + { + let shard = self.lock_shard(entry.hash); + if shard + .find(entry.hash, |stored| Arc::ptr_eq(stored, entry)) + .is_some() + { + self.index(entry); + } + } + + Some(value) + } + + fn remove_entry(&self, key: &Q) -> Option>> + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let mut shard = self.lock_shard(hash); + + let occupied = shard + .find_entry(hash, |entry| entry.key.borrow() == key) + .ok()?; + let (entry, _) = occupied.remove(); + + self.remove_from_index(&entry); + Some(entry) + } + + fn cleanup_abandoned_entry(&self, entry: Arc>) { + let mut shard = self.lock_shard(entry.hash); + let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, &entry)) + else { + // A concurrent remove detached the entry. It may be the final owner, so release it + // after unlocking rather than running user destructors under the shard lock. + drop(shard); + drop(entry); + return; + }; + + // With map ownership confirmed and the shard locked against new callers, two owners means + // the map and this cleanup guard are the only remaining references. + if Arc::strong_count(&entry) == 2 && !entry.cell.initialized() { + let (stored, _) = occupied.remove(); + drop(shard); + drop(entry); + drop(stored); + } else { + // A waiting cleanup must observe this call's reference being released while no new + // caller can clone the map's reference. + drop(entry); + } + } + + fn insert(&self, key: K, value: V) { + let hash = self.hasher.hash_one(&key); + let entry = Arc::new(Entry { + hash, + key, + cell: OnceCell::from_value(value), + was_indexed: AtomicBool::new(false), + }); + + let mut shard = self.lock_shard(hash); + let replaced = shard + .find_entry(hash, |stored| stored.key.eq(&entry.key)) + .ok() + .map(|occupied| occupied.remove().0); + if let Some(replaced) = &replaced { + self.remove_from_index(replaced); + } + shard.insert_unique(hash, Arc::clone(&entry), |entry| entry.hash); + + self.index(&entry); + drop(shard); + drop(replaced); + } + + fn index(&self, entry: &Arc>) { + self.ready.insert(entry); + entry.was_indexed.store(true, Ordering::Relaxed); + } + + fn remove_from_index(&self, entry: &Arc>) { + if entry.was_indexed.load(Ordering::Relaxed) { + self.ready.remove(entry); + } + } +} + +impl FromIterator<(K, V)> for OnceMap +where + K: Eq + Hash, + V: Clone, + S: BuildHasher + Default, +{ + fn from_iter>(iter: T) -> Self { + let iter = iter.into_iter(); + let map = Self::with_config(iter.size_hint().0, S::default(), default_shard_count()); + for (key, value) in iter { + map.insert(key, value); + } + + map + } } // Holds one call's entry so Drop can clean it up if the computation is abandoned. @@ -79,7 +412,7 @@ where return; }; - self.once_map.map.cleanup_abandoned_entry(entry); + self.once_map.cleanup_abandoned_entry(entry); } } @@ -101,16 +434,12 @@ where { /// Creates a new OnceMap with the default hasher. pub fn new() -> Self { - Self { - map: Table::with_hasher(RandomState::new()), - } + Self::with_config(0, RandomState::new(), default_shard_count()) } /// Creates a new OnceMap with the default hasher and the specified capacity. pub fn with_capacity(capacity: usize) -> Self { - Self { - map: Table::with_capacity_and_hasher(capacity, RandomState::new()), - } + Self::with_config(capacity, RandomState::new(), default_shard_count()) } /// Creates a new OnceMap with the default hasher and the specified shard amount. @@ -119,9 +448,7 @@ where /// /// Panics if `shard_amount` is zero or is not a power of two. pub fn with_shard_amount(shard_amount: usize) -> Self { - Self { - map: Table::with_hasher_and_shard_amount(RandomState::new(), shard_amount), - } + Self::with_config(0, RandomState::new(), shard_amount) } /// Creates a new OnceMap with the default hasher, the specified capacity, and the specified @@ -131,13 +458,7 @@ where /// /// Panics if `shard_amount` is zero or is not a power of two. pub fn with_capacity_and_shard_amount(capacity: usize, shard_amount: usize) -> Self { - Self { - map: Table::with_capacity_and_hasher_and_shard_amount( - capacity, - RandomState::new(), - shard_amount, - ), - } + Self::with_config(capacity, RandomState::new(), shard_amount) } } @@ -149,16 +470,12 @@ where { /// Creates a new OnceMap with the given hasher. pub fn with_hasher(hasher: S) -> Self { - Self { - map: Table::with_hasher(hasher), - } + Self::with_config(0, hasher, default_shard_count()) } /// Create a OnceMap with the specified capacity and hasher. pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { - Self { - map: Table::with_capacity_and_hasher(capacity, hasher), - } + Self::with_config(capacity, hasher, default_shard_count()) } /// Creates a new OnceMap with the given hasher and the specified shard amount. @@ -167,9 +484,7 @@ where /// /// Panics if `shard_amount` is zero or is not a power of two. pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { - Self { - map: Table::with_hasher_and_shard_amount(hasher, shard_amount), - } + Self::with_config(0, hasher, shard_amount) } /// Creates a new OnceMap with the specified capacity, hasher, and shard amount. @@ -182,9 +497,7 @@ where hasher: S, shard_amount: usize, ) -> Self { - Self { - map: Table::with_capacity_and_hasher_and_shard_amount(capacity, hasher, shard_amount), - } + Self::with_config(capacity, hasher, shard_amount) } /// Compute the value for the given key if absent. @@ -198,13 +511,13 @@ where where F: AsyncFnOnce() -> V, { - let entry = match self.map.get_or_insert(key) { + let entry = match self.get_or_insert(key) { Lookup::Ready(value) => return value, - Lookup::Entry(entry) => entry, + Lookup::Pending(entry) => entry, }; let guard = ComputeCleanupGuard::new(self, entry); - let result = guard.entry().get_or_init(func).await.clone(); + let result = guard.entry().cell.get_or_init(func).await.clone(); guard.dismiss(); result } @@ -220,13 +533,13 @@ where where F: AsyncFnOnce() -> Result, { - let entry = match self.map.get_or_insert(key) { + let entry = match self.get_or_insert(key) { Lookup::Ready(value) => return Ok(value), - Lookup::Entry(entry) => entry, + Lookup::Pending(entry) => entry, }; let guard = ComputeCleanupGuard::new(self, entry); - let result = guard.entry().get_or_try_init(func).await?.clone(); + let result = guard.entry().cell.get_or_try_init(func).await?.clone(); guard.dismiss(); Ok(result) } @@ -237,7 +550,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - self.map.get(key) + self.get_value(key) } /// Remove the given key from the map. @@ -250,7 +563,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - self.map.remove(key); + self.remove_entry(key); } /// Remove the given key from the map and return a *clone* of the value if exists. @@ -264,20 +577,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.remove(key)?; - entry.get().cloned() - } -} - -impl FromIterator<(K, V)> for OnceMap -where - K: Eq + Hash, - V: Clone, - S: Default + BuildHasher, -{ - fn from_iter>(iter: T) -> Self { - Self { - map: iter.into_iter().collect(), - } + let entry = self.remove_entry(key)?; + entry.cell.get().cloned() } } diff --git a/asyncband/src/once/once_map/table.rs b/asyncband/src/once/once_map/table.rs deleted file mode 100644 index 57f81bd..0000000 --- a/asyncband/src/once/once_map/table.rs +++ /dev/null @@ -1,314 +0,0 @@ -// 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::borrow::Borrow; -use std::fmt; -use std::hash::BuildHasher; -use std::hash::Hash; -use std::panic::UnwindSafe; -use std::sync::Arc; -use std::sync::MutexGuard; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; - -use hashbrown::HashTable; - -use crate::internal::cache_padded::CachePadded; -use crate::internal::default_shard_count; -use crate::internal::mutex::Mutex; -use crate::once::OnceCell; - -mod ready_index; - -use ready_index::ReadyIndex; - -type Entries = HashTable>>; -type Shard = CachePadded>>; - -pub struct Entry { - hash: u64, - key: K, - cell: OnceCell, - was_indexed: AtomicBool, -} - -pub enum Lookup { - Ready(V), - Entry(Arc>), -} - -impl Entry { - pub fn initialized(&self) -> bool { - self.cell.initialized() - } - - pub fn get(&self) -> Option<&V> { - self.cell.get() - } - - pub async fn get_or_init(&self, init: F) -> &V - where - F: AsyncFnOnce() -> V, - { - self.cell.get_or_init(init).await - } - - pub async fn get_or_try_init(&self, init: F) -> Result<&V, E> - where - F: AsyncFnOnce() -> Result, - { - self.cell.get_or_try_init(init).await - } -} - -pub struct Table { - shards: Box<[Shard]>, - index: ReadyIndex, - hasher: S, -} - -/// Table operations recover poisoned ready-index locks before accessing their contents. -impl UnwindSafe for Table {} - -impl fmt::Debug for Table -where - K: fmt::Debug, - V: fmt::Debug, -{ - 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(); - debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); - } - debug_map.finish() - } -} - -impl Table { - pub fn with_hasher(hasher: S) -> Self { - Self::with_capacity_and_hasher(0, hasher) - } - - pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { - Self::with_capacity_and_hasher_and_shard_amount(capacity, hasher, default_shard_count()) - } - - pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { - Self::with_capacity_and_hasher_and_shard_amount(0, hasher, shard_amount) - } - - pub fn with_capacity_and_hasher_and_shard_amount( - capacity: usize, - hasher: S, - shard_amount: usize, - ) -> Self { - assert!( - shard_amount.is_power_of_two(), - "shard amount must be greater than zero and a power of two" - ); - - let shard_capacity = capacity.div_ceil(shard_amount); - let shards = (0..shard_amount) - .map(|_| CachePadded::new(Mutex::new(HashTable::with_capacity(shard_capacity)))) - .collect(); - - Self { - shards, - index: ReadyIndex::with_capacity_and_shard_amount(capacity, shard_amount), - hasher, - } - } - - fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { - self.shards[(hash as usize) & (self.shards.len() - 1)].lock() - } - - #[cfg(test)] - pub fn len(&self) -> usize { - self.shards.iter().map(|shard| shard.lock().len()).sum() - } - - #[cfg(test)] - pub fn is_empty(&self) -> bool { - self.shards.iter().all(|shard| shard.lock().is_empty()) - } -} - -impl Table -where - K: Eq + Hash, - S: BuildHasher, -{ - fn lookup_index(&self, hash: u64, key: &Q) -> Option - where - K: Borrow, - Q: Eq + ?Sized, - V: Clone, - { - self.index.get(hash, key) - } - - pub fn get_or_insert(&self, key: K) -> Lookup - where - V: Clone, - { - let hash = self.hasher.hash_one(&key); - if let Some(value) = self.lookup_index(hash, &key) { - return Lookup::Ready(value); - } - - let entry = { - let mut shard = self.lock_shard(hash); - shard - .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) - .or_insert_with(|| { - Arc::new(Entry { - hash, - key, - cell: OnceCell::new(), - was_indexed: AtomicBool::new(false), - }) - }) - .into_mut() - .clone() - }; - - match self.index_if_ready(&entry) { - Some(value) => Lookup::Ready(value), - None => Lookup::Entry(entry), - } - } - - pub fn get(&self, key: &Q) -> Option - where - K: Borrow, - Q: Eq + Hash + ?Sized, - V: Clone, - { - let hash = self.hasher.hash_one(key); - - if let Some(value) = self.lookup_index(hash, key) { - Some(value) - } else { - let entry = self - .lock_shard(hash) - .find(hash, |entry| entry.key.borrow() == key) - .map(Arc::clone)?; - - self.index_if_ready(&entry) - } - } - - fn index_if_ready(&self, entry: &Arc>) -> Option - where - V: Clone, - { - let value = entry.get()?.clone(); - - { - let shard = self.lock_shard(entry.hash); - if shard - .find(entry.hash, |stored| Arc::ptr_eq(stored, entry)) - .is_some() - { - self.index(entry); - } - } - - Some(value) - } - - 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.lock_shard(hash); - - let occupied = shard - .find_entry(hash, |entry| entry.key.borrow() == key) - .ok()?; - let (entry, _) = occupied.remove(); - - self.remove_from_index(&entry); - Some(entry) - } - - pub fn cleanup_abandoned_entry(&self, entry: Arc>) { - let mut shard = self.lock_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. The pointer comparison rejects a detached or replaced entry. - 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); - } - - fn insert(&self, key: K, value: V) { - let hash = self.hasher.hash_one(&key); - let entry = Arc::new(Entry { - hash, - key, - cell: OnceCell::from_value(value), - was_indexed: AtomicBool::new(false), - }); - - let mut shard = self.lock_shard(hash); - if let Ok(occupied) = shard.find_entry(hash, |stored| stored.key.eq(&entry.key)) { - let (replaced, _) = occupied.remove(); - self.remove_from_index(&replaced); - } - shard.insert_unique(hash, Arc::clone(&entry), |entry| entry.hash); - - self.index(&entry); - } - - fn index(&self, entry: &Arc>) { - self.index.insert(entry); - entry.was_indexed.store(true, Ordering::Relaxed); - } - - fn remove_from_index(&self, entry: &Arc>) { - if entry.was_indexed.load(Ordering::Relaxed) { - self.index.remove(entry); - } - } -} - -impl FromIterator<(K, V)> for Table -where - K: Eq + Hash, - S: BuildHasher + Default, -{ - fn from_iter>(iter: T) -> Self { - let iter = iter.into_iter(); - let table = Self::with_capacity_and_hasher(iter.size_hint().0, S::default()); - for (key, value) in iter { - table.insert(key, value); - } - - table - } -} diff --git a/asyncband/src/once/once_map/table/ready_index.rs b/asyncband/src/once/once_map/table/ready_index.rs deleted file mode 100644 index 27bc94f..0000000 --- a/asyncband/src/once/once_map/table/ready_index.rs +++ /dev/null @@ -1,88 +0,0 @@ -// 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::borrow::Borrow; -use std::sync::Arc; - -use hashbrown::HashTable; - -use super::Entry; -use crate::internal::cache_padded::CachePadded; -use crate::internal::rwlock::RwLock; - -const TARGET_ENTRIES_PER_BUCKET: usize = 2; - -type Bucket = CachePadded>>>>; - -pub struct ReadyIndex { - buckets: Box<[Bucket]>, -} - -impl ReadyIndex { - pub fn with_capacity_and_shard_amount(capacity: usize, shard_amount: usize) -> Self { - let bucket_count = capacity - .div_ceil(TARGET_ENTRIES_PER_BUCKET) - .max(shard_amount) - .next_power_of_two(); - - Self { - buckets: (0..bucket_count) - .map(|_| CachePadded::new(RwLock::new(HashTable::new()))) - .collect(), - } - } - - fn bucket(&self, hash: u64) -> &Bucket { - &self.buckets[(hash as usize) & (self.buckets.len() - 1)] - } - - pub fn get(&self, hash: u64, key: &Q) -> Option - where - K: Borrow, - Q: Eq + ?Sized, - V: Clone, - { - self.bucket(hash) - .read() - .find(hash, |entry| entry.key.borrow() == key) - .and_then(|entry| entry.get().cloned()) - } - - pub fn insert(&self, entry: &Arc>) - where - K: Eq, - { - let mut entries = self.bucket(entry.hash).write(); - - if let Ok(occupied) = entries.find_entry(entry.hash, |stored| stored.key == entry.key) { - if Arc::ptr_eq(occupied.get(), entry) { - return; - } - drop(occupied.remove()); - } - - entries.insert_unique(entry.hash, Arc::clone(entry), |entry| entry.hash); - } - - pub fn remove(&self, entry: &Arc>) { - let mut entries = self.bucket(entry.hash).write(); - - if let Ok(occupied) = entries.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) { - drop(occupied.remove()); - } - } -} diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index b3cb88d..52ee05d 100644 --- a/asyncband/src/once/once_map/tests.rs +++ b/asyncband/src/once/once_map/tests.rs @@ -18,10 +18,19 @@ use std::sync::Arc; use super::OnceMap; +use super::ReadyIndex; use crate::test_support::poll_once; // These tests stay next to the implementation because they inspect private state. +#[test] +fn ready_index_bucket_count_is_bounded_by_shards() { + let shard_amount = 8; + let index = ReadyIndex::::new(1_000_000, shard_amount); + + assert_eq!(index.buckets.len(), shard_amount * 4); +} + #[tokio::test] async fn failed_compute_removes_empty_entry() { let map = OnceMap::new(); @@ -29,7 +38,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.is_empty()); + assert!(map.is_empty()); } #[tokio::test] @@ -46,7 +55,7 @@ async fn panicked_compute_removes_empty_entry() { }); assert!(task.await.unwrap_err().is_panic()); - assert!(map.map.is_empty()); + assert!(map.is_empty()); } #[tokio::test] @@ -65,11 +74,11 @@ async fn cancelled_compute_removes_empty_entry() { }); started_rx.await.unwrap(); - assert_eq!(map.map.len(), 1); + assert_eq!(map.len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(map.map.is_empty()); + assert!(map.is_empty()); } #[tokio::test] @@ -91,7 +100,7 @@ async fn failed_compute_preserves_entry_for_waiter_retry() { release_tx.send(()).unwrap(); assert_eq!(first.await, Err("fail")); - assert_eq!(map.map.len(), 1); + assert_eq!(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 1e602ec..88a1e79 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -18,23 +18,156 @@ //! Singleflight provides a duplicate function call suppression mechanism. use std::borrow::Borrow; +use std::fmt; use std::hash::BuildHasher; use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; +use std::sync::MutexGuard; -use table::Entry; -use table::Table; +use hashbrown::HashTable; + +use crate::internal::cache_padded::CachePadded; +use crate::internal::default_shard_count; +use crate::internal::mutex::Mutex; +use crate::once::OnceCell; -mod table; #[cfg(test)] mod tests; +type Entries = HashTable>>; +type Shard = CachePadded>>; + +struct Entry { + hash: u64, + key: K, + cell: OnceCell, +} + /// Group represents a class of work and forms a namespace in which /// units of work can be executed with duplicate suppression. -#[derive(Debug)] pub struct Group { - map: Table, + shards: Box<[Shard]>, + hasher: S, +} + +impl fmt::Debug for Group +where + K: fmt::Debug, + V: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Write::write_str(f, "Group ")?; + 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 Group { + fn with_config(hasher: S, shard_amount: usize) -> Self { + assert!( + shard_amount.is_power_of_two(), + "shard amount must be greater than zero and a power of two" + ); + + let shards = (0..shard_amount) + .map(|_| CachePadded::new(Mutex::new(HashTable::new()))) + .collect(); + Self { shards, hasher } + } + + fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { + self.shards[(hash as usize) & (self.shards.len() - 1)].lock() + } + + #[cfg(test)] + fn len(&self) -> usize { + self.shards.iter().map(|shard| shard.lock().len()).sum() + } + + #[cfg(test)] + fn is_empty(&self) -> bool { + self.shards.iter().all(|shard| shard.lock().is_empty()) + } +} + +impl Group +where + K: Eq + Hash, + S: BuildHasher, +{ + fn get_or_insert(&self, key: K) -> Arc> { + let hash = self.hasher.hash_one(&key); + let mut shard = self.lock_shard(hash); + shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(Entry { + hash, + key, + cell: OnceCell::new(), + }) + }) + .into_mut() + .clone() + } + + fn remove(&self, key: &Q) + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let removed = { + let mut shard = self.lock_shard(hash); + let Ok(occupied) = shard.find_entry(hash, |entry| entry.key.borrow() == key) else { + return; + }; + occupied.remove().0 + }; + drop(removed); + } + + fn remove_if_current(&self, entry: &Arc>) { + let removed = { + let mut shard = self.lock_shard(entry.hash); + let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) + else { + return; + }; + occupied.remove().0 + }; + drop(removed); + } + + fn cleanup_abandoned_entry(&self, entry: Arc>) { + let mut shard = self.lock_shard(entry.hash); + let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, &entry)) + else { + // `forget` detached the entry. It may be the final owner, so release it after + // unlocking rather than running user destructors under the shard lock. + drop(shard); + drop(entry); + return; + }; + + // With group ownership confirmed and the shard locked against new callers, two owners + // means the group and this cleanup guard are the only remaining references. + if Arc::strong_count(&entry) == 2 && !entry.cell.initialized() { + let (stored, _) = occupied.remove(); + drop(shard); + drop(entry); + drop(stored); + } else { + // A waiting cleanup must observe this call's reference being released while no new + // caller can clone the group's reference. + drop(entry); + } + } } // Holds one call's entry so Drop can clean it up if the work is abandoned. @@ -53,7 +186,7 @@ where S: BuildHasher, { fn new(group: &'a Group, key: K) -> Self { - let entry = group.map.get_or_insert(key); + let entry = group.get_or_insert(key); Self { group, @@ -80,7 +213,7 @@ where return; }; - self.group.map.cleanup_abandoned_entry(entry); + self.group.cleanup_abandoned_entry(entry); } } @@ -102,9 +235,7 @@ where { /// Creates a new Group with the default hasher. pub fn new() -> Self { - Self { - map: Table::with_hasher(RandomState::new()), - } + Self::with_config(RandomState::new(), default_shard_count()) } /// Creates a new Group with the default hasher and the specified shard amount. @@ -113,9 +244,7 @@ where /// /// Panics if `shard_amount` is zero or is not a power of two. pub fn with_shard_amount(shard_amount: usize) -> Self { - Self { - map: Table::with_hasher_and_shard_amount(RandomState::new(), shard_amount), - } + Self::with_config(RandomState::new(), shard_amount) } } @@ -127,9 +256,7 @@ where { /// Creates a new Group with the given hasher. pub fn with_hasher(hasher: S) -> Self { - Self { - map: Table::with_hasher(hasher), - } + Self::with_config(hasher, default_shard_count()) } /// Creates a new Group with the given hasher and the specified shard amount. @@ -138,9 +265,7 @@ where /// /// Panics if `shard_amount` is zero or is not a power of two. pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { - Self { - map: Table::with_hasher_and_shard_amount(hasher, shard_amount), - } + Self::with_config(hasher, shard_amount) } /// Executes and returns the results of the given function, making sure that only one execution @@ -202,9 +327,10 @@ where let guard = WorkCleanupGuard::new(self, key); let entry = guard.entry(); let result = entry + .cell .get_or_init(async || { let result = func().await; - self.map.remove_if_current(entry); + self.remove_if_current(entry); result }) .await @@ -266,9 +392,10 @@ where let guard = WorkCleanupGuard::new(self, key); let entry = guard.entry(); let result = entry + .cell .get_or_try_init(async || { let result = func().await?; - self.map.remove_if_current(entry); + self.remove_if_current(entry); Ok(result) }) .await? @@ -286,6 +413,6 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - self.map.remove(key); + self.remove(key); } } diff --git a/asyncband/src/singleflight/table.rs b/asyncband/src/singleflight/table.rs deleted file mode 100644 index 73699ef..0000000 --- a/asyncband/src/singleflight/table.rs +++ /dev/null @@ -1,169 +0,0 @@ -// 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::borrow::Borrow; -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::cache_padded::CachePadded; -use crate::internal::default_shard_count; -use crate::internal::mutex::Mutex; -use crate::once::OnceCell; - -type Entries = HashTable>>; -type Shard = CachePadded>>; - -pub struct Entry { - hash: u64, - key: K, - cell: OnceCell, -} - -impl Entry { - fn initialized(&self) -> bool { - self.cell.initialized() - } - - pub async fn get_or_init(&self, init: F) -> &V - where - F: AsyncFnOnce() -> V, - { - self.cell.get_or_init(init).await - } - - pub async fn get_or_try_init(&self, init: F) -> Result<&V, E> - where - F: AsyncFnOnce() -> Result, - { - self.cell.get_or_try_init(init).await - } -} - -/// Storage for one in-flight call per key. -pub struct Table { - shards: Box<[Shard]>, - hasher: S, -} - -impl fmt::Debug for Table -where - K: fmt::Debug, - V: fmt::Debug, -{ - 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(); - debug_map.entries(entries.iter().map(|entry| (&entry.key, &entry.cell))); - } - debug_map.finish() - } -} - -impl Table { - pub fn with_hasher(hasher: S) -> Self { - Self::with_hasher_and_shard_amount(hasher, default_shard_count()) - } - - pub fn with_hasher_and_shard_amount(hasher: S, shard_amount: usize) -> Self { - assert!( - shard_amount.is_power_of_two(), - "shard amount must be greater than zero and a power of two" - ); - - let shards = (0..shard_amount) - .map(|_| CachePadded::new(Mutex::new(HashTable::new()))) - .collect(); - Self { shards, hasher } - } - - fn lock_shard(&self, hash: u64) -> MutexGuard<'_, Entries> { - self.shards[(hash as usize) & (self.shards.len() - 1)].lock() - } - - #[cfg(test)] - pub fn len(&self) -> usize { - self.shards.iter().map(|shard| shard.lock().len()).sum() - } - - #[cfg(test)] - pub fn is_empty(&self) -> bool { - self.shards.iter().all(|shard| shard.lock().is_empty()) - } -} - -impl Table -where - K: Eq + Hash, - S: BuildHasher, -{ - pub fn get_or_insert(&self, key: K) -> Arc> { - let hash = self.hasher.hash_one(&key); - let mut shard = self.lock_shard(hash); - shard - .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) - .or_insert_with(|| { - Arc::new(Entry { - hash, - key, - cell: OnceCell::new(), - }) - }) - .into_mut() - .clone() - } - - pub fn remove(&self, key: &Q) - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - let hash = self.hasher.hash_one(key); - let mut shard = self.lock_shard(hash); - if let Ok(occupied) = shard.find_entry(hash, |entry| entry.key.borrow() == key) { - drop(occupied.remove()); - } - } - - pub fn remove_if_current(&self, entry: &Arc>) { - let mut shard = self.lock_shard(entry.hash); - if let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry)) { - drop(occupied.remove()); - } - } - - pub fn cleanup_abandoned_entry(&self, entry: Arc>) { - let mut shard = self.lock_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. The pointer comparison rejects a detached or replaced entry. - 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); - } -} diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index 61eb3b1..c5792c3 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.is_empty()); + assert!(group.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.len(), 1); + assert_eq!(group.len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(group.map.is_empty()); + assert!(group.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.is_empty()); + assert!(group.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.len(), 1); + assert_eq!(group.len(), 1); assert_eq!(retry.await, Ok("success")); - assert!(group.map.is_empty()); + assert!(group.is_empty()); } diff --git a/tests-integration/tests/once_map_test.rs b/tests-integration/tests/once_map_test.rs index c1bf993..0528313 100644 --- a/tests-integration/tests/once_map_test.rs +++ b/tests-integration/tests/once_map_test.rs @@ -114,6 +114,25 @@ async fn get_remove_and_discard() { assert_eq!(map.get("key"), None); } +#[test] +fn discard_releases_the_removed_value() { + #[derive(Clone)] + struct DropCounter(Arc); + + impl Drop for DropCounter { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let map: OnceMap<_, _> = [(0, DropCounter(Arc::clone(&drops)))].into_iter().collect(); + + map.discard(&0); + + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + #[tokio::test] async fn remove_while_computing_detaches_entry() { let map = Arc::new(OnceMap::new());