From 8317a740f206cbd929496ed950166c359b7adc2c Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Fri, 28 Aug 2026 03:11:55 +0800 Subject: [PATCH 1/3] 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/3] 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/3] 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); + } +}