diff --git a/CHANGELOG.md b/CHANGELOG.md index 78bc8acf..ddc9391d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ All notable changes to this project will be documented in this file. * Replace `Semaphore::forget` with `Semaphore::drain_permits` and `Semaphore::forget_exact` with `Semaphore::reduce_permits`; permit-level `forget` methods are unchanged. * Rename `ShutdownSend` and `ShutdownRecv` to `Shutdown` and `ShutdownGuard`; rename `shutdown::new_pair` to `shutdown::new`; make `Shutdown` awaitable for requesting shutdown and awaiting completion; and rename the remaining operations to `request_shutdown`, `watch`, `into_watch`, `is_shutdown_requested`, `shutdown_requested`, and `shutdown_requested_owned`. * Raise the minimum supported Rust version from 1.85.0 to 1.86.0. +* Require the hasher to be `Sync` for `OnceMap` and `singleflight::Group` to be `Sync`; lookups now hash keys and pick shards outside the exclusive lock, so concurrent readers share the hasher. The default `RandomState` and other common hashers are unaffected. ### Bug fixes @@ -31,3 +32,4 @@ All notable changes to this project will be documented in this file. ### Improvements * Remove the `slab` dependency in favor of a focused internal waiter arena. +* Shard the keyed table behind `OnceMap` and `singleflight::Group` and serve initialized hits and duplicate waiters under a shard read lock, so concurrent lookups no longer serialize and operations on different keys proceed in parallel. diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index cdd33e34..8b0142e9 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -44,6 +44,9 @@ pub(crate) mod countdown; #[allow(dead_code)] pub(crate) mod once_table; +#[cfg(any(feature = "once-map", feature = "singleflight"))] +pub(crate) mod rwlock; + #[cfg(any(feature = "lazy-cell", feature = "once-cell"))] // `LazyCell` and `OnceCell` use different subsets of `ValueCell`, so single-feature builds leave // some operations in the shared implementation unused. diff --git a/asyncband/src/internal/once_table.rs b/asyncband/src/internal/once_table.rs index 27865ed3..fdbd3b7a 100644 --- a/asyncband/src/internal/once_table.rs +++ b/asyncband/src/internal/once_table.rs @@ -20,11 +20,18 @@ use std::fmt; use std::hash::BuildHasher; use std::hash::Hash; use std::sync::Arc; +use std::sync::RwLockReadGuard; +use std::sync::RwLockWriteGuard; use hashbrown::HashTable; +use crate::internal::rwlock::RwLock; use crate::once::OnceCell; +const SHARD_COUNT: usize = 64; + +type Entries = HashTable>>; + pub struct OnceTableEntry { hash: u64, key: K, @@ -55,9 +62,16 @@ impl OnceTableEntry { } } +/// Outcome of looking a key up for compute: an initialized entry resolves to its value while the +/// shard lock is still held, so contended hits never touch the entry's shared reference count. +pub enum OnceTableLookup { + Hit(V), + Pending(Arc>), +} + /// Shared keyed storage that lets once primitives clean up an exact entry without cloning its key. pub struct OnceTable { - entries: HashTable>>, + shards: Box<[RwLock>]>, hasher: S, } @@ -67,35 +81,44 @@ 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.read(); + 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(|_| RwLock::new(HashTable::with_capacity(shard_capacity))) + .collect(); + Self { shards, hasher } + } + + fn shard_read(&self, hash: u64) -> RwLockReadGuard<'_, Entries> { + self.shards[hash as usize & (SHARD_COUNT - 1)].read() + } + + fn shard_write(&self, hash: u64) -> RwLockWriteGuard<'_, Entries> { + self.shards[hash as usize & (SHARD_COUNT - 1)].write() } #[cfg(test)] pub fn len(&self) -> usize { - self.entries.len() + self.shards.iter().map(|shard| shard.read().len()).sum() } #[cfg(test)] pub fn is_empty(&self) -> bool { - self.entries.is_empty() + self.shards.iter().all(|shard| shard.read().is_empty()) } } @@ -104,37 +127,99 @@ 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 shard = self.shard_read(hash); + if let Some(entry) = shard.find(hash, |entry| entry.key.eq(&key)) { + return Arc::clone(entry); + } + } + + let mut shard = self.shard_write(hash); + Arc::clone( + shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .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_read(hash) + .find(hash, |entry| entry.key.borrow() == key) + .map(Arc::clone) } - pub fn remove(&mut self, key: &Q) -> Option>> + /// Clones the value of an initialized entry under the shard read lock, so hits never touch + /// the entry's shared reference count. + pub fn get_value(&self, key: &Q) -> Option where K: Borrow, Q: Eq + Hash + ?Sized, + V: Clone, { let hash = self.hasher.hash_one(key); - let entry = self - .entries + self.shard_read(hash) + .find(hash, |entry| entry.key.borrow() == key)? + .get() + .cloned() + } + + /// Looks the key up under the shard read lock: an initialized hit resolves to its value + /// without cloning the entry, a pending entry is returned to wait on, and only an absent key + /// takes the shard write lock to insert. + pub fn lookup_or_insert(&self, key: K) -> OnceTableLookup + where + V: Clone, + { + let hash = self.hasher.hash_one(&key); + { + let shard = self.shard_read(hash); + if let Some(entry) = shard.find(hash, |entry| entry.key.eq(&key)) { + if let Some(value) = entry.get() { + return OnceTableLookup::Hit(value.clone()); + } + return OnceTableLookup::Pending(Arc::clone(entry)); + } + } + + let mut shard = self.shard_write(hash); + let entry = shard + .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash) + .or_insert_with(|| { + Arc::new(OnceTableEntry { + hash, + key, + cell: OnceCell::new(), + }) + }) + .into_mut(); + if let Some(value) = entry.get() { + return OnceTableLookup::Hit(value.clone()); + } + OnceTableLookup::Pending(Arc::clone(entry)) + } + + pub fn remove(&self, key: &Q) -> Option>> + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let mut shard = self.shard_write(hash); + let entry = shard .find_entry(hash, |entry| entry.key.borrow() == key) .ok()?; let (entry, _) = entry.remove(); @@ -142,18 +227,35 @@ 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_write(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_write(entry.hash); + // If the table still owns this entry, a count of two means the current call is its only + // owner outside the table: entries are only cloned out of their shard while holding the + // shard lock, so the write lock excludes new owners while the count is checked, and + // owners that release outside the lock do so only after the cell is initialized or the + // entry was detached. The ptr_eq probe rejects an entry that was detached or replaced. + if Arc::strong_count(&entry) == 2 && !entry.initialized() { + if let Ok(occupied) = shard.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, &entry)) + { + 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 +264,7 @@ where key, cell: OnceCell::from_value(value), }); - self.entries.insert_unique(hash, entry, |entry| entry.hash); + self.shard_write(hash) + .insert_unique(hash, entry, |entry| entry.hash); } } diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs new file mode 100644 index 00000000..bedf51bc --- /dev/null +++ b/asyncband/src/internal/rwlock.rs @@ -0,0 +1,62 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt; +use std::sync::PoisonError; + +pub struct RwLock(std::sync::RwLock); + +impl fmt::Debug for RwLock { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl RwLock { + pub const fn new(t: T) -> Self { + Self(std::sync::RwLock::new(t)) + } + + pub fn read(&self) -> std::sync::RwLockReadGuard<'_, T> { + self.0.read().unwrap_or_else(PoisonError::into_inner) + } + + pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, T> { + self.0.write().unwrap_or_else(PoisonError::into_inner) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::internal::rwlock::RwLock; + + #[test] + fn test_poison_rwlock() { + let rwlock = Arc::new(RwLock::new(42)); + let r = rwlock.clone(); + let handle = std::thread::spawn(move || { + let _guard = r.write(); + panic!("poison"); + }); + let _ = handle.join(); + assert_eq!(*rwlock.read(), 42); + let guard = rwlock.write(); + assert_eq!(*guard, 42); + } +} diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index d704ac6a..02d46c42 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -21,9 +21,9 @@ 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; +use crate::internal::once_table::OnceTableLookup; #[cfg(test)] mod tests; @@ -34,7 +34,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 +78,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 +101,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 +122,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,13 +144,9 @@ 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.map.lookup_or_insert(key) { + OnceTableLookup::Hit(value) => return value, + OnceTableLookup::Pending(entry) => entry, }; let guard = ComputeCleanupGuard::new(self, entry); @@ -181,13 +166,9 @@ 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.map.lookup_or_insert(key) { + OnceTableLookup::Hit(value) => return Ok(value), + OnceTableLookup::Pending(entry) => entry, }; let guard = ComputeCleanupGuard::new(self, entry); @@ -202,9 +183,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let map = self.map.lock(); - let entry = map.get(key)?; - entry.get().cloned() + self.map.get_value(key) } /// Remove the given key from the map. @@ -217,8 +196,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); - map.remove(key); + self.map.remove(key); } /// Remove the given key from the map and return a *clone* of the value if exists. @@ -232,7 +210,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.lock().remove(key)?; + let entry = self.map.remove(key)?; entry.get().cloned() } } @@ -244,13 +222,11 @@ where S: Default + BuildHasher, { fn from_iter>(iter: T) -> Self { - let mut map = OnceTable::with_hasher(S::default()); + let map = OnceTable::with_hasher(S::default()); for (key, value) in iter { map.insert(key, value); } - Self { - map: Mutex::new(map), - } + Self { map } } } diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index 1dafc702..ca4341b5 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,74 @@ 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)); } + +#[tokio::test] +async fn cancelled_waiter_preserves_pending_entry() { + let map = OnceMap::<&str, i32>::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let leader = map.compute("key", async move || { + release_rx.await.unwrap(); + 1 + }); + tokio::pin!(leader); + assert!(poll_once(leader.as_mut()).is_pending()); + + let mut waiter = Box::pin(map.compute("key", async || unreachable!())); + assert!(poll_once(waiter.as_mut()).is_pending()); + drop(waiter); + + assert_eq!(map.map.len(), 1); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + assert_eq!(map.get("key"), Some(1)); +} + +#[tokio::test] +async fn cancelled_waiter_preserves_initialized_entry() { + let map = OnceMap::<&str, i32>::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let leader = map.compute("key", async move || { + release_rx.await.unwrap(); + 1 + }); + tokio::pin!(leader); + assert!(poll_once(leader.as_mut()).is_pending()); + + let mut waiter = Box::pin(map.compute("key", async || unreachable!())); + assert!(poll_once(waiter.as_mut()).is_pending()); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + + drop(waiter); + + assert_eq!(map.map.len(), 1); + assert_eq!(map.get("key"), Some(1)); +} + +#[tokio::test] +async fn cancelled_compute_preserves_replacement_entry() { + let map = OnceMap::<&str, i32>::new(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + + let mut first = Box::pin(map.compute("key", async move || { + release_rx.await.unwrap(); + 1 + })); + assert!(poll_once(first.as_mut()).is_pending()); + + map.discard("key"); + assert_eq!(map.compute("key", async || 2).await, 2); + + drop(first); + + assert_eq!(map.map.len(), 1); + assert_eq!(map.get("key"), Some(2)); +} diff --git a/asyncband/src/singleflight/mod.rs b/asyncband/src/singleflight/mod.rs index cb117304..cbb37528 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -23,7 +23,6 @@ use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; -use crate::internal::mutex::Mutex; use crate::internal::once_table::OnceTable; use crate::internal::once_table::OnceTableEntry; @@ -34,7 +33,7 @@ mod tests; /// units of work can be executed with duplicate suppression. #[derive(Debug)] pub struct Group { - map: Mutex>, + map: OnceTable, } // Holds one call's entry so Drop can clean it up if the work is abandoned. @@ -53,10 +52,7 @@ where S: BuildHasher, { fn new(group: &'a Group, key: K) -> Self { - let entry = { - let mut map = group.map.lock(); - Arc::clone(map.get_or_insert(key)) - }; + let entry = group.map.get_or_insert(key); Self { group, @@ -83,15 +79,7 @@ where return; }; - let mut table = self.group.map.lock(); - // If the table still owns this entry, a count of two means the current call is its only - // owner outside the table. remove_entry rejects an entry that was detached or replaced. - if Arc::strong_count(&entry) == 2 && !entry.initialized() { - table.remove_entry(&entry); - } - // Drop this call's reference before unlocking so a waiting cleanup observes the updated - // reference count. - drop(entry); + self.group.map.cleanup_abandoned_entry(entry); } } @@ -114,7 +102,7 @@ where /// Creates a new Group with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(RandomState::new())), + map: OnceTable::with_hasher(RandomState::new()), } } } @@ -128,7 +116,7 @@ where /// Creates a new Group with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(hasher)), + map: OnceTable::with_hasher(hasher), } } @@ -193,7 +181,7 @@ where let result = entry .get_or_init(async || { let result = func().await; - self.map.lock().remove_entry(entry); + self.map.remove_entry(entry); result }) .await @@ -257,7 +245,7 @@ where let result = entry .get_or_try_init(async || { let result = func().await?; - self.map.lock().remove_entry(entry); + self.map.remove_entry(entry); Ok(result) }) .await? @@ -275,7 +263,6 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); - map.remove(key); + self.map.remove(key); } } diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index 9f3dbb15..0139e3fd 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,59 @@ 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()); +} + +#[tokio::test] +async fn cancelled_waiter_preserves_inflight_entry() { + let group = Group::<&str, i32>::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let leader = group.work("key", || async move { + release_rx.await.unwrap(); + 1 + }); + tokio::pin!(leader); + assert!(poll_once(leader.as_mut()).is_pending()); + + let mut waiter = Box::pin(group.work("key", || async { unreachable!() })); + assert!(poll_once(waiter.as_mut()).is_pending()); + drop(waiter); + + assert_eq!(group.map.len(), 1); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + assert!(group.map.is_empty()); +} + +#[tokio::test] +async fn cancelled_work_preserves_replacement_entry() { + let group = Group::<&str, i32>::new(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + let (replacement_tx, replacement_rx) = tokio::sync::oneshot::channel(); + + let mut first = Box::pin(group.work("key", || async move { + release_rx.await.unwrap(); + 1 + })); + assert!(poll_once(first.as_mut()).is_pending()); + + group.forget("key"); + + let second = group.work("key", || async move { + replacement_rx.await.unwrap(); + 2 + }); + tokio::pin!(second); + assert!(poll_once(second.as_mut()).is_pending()); + + drop(first); + assert_eq!(group.map.len(), 1); + + replacement_tx.send(()).unwrap(); + assert_eq!(second.await, 2); + assert!(group.map.is_empty()); } diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index db52c620..20a6ba5c 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -16,6 +16,8 @@ // under the License. use std::cell::Cell; +use std::collections::hash_map::DefaultHasher; +use std::hash::BuildHasher; use asyncband::barrier::Barrier; use asyncband::condvar::Condvar; @@ -45,6 +47,30 @@ use asyncband::waitgroup::WaitGroup; struct PoolManager; +struct LocalBuildHasher(Cell); + +impl BuildHasher for LocalBuildHasher { + type Hasher = DefaultHasher; + + fn build_hasher(&self) -> Self::Hasher { + self.0.set(self.0.get().wrapping_add(1)); + DefaultHasher::new() + } +} + +// A hasher that is Send but not Sync. OnceMap and Group hash keys and pick shards outside the +// exclusive lock, so being Sync requires the hasher to be Sync, while being Send does not. +#[allow(dead_code)] +struct SendOnlyState(std::marker::PhantomData>); + +impl BuildHasher for SendOnlyState { + type Hasher = DefaultHasher; + + fn build_hasher(&self) -> Self::Hasher { + DefaultHasher::new() + } +} + impl ManageObject for PoolManager { type Object = i64; type Error = std::convert::Infallible; @@ -104,11 +130,22 @@ fn movable_public_types_are_send() { fn assert_send() {} assert_send::>>(); + assert_send::>(); + assert_send::>(); assert_send::>(); assert_send::>(); assert_send::>>(); } +#[tokio::test] +async fn keyed_types_accept_non_sync_hashers_for_local_use() { + let map = OnceMap::<&str, u32, _>::with_hasher(LocalBuildHasher(Cell::new(0))); + assert_eq!(map.compute("key", async || 1).await, 1); + + let group = singleflight::Group::<&str, u32, _>::with_hasher(LocalBuildHasher(Cell::new(0))); + assert_eq!(group.work("key", || async { 1 }).await, 1); +} + #[test] fn public_types_are_unpin() { fn assert_unpin() {}