perf: take the keyed table read lock on hits and duplicate waiters

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.

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",

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<T: ?Sized>(std::sync::RwLock<T>);
+
+impl<T: ?Sized + fmt::Debug> fmt::Debug for RwLock<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<T> RwLock<T> {
+    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);
+    }
+}

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<V>` to make cloning cheap.
 #[derive(Debug)]
 pub struct OnceMap<K, V, S = RandomState> {
-    map: Mutex<OnceTable<K, V, S>>,
+    map: RwLock<OnceTable<K, V, S>>,
 }
 
 // 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<K, V> {
+    Hit(V),
+    Pending(Arc<OnceTableEntry<K, V>>),
+}
+
 impl<K, V, S> Default for OnceMap<K, V, S>
 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<K, V> {
+        {
+            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<V, E>,
     {
-        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>,
         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>,
         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>,
         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),
         }
     }
 }

asyncband/src/once/once_map/tests.rs

@@ -29,7 +29,7 @@ async fn failed_compute_removes_empty_entry() {
     let result: Result<i32, &str> = 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));
 }