diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs new file mode 100644 index 00000000..f4400ebc --- /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 e2676a84..df5b5043 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -41,10 +41,7 @@ pub(crate) mod arena; 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; +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 @@ -65,6 +62,9 @@ pub(crate) mod value_cell; ))] pub(crate) mod mutex; +#[cfg(feature = "once-map")] +pub(crate) mod rwlock; + #[cfg(any( feature = "mpsc", feature = "mutex", @@ -97,3 +97,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 9477e26b..ce5229f8 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/once_table.rs b/asyncband/src/internal/once_table.rs deleted file mode 100644 index 27865ed3..00000000 --- a/asyncband/src/internal/once_table.rs +++ /dev/null @@ -1,167 +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 hashbrown::HashTable; - -use crate::once::OnceCell; - -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 { - entries: HashTable>>, - hasher: S, -} - -impl fmt::Debug for OnceTable -where - K: fmt::Debug, - 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() - } -} - -impl OnceTable { - pub fn with_hasher(hasher: S) -> Self { - Self { - entries: HashTable::new(), - hasher, - } - } - - pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { - Self { - entries: HashTable::with_capacity(capacity), - hasher, - } - } - - #[cfg(test)] - pub fn len(&self) -> usize { - self.entries.len() - } - - #[cfg(test)] - pub fn is_empty(&self) -> bool { - self.entries.is_empty() - } -} - -impl OnceTable -where - K: Eq + Hash, - S: BuildHasher, -{ - pub fn get_or_insert(&mut 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(), - }) - }) - .into_mut() - } - - pub fn get(&self, key: &Q) -> Option<&Arc>> - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - let hash = self.hasher.hash_one(key); - self.entries.find(hash, |entry| entry.key.borrow() == key) - } - - pub fn remove(&mut self, key: &Q) -> Option>> - where - K: Borrow, - Q: Eq + Hash + ?Sized, - { - let hash = self.hasher.hash_one(key); - let entry = self - .entries - .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(&mut self, entry: &Arc>) { - let Ok(occupied) = self - .entries - .find_entry(entry.hash, |existing| Arc::ptr_eq(existing, entry)) - else { - return; - }; - - drop(occupied.remove()); - } - - pub fn insert(&mut 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.entries.insert_unique(hash, entry, |entry| entry.hash); - } -} diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs new file mode 100644 index 00000000..906947fa --- /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/internal/waitset.rs b/asyncband/src/internal/waitset.rs index ddf8bf2e..8f4ab1f2 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()); diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 87c7c8fe..45bc7406 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/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index d704ac6a..bc361e39 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -16,25 +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 hashbrown::HashTable; + +use crate::internal::cache_padded::CachePadded; +use crate::internal::default_shard_count; use crate::internal::mutex::Mutex; -use crate::internal::once_table::OnceTable; -use crate::internal::once_table::OnceTableEntry; +use crate::internal::rwlock::RwLock; +use crate::once::OnceCell; #[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: Mutex>, + // 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. @@ -44,7 +378,7 @@ where S: BuildHasher, { once_map: &'a OnceMap, - entry: Option>>, + entry: Option>>, } impl<'a, K, V, S> ComputeCleanupGuard<'a, K, V, S> @@ -52,14 +386,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() } @@ -78,15 +412,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.cleanup_abandoned_entry(entry); } } @@ -108,19 +434,31 @@ where { /// Creates a new OnceMap with the default hasher. pub fn new() -> Self { - Self { - map: Mutex::new(OnceTable::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: Mutex::new(OnceTable::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. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_shard_amount(shard_amount: usize) -> Self { + Self::with_config(0, 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::with_config(capacity, RandomState::new(), shard_amount) } } @@ -132,16 +470,34 @@ where { /// Creates a new OnceMap with the given hasher. pub fn with_hasher(hasher: S) -> Self { - Self { - map: Mutex::new(OnceTable::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: Mutex::new(OnceTable::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. + /// + /// # 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::with_config(0, 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::with_config(capacity, hasher, shard_amount) } /// Compute the value for the given key if absent. @@ -155,17 +511,13 @@ 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 = match self.get_or_insert(key) { + Lookup::Ready(value) => return value, + 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 } @@ -181,17 +533,13 @@ 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 = match self.get_or_insert(key) { + Lookup::Ready(value) => return Ok(value), + 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) } @@ -202,9 +550,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let map = self.map.lock(); - let entry = map.get(key)?; - entry.get().cloned() + self.get_value(key) } /// Remove the given key from the map. @@ -217,8 +563,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); - map.remove(key); + self.remove_entry(key); } /// Remove the given key from the map and return a *clone* of the value if exists. @@ -232,25 +577,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.lock().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 { - let mut map = OnceTable::with_hasher(S::default()); - for (key, value) in iter { - map.insert(key, value); - } - - Self { - map: Mutex::new(map), - } + let entry = self.remove_entry(key)?; + entry.cell.get().cloned() } } diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index 1dafc702..52ee05dc 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.lock().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.lock().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.lock().len(), 1); + assert_eq!(map.len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(map.map.lock().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.lock().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/rwlock/mod.rs b/asyncband/src/rwlock/mod.rs index 9f048473..36e97ca8 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. /// diff --git a/asyncband/src/singleflight/mod.rs b/asyncband/src/singleflight/mod.rs index cb117304..88a1e792 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 hashbrown::HashTable; + +use crate::internal::cache_padded::CachePadded; +use crate::internal::default_shard_count; use crate::internal::mutex::Mutex; -use crate::internal::once_table::OnceTable; -use crate::internal::once_table::OnceTableEntry; +use crate::once::OnceCell; #[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: Mutex>, + 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. @@ -44,7 +177,7 @@ where S: BuildHasher, { group: &'a Group, - entry: Option>>, + entry: Option>>, } impl<'a, K, V, S> WorkCleanupGuard<'a, K, V, S> @@ -53,10 +186,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.get_or_insert(key); Self { group, @@ -64,7 +194,7 @@ where } } - fn entry(&self) -> &Arc> { + fn entry(&self) -> &Arc> { self.entry.as_ref().unwrap() } @@ -83,15 +213,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.cleanup_abandoned_entry(entry); } } @@ -113,9 +235,16 @@ where { /// Creates a new Group with the default hasher. pub fn new() -> Self { - Self { - map: Mutex::new(OnceTable::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. + /// + /// # Panics + /// + /// Panics if `shard_amount` is zero or is not a power of two. + pub fn with_shard_amount(shard_amount: usize) -> Self { + Self::with_config(RandomState::new(), shard_amount) } } @@ -127,9 +256,16 @@ where { /// Creates a new Group with the given hasher. pub fn with_hasher(hasher: S) -> Self { - Self { - map: Mutex::new(OnceTable::with_hasher(hasher)), - } + Self::with_config(hasher, default_shard_count()) + } + + /// 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::with_config(hasher, shard_amount) } /// Executes and returns the results of the given function, making sure that only one execution @@ -191,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.lock().remove_entry(entry); + self.remove_if_current(entry); result }) .await @@ -255,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.lock().remove_entry(entry); + self.remove_if_current(entry); Ok(result) }) .await? @@ -275,7 +413,6 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); - map.remove(key); + self.remove(key); } } diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index 9f3dbb15..c5792c38 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.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.len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(group.map.lock().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.lock().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.lock().len(), 1); + assert_eq!(group.len(), 1); assert_eq!(retry.await, Ok("success")); - assert!(group.map.lock().is_empty()); + assert!(group.is_empty()); } diff --git a/benchmarks/asyncband/once_map/compute.rs b/benchmarks/asyncband/once_map/compute.rs index e5733fc2..71ffd2c8 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 cef31a02..287eb145 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 44c59c8a..b5b155b5 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 62ffefb2..7d706008 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 ce12a638..cba7cbd6 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) { diff --git a/tests-integration/tests/once_map_test.rs b/tests-integration/tests/once_map_test.rs index c1bf9936..05283133 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());