From dba8ecd1bbc5cf338af7038135a42973a7067814 Mon Sep 17 00:00:00 2001 From: An Long Date: Thu, 27 Aug 2026 00:01:45 +0900 Subject: [PATCH 1/2] perf: take the keyed table read lock on hits and duplicate waiters Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 2 + asyncband/src/internal/mod.rs | 3 + asyncband/src/internal/rwlock.rs | 62 +++++++++++++++++++++ asyncband/src/once/once_map/mod.rs | 65 ++++++++++++---------- asyncband/src/once/once_map/tests.rs | 77 ++++++++++++++++++++++++-- asyncband/src/singleflight/mod.rs | 31 +++++++---- asyncband/src/singleflight/tests.rs | 64 +++++++++++++++++++-- tests-integration/tests/traits_test.rs | 15 +++++ 8 files changed, 269 insertions(+), 50 deletions(-) create mode 100644 asyncband/src/internal/rwlock.rs 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..6452a70 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); } @@ -109,14 +112,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,17 +136,31 @@ 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) -> Arc> { + { + let map = self.map.read(); + if let Some(entry) = map.get(&key) { + return Arc::clone(entry); + } + } + + let mut map = self.map.write(); + Arc::clone(map.get_or_insert(key)) + } + /// Compute the value for the given key if absent. /// /// If the value for the key is already being computed by another task, this task will wait for @@ -155,14 +172,10 @@ where where F: AsyncFnOnce() -> V, { - let entry = { - let mut map = self.map.lock(); - let entry = map.get_or_insert(key); - if let Some(value) = entry.get() { - return value.clone(); - } - Arc::clone(entry) - }; + let entry = self.compute_entry(key); + if let Some(value) = entry.get() { + return value.clone(); + } let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_init(func).await.clone(); @@ -181,14 +194,10 @@ where where F: AsyncFnOnce() -> Result, { - let entry = { - let mut map = self.map.lock(); - let entry = map.get_or_insert(key); - if let Some(value) = entry.get() { - return Ok(value.clone()); - } - Arc::clone(entry) - }; + let entry = self.compute_entry(key); + if let Some(value) = entry.get() { + return Ok(value.clone()); + } let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_try_init(func).await?.clone(); @@ -202,7 +211,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 +226,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 +241,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 +259,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::>>(); From 6c55ca45befc33ce9857ae4b049b88802a68cffe Mon Sep 17 00:00:00 2001 From: An Long Date: Thu, 27 Aug 2026 20:43:29 +0900 Subject: [PATCH 2/2] perf: resolve initialized compute hits under the read lock An initialized hit in compute previously cloned the entry Arc under the read lock and read the value only after releasing it, costing two atomic RMWs on the entry's shared reference count on the hottest path. Resolve the hit to its value while the read lock is still held, so contended hits never touch the entry reference count. The cleanup protocol is unchanged: entry Arcs are still cloned only under the table lock, and only initialized values escape it. Co-Authored-By: Claude Fable 5 --- asyncband/src/once/once_map/mod.rs | 36 +++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index 6452a70..3ed9321 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -93,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, @@ -149,16 +156,23 @@ where // 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) -> Arc> { + fn compute_entry(&self, key: K) -> ComputeLookup { { let map = self.map.read(); if let Some(entry) = map.get(&key) { - return Arc::clone(entry); + if let Some(value) = entry.get() { + return ComputeLookup::Hit(value.clone()); + } + return ComputeLookup::Pending(Arc::clone(entry)); } } let mut map = self.map.write(); - Arc::clone(map.get_or_insert(key)) + 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. @@ -172,10 +186,10 @@ where where F: AsyncFnOnce() -> V, { - let entry = self.compute_entry(key); - if let Some(value) = entry.get() { - return value.clone(); - } + let entry = match self.compute_entry(key) { + ComputeLookup::Hit(value) => return value, + ComputeLookup::Pending(entry) => entry, + }; let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_init(func).await.clone(); @@ -194,10 +208,10 @@ where where F: AsyncFnOnce() -> Result, { - let entry = self.compute_entry(key); - if let Some(value) = entry.get() { - return Ok(value.clone()); - } + let entry = match self.compute_entry(key) { + ComputeLookup::Hit(value) => return Ok(value), + ComputeLookup::Pending(entry) => entry, + }; let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().get_or_try_init(func).await?.clone();