diff --git a/CHANGELOG.md b/CHANGELOG.md index 78bc8ac..848d341 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 under a shared read 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. +* Serve `OnceMap` and `singleflight::Group` lookups for initialized hits and duplicate waiters under a shared read lock instead of the exclusive table lock, so concurrent lookups no longer serialize. diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index cdd33e3..98e73e2 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -61,6 +61,9 @@ pub(crate) mod value_cell; ))] pub(crate) mod mutex; +#[cfg(any(feature = "once-map", feature = "singleflight"))] +pub(crate) mod rwlock; + #[cfg(any( feature = "mpsc", feature = "mutex", diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs new file mode 100644 index 0000000..bedf51b --- /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 d704ac6..3ed9321 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::rwlock::RwLock; #[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: RwLock>, } // Holds one call's entry so Drop can clean it up if the computation is abandoned. @@ -78,9 +78,12 @@ where return; }; - let mut table = self.once_map.map.lock(); + let mut table = self.once_map.map.write(); // If the table still owns this entry, a count of two means the current call is its only - // owner outside the table. remove_entry rejects an entry that was detached or replaced. + // owner outside the table: entries are only cloned out of the table under the table 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. remove_entry rejects an entry that was detached or replaced. if Arc::strong_count(&entry) == 2 && !entry.initialized() { table.remove_entry(&entry); } @@ -90,6 +93,13 @@ where } } +// Outcome of looking a key up for compute: an initialized entry resolves to its value while the +// table lock is still held, so contended hits never touch the entry's shared reference count. +enum ComputeLookup { + Hit(V), + Pending(Arc>), +} + impl Default for OnceMap where K: Eq + Hash, @@ -109,14 +119,14 @@ where /// Creates a new OnceMap with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(RandomState::new())), + map: RwLock::new(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( + map: RwLock::new(OnceTable::with_capacity_and_hasher( capacity, RandomState::new(), )), @@ -133,15 +143,36 @@ 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: RwLock::new(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: RwLock::new(OnceTable::with_capacity_and_hasher(capacity, hasher)), + } + } + + // Looks the key up under the read lock so initialized hits and duplicate waiters stay off the + // exclusive lock, and inserts under the write lock only when the key is absent. + fn compute_entry(&self, key: K) -> ComputeLookup { + { + let map = self.map.read(); + if let Some(entry) = map.get(&key) { + if let Some(value) = entry.get() { + return ComputeLookup::Hit(value.clone()); + } + return ComputeLookup::Pending(Arc::clone(entry)); + } + } + + let mut map = self.map.write(); + let entry = map.get_or_insert(key); + if let Some(value) = entry.get() { + return ComputeLookup::Hit(value.clone()); } + ComputeLookup::Pending(Arc::clone(entry)) } /// Compute the value for the given key if absent. @@ -155,13 +186,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.compute_entry(key) { + ComputeLookup::Hit(value) => return value, + ComputeLookup::Pending(entry) => entry, }; let guard = ComputeCleanupGuard::new(self, entry); @@ -181,13 +208,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.compute_entry(key) { + ComputeLookup::Hit(value) => return Ok(value), + ComputeLookup::Pending(entry) => entry, }; let guard = ComputeCleanupGuard::new(self, entry); @@ -202,7 +225,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let map = self.map.lock(); + let map = self.map.read(); let entry = map.get(key)?; entry.get().cloned() } @@ -217,7 +240,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); + let mut map = self.map.write(); map.remove(key); } @@ -232,7 +255,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let entry = self.map.lock().remove(key)?; + let entry = self.map.write().remove(key)?; entry.get().cloned() } } @@ -250,7 +273,7 @@ where } Self { - map: Mutex::new(map), + map: RwLock::new(map), } } } diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index 1dafc70..1e89329 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.read().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.read().is_empty()); } #[tokio::test] @@ -65,11 +65,78 @@ async fn cancelled_compute_removes_empty_entry() { }); started_rx.await.unwrap(); - assert_eq!(map.map.lock().len(), 1); + assert_eq!(map.map.read().len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(map.map.lock().is_empty()); + assert!(map.map.read().is_empty()); +} + +#[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.read().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.read().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.read().len(), 1); + assert_eq!(map.get("key"), Some(2)); } #[tokio::test] @@ -91,7 +158,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.read().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..0ca7d58 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -23,9 +23,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::rwlock::RwLock; #[cfg(test)] mod tests; @@ -34,7 +34,7 @@ mod tests; /// units of work can be executed with duplicate suppression. #[derive(Debug)] pub struct Group { - map: Mutex>, + map: RwLock>, } // Holds one call's entry so Drop can clean it up if the work is abandoned. @@ -53,9 +53,15 @@ where S: BuildHasher, { fn new(group: &'a Group, key: K) -> Self { + // Looks the key up under the read lock so duplicate waiters stay off the exclusive lock, + // and inserts under the write lock only when no call is in flight for the key. let entry = { - let mut map = group.map.lock(); - Arc::clone(map.get_or_insert(key)) + let map = group.map.read(); + map.get(&key).map(Arc::clone) + }; + let entry = match entry { + Some(entry) => entry, + None => Arc::clone(group.map.write().get_or_insert(key)), }; Self { @@ -83,9 +89,12 @@ where return; }; - let mut table = self.group.map.lock(); + let mut table = self.group.map.write(); // If the table still owns this entry, a count of two means the current call is its only - // owner outside the table. remove_entry rejects an entry that was detached or replaced. + // owner outside the table: entries are only cloned out of the table under the table 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. remove_entry rejects an entry that was detached or replaced. if Arc::strong_count(&entry) == 2 && !entry.initialized() { table.remove_entry(&entry); } @@ -114,7 +123,7 @@ where /// Creates a new Group with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(OnceTable::with_hasher(RandomState::new())), + map: RwLock::new(OnceTable::with_hasher(RandomState::new())), } } } @@ -128,7 +137,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: RwLock::new(OnceTable::with_hasher(hasher)), } } @@ -193,7 +202,7 @@ where let result = entry .get_or_init(async || { let result = func().await; - self.map.lock().remove_entry(entry); + self.map.write().remove_entry(entry); result }) .await @@ -257,7 +266,7 @@ where let result = entry .get_or_try_init(async || { let result = func().await?; - self.map.lock().remove_entry(entry); + self.map.write().remove_entry(entry); Ok(result) }) .await? @@ -275,7 +284,7 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let mut map = self.map.lock(); + let mut map = self.map.write(); map.remove(key); } } diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index 9f3dbb1..7c0764c 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.read().is_empty()); let result = group.work("key", || async { "success".to_owned() }).await; assert_eq!(result, "success"); @@ -58,11 +58,63 @@ async fn cancelled_work_removes_empty_entry() { }); started_rx.await.unwrap(); - assert_eq!(group.map.lock().len(), 1); + assert_eq!(group.map.read().len(), 1); task.abort(); assert!(task.await.unwrap_err().is_cancelled()); - assert!(group.map.lock().is_empty()); + assert!(group.map.read().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.read().len(), 1); + + release_tx.send(()).unwrap(); + assert_eq!(leader.await, 1); + assert!(group.map.read().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.read().len(), 1); + + replacement_tx.send(()).unwrap(); + assert_eq!(second.await, 2); + assert!(group.map.read().is_empty()); } #[tokio::test] @@ -73,7 +125,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.read().is_empty()); let retry = group .try_work("key", || async { Ok::<&str, ()>("success") }) @@ -100,7 +152,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.read().len(), 1); assert_eq!(retry.await, Ok("success")); - assert!(group.map.lock().is_empty()); + assert!(group.map.read().is_empty()); } diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index db52c62..ad3d512 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -62,6 +62,19 @@ impl ManageObject for PoolManager { } } +// A hasher that is Send but not Sync. OnceMap and Group hash keys under a shared read lock, so +// being Sync requires the hasher to be Sync, while being Send does not. +#[allow(dead_code)] +struct SendOnlyState(std::marker::PhantomData>); + +impl std::hash::BuildHasher for SendOnlyState { + type Hasher = std::hash::DefaultHasher; + + fn build_hasher(&self) -> Self::Hasher { + std::hash::DefaultHasher::new() + } +} + #[test] fn public_types_are_send_and_sync() { fn assert_send_and_sync() {} @@ -104,6 +117,8 @@ fn movable_public_types_are_send() { fn assert_send() {} assert_send::>>(); + assert_send::>(); + assert_send::>(); assert_send::>(); assert_send::>(); assert_send::>>();