refactor(once_map): build the table before locking it and simplify lookups

asyncband/src/once/once_map/mod.rs

@@ -30,26 +30,19 @@ use crate::once::OnceCell;
 #[cfg(test)]
 mod tests;
 
-type Entries<K, V> = HashTable<Arc<Entry<K, V>>>;
-
 struct Entry<K, V> {
     hash: u64,
     key: K,
     cell: OnceCell<V>,
 }
 
-enum Lookup<K, V> {
-    Ready(V),
-    Pending(Arc<Entry<K, V>>),
-}
-
 /// A hash map that runs computation only once for each key and stores the result.
 ///
 /// Every successful lookup returns an owned clone of `V`. Choose a cheaply cloned value type, such
 /// as `Arc<T>`, when entries are large or frequently shared.
 pub struct OnceMap<K, V, S = RandomState> {
     // Hashbrown allocates the table lazily, and computation always runs after releasing this lock.
-    entries: Mutex<Entries<K, V>>,
+    entries: Mutex<HashTable<Arc<Entry<K, V>>>>,
     hasher: S,
 }
 
@@ -75,62 +68,21 @@ where
     K: Eq + Hash,
     S: BuildHasher,
 {
-    fn get_or_insert(&self, key: K) -> Lookup<K, V>
-    where
-        V: Clone,
-    {
+    fn get_or_insert(&self, key: K) -> Arc<Entry<K, V>> {
         let hash = self.hasher.hash_one(&key);
-        let entry = {
-            let mut entries = self.entries.lock();
-            if let Some(entry) = entries
-                .find(hash, |entry| entry.key.eq(&key))
-                .map(Arc::clone)
-            {
-                entry
-            } else {
-                let entry = Arc::new(Entry {
-                    hash,
-                    key,
-                    cell: OnceCell::new(),
-                });
-                entries.insert_unique(hash, Arc::clone(&entry), |entry| entry.hash);
-                entry
-            }
-        };
-
-        Self::classify(entry)
-    }
-
-    fn classify(entry: Arc<Entry<K, V>>) -> Lookup<K, V>
-    where
-        V: Clone,
-    {
-        match entry.cell.get().cloned() {
-            Some(value) => Lookup::Ready(value),
-            None => Lookup::Pending(entry),
+        let mut entries = self.entries.lock();
+        // Drop duplicate keys after unlocking: their destructors may reenter the map.
+        if let Some(entry) = entries.find(hash, |entry| entry.key.eq(&key)) {
+            return entry.clone();
         }
-    }
 
-    fn find_entry(
-        &self,
-        hash: u64,
-        matches: impl Fn(&Entry<K, V>) -> bool,
-    ) -> Option<Arc<Entry<K, V>>> {
-        self.entries
-            .lock()
-            .find(hash, |entry| matches(entry))
-            .cloned()
-    }
-
-    fn get_value<Q>(&self, key: &Q) -> Option<V>
-    where
-        K: Borrow<Q>,
-        Q: Eq + Hash + ?Sized,
-        V: Clone,
-    {
-        let hash = self.hasher.hash_one(key);
-        let entry = self.find_entry(hash, |entry| entry.key.borrow() == key)?;
-        entry.cell.get().cloned()
+        let entry = Arc::new(Entry {
+            hash,
+            key,
+            cell: OnceCell::new(),
+        });
+        entries.insert_unique(hash, entry.clone(), |entry| entry.hash);
+        entry
     }
 
     fn remove_entry<Q>(&self, key: &Q) -> Option<Arc<Entry<K, V>>>
@@ -143,9 +95,7 @@ where
         let occupied = entries
             .find_entry(hash, |entry| entry.key.borrow() == key)
             .ok()?;
-        let (entry, _) = occupied.remove();
-        drop(entries);
-        Some(entry)
+        Some(occupied.remove().0)
     }
 
     fn cleanup_abandoned_entry(&self, entry: Arc<Entry<K, V>>) {
@@ -164,32 +114,14 @@ where
                 Some(occupied.remove().0)
             } else {
                 // A waiting cleanup must observe this call's reference being released before it
-                // can inspect the count while holding the write lock.
+                // can inspect the count while holding the table lock.
                 drop(entry);
                 None
             }
         };
         // Key and value destructors must not run while the table is locked.
         drop(removed);
     }
-
-    fn insert(&mut self, key: K, value: V) {
-        let hash = self.hasher.hash_one(&key);
-        let entry = Arc::new(Entry {
-            hash,
-            key,
-            cell: OnceCell::from_value(value),
-        });
-
-        let mut entries = self.entries.lock();
-        let replaced = entries
-            .find_entry(hash, |stored| stored.key.eq(&entry.key))
-            .ok()
-            .map(|occupied| occupied.remove().0);
-        entries.insert_unique(hash, entry, |entry| entry.hash);
-        drop(entries);
-        drop(replaced);
-    }
 }
 
 impl<K, V, S> FromIterator<(K, V)> for OnceMap<K, V, S>
@@ -200,14 +132,22 @@ where
 {
     fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
         let iter = iter.into_iter();
-        let mut map = Self {
-            entries: Mutex::new(HashTable::with_capacity(iter.size_hint().0)),
-            hasher: S::default(),
-        };
+        let hasher = S::default();
+        let mut entries = HashTable::<Arc<Entry<K, V>>>::with_capacity(iter.size_hint().0);
         for (key, value) in iter {
-            map.insert(key, value);
+            let hash = hasher.hash_one(&key);
+            entries
+                .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash)
+                .insert(Arc::new(Entry {
+                    hash,
+                    key,
+                    cell: OnceCell::from_value(value),
+                }));
+        }
+        Self {
+            entries: Mutex::new(entries),
+            hasher,
         }
-        map
     }
 }
 
@@ -307,12 +247,14 @@ where
     where
         F: AsyncFnOnce() -> V,
     {
-        let entry = match self.get_or_insert(key) {
-            Lookup::Ready(value) => return value,
-            Lookup::Pending(entry) => entry,
+        // The block ends the entry's scope before the await, so the future does not store it.
+        let guard = {
+            let entry = self.get_or_insert(key);
+            if let Some(value) = entry.cell.get() {
+                return value.clone();
+            }
+            ComputeCleanupGuard::new(self, entry)
         };
-
-        let guard = ComputeCleanupGuard::new(self, entry);
         let result = guard.entry().cell.get_or_init(func).await.clone();
         guard.dismiss();
         result
@@ -334,12 +276,14 @@ where
     where
         F: AsyncFnOnce() -> Result<V, E>,
     {
-        let entry = match self.get_or_insert(key) {
-            Lookup::Ready(value) => return Ok(value),
-            Lookup::Pending(entry) => entry,
+        // The block ends the entry's scope before the await, so the future does not store it.
+        let guard = {
+            let entry = self.get_or_insert(key);
+            if let Some(value) = entry.cell.get() {
+                return Ok(value.clone());
+            }
+            ComputeCleanupGuard::new(self, entry)
         };
-
-        let guard = ComputeCleanupGuard::new(self, entry);
         let result = guard.entry().cell.get_or_try_init(func).await?.clone();
         guard.dismiss();
         Ok(result)
@@ -353,7 +297,13 @@ where
         K: Borrow<Q>,
         Q: Hash + Eq + ?Sized,
     {
-        self.get_value(key)
+        let hash = self.hasher.hash_one(key);
+        let entry = self
+            .entries
+            .lock()
+            .find(hash, |entry| entry.key.borrow() == key)
+            .cloned();
+        entry?.cell.get().cloned()
     }
 
     /// Remove the given key from the map.

asyncband/src/once/once_map/tests.rs

@@ -25,7 +25,6 @@ use std::pin::pin;
 use std::sync::Arc;
 use std::task::Poll;
 
-use super::Lookup;
 use super::OnceMap;
 use crate::test_support::poll_once;
 
@@ -90,13 +89,12 @@ fn cancelled_compute_removes_empty_entry() {
 fn pending_computation_does_not_block_another_key() {
     let map = OnceMap::new();
     {
-        let mut pending = std::pin::pin!(
-            map.compute("pending", async || { std::future::pending::<i32>().await })
-        );
+        let mut pending =
+            pin!(map.compute("pending", async || std::future::pending::<i32>().await));
         assert!(poll_once(pending.as_mut()).is_pending());
 
-        let mut ready = std::pin::pin!(map.compute("ready", async || 1));
-        assert_eq!(poll_once(ready.as_mut()), std::task::Poll::Ready(1));
+        let mut ready = pin!(map.compute("ready", async || 1));
+        assert_eq!(poll_once(ready.as_mut()), Poll::Ready(1));
     }
 
     assert_eq!(map.len(), 1);
@@ -107,21 +105,20 @@ fn failed_compute_preserves_entry_for_waiter_retry() {
     let map = OnceMap::new();
     let released = Cell::new(false);
 
-    let first = map.try_compute("key", async || {
+    let mut first = pin!(map.try_compute("key", async || {
         poll_fn(|_| {
-            released
-                .get()
-                .then_some(())
-                .map_or(Poll::Pending, Poll::Ready)
+            if released.get() {
+                Poll::Ready(())
+            } else {
+                Poll::Pending
+            }
         })
         .await;
         Err::<i32, &str>("fail")
-    });
-    let mut first = pin!(first);
+    }));
     assert!(poll_once(first.as_mut()).is_pending());
 
-    let retry = map.try_compute("key", async || Ok::<i32, &str>(1));
-    let mut retry = pin!(retry);
+    let mut retry = pin!(map.try_compute("key", async || Ok::<i32, &str>(1)));
     assert!(poll_once(retry.as_mut()).is_pending());
 
     released.set(true);
@@ -135,15 +132,9 @@ fn failed_compute_preserves_entry_for_waiter_retry() {
 #[test]
 fn colliding_pending_entries_are_tracked_independently() {
     let map: OnceMap<usize, usize, BuildHasherDefault<ConstantHasher>> = OnceMap::default();
-    let Lookup::Pending(first) = map.get_or_insert(1) else {
-        unreachable!()
-    };
-    let Lookup::Pending(first_waiter) = map.get_or_insert(1) else {
-        unreachable!()
-    };
-    let Lookup::Pending(second) = map.get_or_insert(2) else {
-        unreachable!()
-    };
+    let first = map.get_or_insert(1);
+    let first_waiter = map.get_or_insert(1);
+    let second = map.get_or_insert(2);
 
     assert!(Arc::ptr_eq(&first, &first_waiter));
     assert!(!Arc::ptr_eq(&first, &second));

benchmarks/benches/primitives/once_map/build.rs

@@ -0,0 +1,44 @@
+// 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 divan::Bencher;
+use divan::black_box;
+
+use super::support::BATCH_SAMPLE_SIZE;
+use super::support::BenchMap;
+use super::support::HOT_ENTRY_COUNTS;
+
+// Builds a map from distinct keys. The input vector is prepared outside the timed section and the
+// map is dropped after it.
+#[divan::bench(args = HOT_ENTRY_COUNTS, sample_size = BATCH_SAMPLE_SIZE)]
+fn collect_distinct_keys(bencher: Bencher, entries: usize) {
+    bencher
+        .with_inputs(|| (0..entries).map(|key| (key, key)).collect::<Vec<_>>())
+        .bench_local_values(|items| black_box(items.into_iter().collect::<BenchMap>()));
+}
+
+// Every key appears twice, so half of the items replace an entry inserted moments earlier.
+#[divan::bench(args = HOT_ENTRY_COUNTS, sample_size = BATCH_SAMPLE_SIZE)]
+fn collect_duplicate_keys(bencher: Bencher, entries: usize) {
+    bencher
+        .with_inputs(|| {
+            (0..entries)
+                .map(|item| (item % (entries / 2), item))
+                .collect::<Vec<_>>()
+        })
+        .bench_local_values(|items| black_box(items.into_iter().collect::<BenchMap>()));
+}

benchmarks/benches/primitives/once_map/lookup.rs

@@ -46,8 +46,8 @@ fn get_hit_distributed(bencher: Bencher, ready_entries: usize) {
         .bench_values(|key| black_box(map.get(black_box(&key))));
 }
 
-// Negative get is still a read-only steady-state path: it traverses the ready index but never
-// creates pending state. Varying absent hashes avoids measuring one unusually cache-hot miss.
+// Negative get is still a read-only steady-state path: it probes the table but never creates
+// pending state. Varying absent hashes avoids measuring one unusually cache-hot miss.
 #[divan::bench(
     threads = THREAD_COUNTS,
     args = READY_ENTRY_COUNTS,

benchmarks/benches/primitives/once_map/mod.rs

@@ -16,8 +16,9 @@
 // under the License.
 
 //! Uses usize keys, a deterministic hasher, and trivial initializers.
-//! Lookups use threads; coalescing uses controlled polling.
+//! Lookups use threads; coalescing uses controlled polling; builds collect prepared vectors.
 
+mod build;
 mod initialize;
 mod lookup;
 mod maintenance;