diff --git a/CHANGELOG.md b/CHANGELOG.md index 1421c34..7447eff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,11 @@ All notable changes to this project will be documented in this file. * Retry spurious atomic failures when completing `Once` initialization. ([#126](https://github.com/fast/mea/pull/126)) * Make cloning a `WaitGroup` panic on counter overflow instead of silently losing track of a handle. * Align `Condvar` with standard condition-variable semantics by notifying only current waiters and passing a cancelled `notify_one` wakeup to another current waiter instead of storing a permit. +* Clean up uninitialized `OnceMap` and `singleflight` entries once no callers remain after a failure, panic, or cancellation. + +### Improvements + +* Remove unnecessary `Clone` bounds from `singleflight` keys and custom hashers used by keyed once primitives. ## v0.6.5 (2026-07-30) diff --git a/Cargo.lock b/Cargo.lock index b805a3c..ffee835 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -464,6 +464,7 @@ name = "mea" version = "0.6.5" dependencies = [ "divan", + "hashbrown", "pollster", "slab", "tokio", diff --git a/mea/Cargo.toml b/mea/Cargo.toml index f077f45..3149ef7 100644 --- a/mea/Cargo.toml +++ b/mea/Cargo.toml @@ -33,6 +33,9 @@ all-features = true rustdoc-args = ["--cfg", "docsrs"] [dependencies] +hashbrown = { version = "0.17.1", default-features = false, features = [ + "inline-more", +] } slab = { version = "0.4.11" } [dev-dependencies] diff --git a/mea/benches/primitives/main.rs b/mea/benches/primitives/main.rs index c262a72..be86cda 100644 --- a/mea/benches/primitives/main.rs +++ b/mea/benches/primitives/main.rs @@ -12,7 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod once_map; mod oneshot; +mod singleflight; +mod support; fn main() { divan::main(); diff --git a/mea/benches/primitives/once_map.rs b/mea/benches/primitives/once_map.rs new file mode 100644 index 0000000..9c94715 --- /dev/null +++ b/mea/benches/primitives/once_map.rs @@ -0,0 +1,69 @@ +// Copyright 2024 tison +// +// Licensed 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 mea::once::OnceMap; + +use super::support::defer_input_drop; +use super::support::noop_context; +use super::support::poll_ready; + +const CACHED_ENTRY_COUNTS: &[usize] = &[0, 64, 1024]; + +#[divan::bench] +fn compute_vacant(bencher: Bencher) { + let mut context = noop_context(); + bencher + .with_inputs(OnceMap::::new) + .bench_local_values(|map| { + let result = black_box(poll_ready( + map.compute(black_box(0), || async { black_box(1) }), + &mut context, + )); + defer_input_drop(map, result) + }); +} + +#[divan::bench] +fn compute_occupied(bencher: Bencher) { + let mut context = noop_context(); + bencher + .with_inputs(|| [(0, 1)].into_iter().collect::>()) + .bench_local_values(|map| { + let result = black_box(poll_ready( + map.compute(black_box(0), || async { black_box(2) }), + &mut context, + )); + defer_input_drop(map, result) + }); +} + +#[divan::bench(args = CACHED_ENTRY_COUNTS)] +fn try_compute_error(bencher: Bencher, cached_entries: usize) { + let mut context = noop_context(); + bencher + .with_inputs(|| { + (0..cached_entries) + .map(|key| (key, key)) + .collect::>() + }) + .bench_local_values(|map| { + let result = black_box(poll_ready( + map.try_compute(black_box(usize::MAX), || async { Err::(()) }), + &mut context, + )); + defer_input_drop(map, result) + }); +} diff --git a/mea/benches/primitives/singleflight.rs b/mea/benches/primitives/singleflight.rs new file mode 100644 index 0000000..ca65c8e --- /dev/null +++ b/mea/benches/primitives/singleflight.rs @@ -0,0 +1,49 @@ +// Copyright 2024 tison +// +// Licensed 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 mea::singleflight::Group; + +use super::support::defer_input_drop; +use super::support::noop_context; +use super::support::poll_ready; + +#[divan::bench] +fn work_ready(bencher: Bencher) { + let mut context = noop_context(); + bencher + .with_inputs(Group::::new) + .bench_local_values(|group| { + let result = black_box(poll_ready( + group.work(black_box(0), || async { black_box(1) }), + &mut context, + )); + defer_input_drop(group, result) + }); +} + +#[divan::bench] +fn try_work_error(bencher: Bencher) { + let mut context = noop_context(); + bencher + .with_inputs(Group::::new) + .bench_local_values(|group| { + let result = black_box(poll_ready( + group.try_work(black_box(0), || async { Err::(()) }), + &mut context, + )); + defer_input_drop(group, result) + }); +} diff --git a/mea/benches/primitives/support.rs b/mea/benches/primitives/support.rs new file mode 100644 index 0000000..fec1f20 --- /dev/null +++ b/mea/benches/primitives/support.rs @@ -0,0 +1,37 @@ +// Copyright 2024 tison +// +// Licensed 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::future::Future; +use std::pin::pin; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +pub(super) fn noop_context() -> Context<'static> { + Context::from_waker(Waker::noop()) +} + +pub(super) fn poll_ready(future: F, context: &mut Context<'_>) -> F::Output { + let mut future = pin!(future); + match future.as_mut().poll(context) { + Poll::Ready(output) => output, + Poll::Pending => panic!("benchmark future should be ready"), + } +} + +// Move the input into the benchmark output so Divan drops it outside the timed section. +#[inline] +pub(super) fn defer_input_drop(input: I, output: O) -> (I, O) { + (input, output) +} diff --git a/mea/src/admission/tests.rs b/mea/src/admission/tests.rs index 8be10ca..10c775a 100644 --- a/mea/src/admission/tests.rs +++ b/mea/src/admission/tests.rs @@ -13,25 +13,15 @@ // limitations under the License. use std::collections::hash_map::DefaultHasher; -use std::future::Future; use std::hash::BuildHasherDefault; -use std::pin::Pin; use std::pin::pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use std::task::Context; use std::task::Poll; -use std::task::Waker; use super::FairShare; - -fn poll_once(future: Pin<&mut F>) -> Poll -where - F: Future, -{ - future.poll(&mut Context::from_waker(Waker::noop())) -} +use crate::poll_once; #[test] #[should_panic(expected = "FairShare requires at least one permit")] diff --git a/mea/src/condvar/tests.rs b/mea/src/condvar/tests.rs index 6b3eeb6..dbdbb0f 100644 --- a/mea/src/condvar/tests.rs +++ b/mea/src/condvar/tests.rs @@ -12,26 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::future::Future; -use std::pin::Pin; use std::sync::Arc; -use std::task::Context; use std::task::Poll; -use std::task::Waker; use tokio::task::JoinHandle; use crate::condvar::Condvar; use crate::mutex::Mutex; +use crate::poll_once; use crate::test_runtime; -fn poll_once(future: Pin<&mut F>) -> Poll -where - F: Future, -{ - future.poll(&mut Context::from_waker(Waker::noop())) -} - fn expect_ready(poll: Poll) -> T { match poll { Poll::Ready(value) => value, diff --git a/mea/src/internal/mod.rs b/mea/src/internal/mod.rs index 9aa7fab..ae0b68a 100644 --- a/mea/src/internal/mod.rs +++ b/mea/src/internal/mod.rs @@ -15,6 +15,9 @@ mod countdown; pub(crate) use countdown::*; +mod once_table; +pub(crate) use once_table::*; + mod mutex; pub(crate) use mutex::*; diff --git a/mea/src/internal/once_table.rs b/mea/src/internal/once_table.rs new file mode 100644 index 0000000..c7786e7 --- /dev/null +++ b/mea/src/internal/once_table.rs @@ -0,0 +1,164 @@ +// Copyright 2024 tison +// +// Licensed 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::borrow::Borrow; +use std::fmt; +use std::hash::BuildHasher; +use std::hash::Hash; +use std::sync::Arc; + +use hashbrown::HashTable; + +use crate::once::OnceCell; + +pub struct OnceTableEntry { + hash: u64, + key: K, + cell: OnceCell, +} + +impl OnceTableEntry { + pub fn initialized(&self) -> bool { + self.cell.initialized() + } + + pub fn get(&self) -> Option<&V> { + self.cell.get() + } + + pub async fn get_or_init(&self, init: F) -> &V + where + F: AsyncFnOnce() -> V, + { + self.cell.get_or_init(init).await + } + + pub async fn get_or_try_init(&self, init: F) -> Result<&V, E> + where + F: AsyncFnOnce() -> Result, + { + self.cell.get_or_try_init(init).await + } +} + +/// Shared keyed storage that lets once primitives clean up an exact entry without cloning its key. +pub struct OnceTable { + entries: HashTable>>, + hasher: S, +} + +impl fmt::Debug for OnceTable +where + K: fmt::Debug, + 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() + } +} + +impl OnceTable { + pub fn with_hasher(hasher: S) -> Self { + Self { + entries: HashTable::new(), + hasher, + } + } + + pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self { + Self { + entries: HashTable::with_capacity(capacity), + hasher, + } + } + + #[cfg(test)] + pub fn len(&self) -> usize { + self.entries.len() + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } +} + +impl OnceTable +where + K: Eq + Hash, + S: BuildHasher, +{ + pub fn get_or_insert(&mut 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(), + }) + }) + .into_mut() + } + + pub fn get(&self, key: &Q) -> Option<&Arc>> + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + self.entries.find(hash, |entry| entry.key.borrow() == key) + } + + pub fn remove(&mut self, key: &Q) -> Option>> + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + let hash = self.hasher.hash_one(key); + let entry = self + .entries + .find_entry(hash, |entry| entry.key.borrow() == key) + .ok()?; + let (entry, _) = entry.remove(); + Some(entry) + } + + /// 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 { + return; + }; + + drop(occupied.remove()); + } + + pub fn insert(&mut self, key: K, value: V) { + self.remove(&key); + + let hash = self.hasher.hash_one(&key); + let entry = Arc::new(OnceTableEntry { + hash, + key, + cell: OnceCell::from_value(value), + }); + self.entries.insert_unique(hash, entry, |entry| entry.hash); + } +} diff --git a/mea/src/lib.rs b/mea/src/lib.rs index bc257bf..28ddcd2 100644 --- a/mea/src/lib.rs +++ b/mea/src/lib.rs @@ -98,6 +98,14 @@ fn test_runtime() -> &'static tokio::runtime::Runtime { RT.get_or_init(|| Runtime::new().unwrap()) } +#[cfg(test)] +pub(crate) fn poll_once( + future: std::pin::Pin<&mut F>, +) -> std::task::Poll { + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + future.poll(&mut context) +} + #[cfg(test)] mod tests { use crate::admission::FairShare; diff --git a/mea/src/once/once_cell/mod.rs b/mea/src/once/once_cell/mod.rs index a104703..0cac8d3 100644 --- a/mea/src/once/once_cell/mod.rs +++ b/mea/src/once/once_cell/mod.rs @@ -86,12 +86,12 @@ impl OnceCell { } /// Returns whether the internal value is set. - fn initialized(&self) -> bool { + pub(crate) fn initialized(&self) -> bool { self.value_set.load(Ordering::Acquire) } /// Returns whether the internal value is set. - fn initialized_mut(&mut self) -> bool { + pub(crate) fn initialized_mut(&mut self) -> bool { *self.value_set.get_mut() } diff --git a/mea/src/once/once_map/mod.rs b/mea/src/once/once_map/mod.rs index e3f1179..1cc3678 100644 --- a/mea/src/once/once_map/mod.rs +++ b/mea/src/once/once_map/mod.rs @@ -13,14 +13,14 @@ // limitations under the License. use std::borrow::Borrow; -use std::collections::HashMap; use std::hash::BuildHasher; use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; use crate::internal::Mutex; -use crate::once::OnceCell; +use crate::internal::OnceTable; +use crate::internal::OnceTableEntry; #[cfg(test)] mod tests; @@ -31,14 +31,67 @@ mod tests; /// to wrap the `V` in an `Arc` to make cloning cheap. #[derive(Debug)] pub struct OnceMap { - map: Mutex>, S>>, + map: Mutex>, +} + +// Holds one call's entry so Drop can clean it up if the computation is abandoned. +struct ComputeCleanupGuard<'a, K, V, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + once_map: &'a OnceMap, + entry: Option>>, +} + +impl<'a, K, V, S> ComputeCleanupGuard<'a, K, V, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + fn new(once_map: &'a OnceMap, entry: Arc>) -> Self { + Self { + once_map, + entry: Some(entry), + } + } + + fn entry(&self) -> &Arc> { + self.entry.as_ref().unwrap() + } + + fn dismiss(mut self) { + drop(self.entry.take()); + } +} + +impl Drop for ComputeCleanupGuard<'_, K, V, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + fn drop(&mut self) { + let Some(entry) = self.entry.take() else { + 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); + } } impl Default for OnceMap where K: Eq + Hash, V: Clone, - S: BuildHasher + Clone + Default, + S: BuildHasher + Default, { fn default() -> Self { Self::with_hasher(S::default()) @@ -53,14 +106,17 @@ where /// Creates a new OnceMap with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(HashMap::new()), + map: Mutex::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(HashMap::with_capacity(capacity)), + map: Mutex::new(OnceTable::with_capacity_and_hasher( + capacity, + RandomState::new(), + )), } } } @@ -69,19 +125,19 @@ impl OnceMap where K: Eq + Hash, V: Clone, - S: BuildHasher + Clone, + S: BuildHasher, { /// Creates a new OnceMap with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: Mutex::new(HashMap::with_hasher(hasher)), + map: Mutex::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(HashMap::with_capacity_and_hasher(capacity, hasher)), + map: Mutex::new(OnceTable::with_capacity_and_hasher(capacity, hasher)), } } @@ -89,22 +145,26 @@ where /// /// If the value for the key is already being computed by another task, this task will wait for /// the computation to finish and return the result. + /// + /// If the computation is cancelled or panics, another caller waiting for the same key may retry + /// it. pub async fn compute(&self, key: K, func: F) -> V where F: AsyncFnOnce() -> V, { - // 1. Get or create the OnceCell. - let cell = { + let entry = { let mut map = self.map.lock(); - map.entry(key) - .or_insert_with(|| Arc::new(OnceCell::new())) - .clone() + let entry = map.get_or_insert(key); + if let Some(value) = entry.get() { + return value.clone(); + } + Arc::clone(entry) }; - // 2. Try to initialize the cell. - // OnceCell::get_or_init guarantees that only one task executes the closure. - let res = cell.get_or_init(func).await; - res.clone() + let guard = ComputeCleanupGuard::new(self, entry); + let result = guard.entry().get_or_init(func).await.clone(); + guard.dismiss(); + result } /// Compute the value for the given key if absent. @@ -112,24 +172,25 @@ where /// If the value for the key is already being computed by another task, this task will wait for /// the computation to finish and return the result. /// - /// If the computation fails, the error is returned and the value is not stored. Other tasks - /// waiting for the value will retry the computation. + /// If the computation returns an error, it is returned to that caller and the value is not + /// stored. After an error, cancellation, or panic, another caller may retry the computation. pub async fn try_compute(&self, key: K, func: F) -> Result where F: AsyncFnOnce() -> Result, { - // 1. Get or create the OnceCell. - let cell = { + let entry = { let mut map = self.map.lock(); - map.entry(key) - .or_insert_with(|| Arc::new(OnceCell::new())) - .clone() + let entry = map.get_or_insert(key); + if let Some(value) = entry.get() { + return Ok(value.clone()); + } + Arc::clone(entry) }; - // 2. Try to initialize the cell. - // OnceCell::get_or_try_init guarantees that only one task executes the closure. - let res = cell.get_or_try_init(func).await?; - Ok(res.clone()) + let guard = ComputeCleanupGuard::new(self, entry); + let result = guard.entry().get_or_try_init(func).await?.clone(); + guard.dismiss(); + Ok(result) } /// Get a clone of the value for the given key if exists. @@ -139,8 +200,8 @@ where Q: Hash + Eq + ?Sized, { let map = self.map.lock(); - let cell = map.get(key)?; - cell.get().cloned() + let entry = map.get(key)?; + entry.get().cloned() } /// Remove the given key from the map. @@ -168,24 +229,25 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - let cell = self.map.lock().remove(key)?; - cell.get().cloned() + let entry = self.map.lock().remove(key)?; + entry.get().cloned() } } impl FromIterator<(K, V)> for OnceMap where - K: Eq + Hash + Clone, + K: Eq + Hash, V: Clone, - S: Default + BuildHasher + Clone, + S: Default + BuildHasher, { fn from_iter>(iter: T) -> Self { + let mut map = OnceTable::with_hasher(S::default()); + for (key, value) in iter { + map.insert(key, value); + } + Self { - map: Mutex::new( - iter.into_iter() - .map(|(k, v)| (k, Arc::new(OnceCell::from_value(v)))) - .collect(), - ), + map: Mutex::new(map), } } } diff --git a/mea/src/once/once_map/tests.rs b/mea/src/once/once_map/tests.rs index b329cb2..ef230ab 100644 --- a/mea/src/once/once_map/tests.rs +++ b/mea/src/once/once_map/tests.rs @@ -14,12 +14,12 @@ use std::collections::hash_map::RandomState; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; use crate::once::OnceMap; +use crate::poll_once; #[test] fn test_default_and_constructors() { @@ -75,6 +75,7 @@ async fn test_try_compute() { // Fail first let res: Result = map.try_compute("key", async || Err("fail")).await; assert_eq!(res, Err("fail")); + assert!(map.map.lock().is_empty()); // Success then let res: Result = map.try_compute("key", async || Ok::(1)).await; @@ -86,58 +87,84 @@ async fn test_try_compute() { } #[tokio::test] -async fn test_try_compute_concurrent_failure_then_success() { - let map = Arc::new(OnceMap::new()); - let success = Arc::new(AtomicBool::new(false)); - let map_clone = map.clone(); - let success_clone = success.clone(); +async fn test_panicked_compute_removes_empty_entry() { + let map = Arc::new(OnceMap::<&str, i32>::new()); - // Spawn a task that fails - let t1 = tokio::spawn(async move { + let map_clone = map.clone(); + let task = tokio::spawn(async move { map_clone - .try_compute("key", async move || { - tokio::time::sleep(Duration::from_millis(50)).await; - Err::("fail") + .compute("key", async || { + panic!("oops"); }) .await }); - // Spawn a task that succeeds, but starts slightly later/runs concurrent - let map_clone2 = map.clone(); - let t2 = tokio::spawn(async move { - // Wait for t1 to start - tokio::time::sleep(Duration::from_millis(10)).await; - // This should block until t1 fails, then retry (conceptually) - map_clone2 - .try_compute("key", async move || { - success_clone.store(true, Ordering::SeqCst); - Ok::(1) + assert!(task.await.unwrap_err().is_panic()); + assert!(map.map.lock().is_empty()); +} + +#[tokio::test] +async fn test_cancelled_compute_removes_empty_entry() { + let map = Arc::new(OnceMap::<&str, i32>::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + + let map_clone = map.clone(); + let task = tokio::spawn(async move { + map_clone + .compute("key", async move || { + started_tx.send(()).unwrap(); + std::future::pending().await }) .await }); - let res1 = t1.await.unwrap(); - assert_eq!(res1, Err("fail")); + started_rx.await.unwrap(); + assert_eq!(map.map.lock().len(), 1); - let res2 = t2.await.unwrap(); - assert_eq!(res2, Ok(1)); - assert!(success.load(Ordering::SeqCst)); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(map.map.lock().is_empty()); } #[tokio::test] -async fn test_get_remove() { +async fn test_try_compute_concurrent_failure_then_success() { let map = OnceMap::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + + let first = map.try_compute("key", async move || { + release_rx.await.unwrap(); + Err::("fail") + }); + tokio::pin!(first); + assert!(poll_once(first.as_mut()).is_pending()); + + let retry = map.try_compute("key", async || Ok::(1)); + tokio::pin!(retry); + assert!(poll_once(retry.as_mut()).is_pending()); + + release_tx.send(()).unwrap(); + assert_eq!(first.await, Err("fail")); + + // The failed caller must not remove the cell while an existing waiter can still retry it. + assert_eq!(map.map.lock().len(), 1); + assert_eq!(retry.await, Ok(1)); + assert_eq!(map.get("key"), Some(1)); +} + +#[tokio::test] +async fn test_get_remove() { + let map = OnceMap::::new(); assert_eq!(map.get("key"), None); assert_eq!(map.remove("key"), None); - map.compute("key", async || 1).await; + map.compute("key".to_owned(), async || 1).await; assert_eq!(map.get("key"), Some(1)); let v = map.remove("key"); assert_eq!(v, Some(1)); assert_eq!(map.get("key"), None); - map.compute("key", async || 2).await; + map.compute("key".to_owned(), async || 2).await; map.discard("key"); assert_eq!(map.get("key"), None); } @@ -193,15 +220,20 @@ async fn test_get_while_computing() { #[tokio::test] async fn test_from_iter() { - let map: OnceMap<_, _> = vec![("a", 1), ("b", 2)].into_iter().collect(); - assert_eq!(map.get("a"), Some(1)); - assert_eq!(map.get("b"), Some(2)); - assert_eq!(map.get("c"), None); + #[derive(Hash, PartialEq, Eq)] + struct Key(&'static str); + + let map: OnceMap<_, _> = vec![(Key("a"), 1), (Key("b"), 2), (Key("a"), 3)] + .into_iter() + .collect(); + assert_eq!(map.get(&Key("a")), Some(3)); + assert_eq!(map.get(&Key("b")), Some(2)); + assert_eq!(map.get(&Key("c")), None); } #[tokio::test] async fn test_complex_key_value() { - #[derive(Hash, PartialEq, Eq, Clone, Debug)] + #[derive(Hash, PartialEq, Eq, Debug)] struct Key(i32); let map = OnceMap::new(); diff --git a/mea/src/singleflight/mod.rs b/mea/src/singleflight/mod.rs index 83b59bc..a7cad3a 100644 --- a/mea/src/singleflight/mod.rs +++ b/mea/src/singleflight/mod.rs @@ -15,14 +15,14 @@ //! Singleflight provides a duplicate function call suppression mechanism. use std::borrow::Borrow; -use std::collections::HashMap; use std::hash::BuildHasher; use std::hash::Hash; use std::hash::RandomState; use std::sync::Arc; use crate::internal::Mutex; -use crate::once::OnceCell; +use crate::internal::OnceTable; +use crate::internal::OnceTableEntry; #[cfg(test)] mod tests; @@ -31,14 +31,72 @@ mod tests; /// units of work can be executed with duplicate suppression. #[derive(Debug)] pub struct Group { - map: Mutex>, S>>, + map: Mutex>, +} + +// Holds one call's entry so Drop can clean it up if the work is abandoned. +struct WorkCleanupGuard<'a, K, V, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + group: &'a Group, + entry: Option>>, +} + +impl<'a, K, V, S> WorkCleanupGuard<'a, K, V, S> +where + K: Eq + Hash, + 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)) + }; + + Self { + group, + entry: Some(entry), + } + } + + fn entry(&self) -> &Arc> { + self.entry.as_ref().unwrap() + } + + fn dismiss(mut self) { + drop(self.entry.take()); + } +} + +impl Drop for WorkCleanupGuard<'_, K, V, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + fn drop(&mut self) { + let Some(entry) = self.entry.take() else { + 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); + } } impl Default for Group where - K: Eq + Hash + Clone, + K: Eq + Hash, V: Clone, - S: BuildHasher + Clone + Default, + S: BuildHasher + Default, { fn default() -> Self { Self::with_hasher(S::default()) @@ -47,27 +105,27 @@ where impl Group where - K: Eq + Hash + Clone, + K: Eq + Hash, V: Clone, { /// Creates a new Group with the default hasher. pub fn new() -> Self { Self { - map: Mutex::new(HashMap::new()), + map: Mutex::new(OnceTable::with_hasher(RandomState::new())), } } } impl Group where - K: Eq + Hash + Clone, + K: Eq + Hash, V: Clone, - S: BuildHasher + Clone, + S: BuildHasher, { /// Creates a new Group with the given hasher. pub fn with_hasher(hasher: S) -> Self { Self { - map: Mutex::new(HashMap::with_hasher(hasher)), + map: Mutex::new(OnceTable::with_hasher(hasher)), } } @@ -77,6 +135,9 @@ where /// If a duplicate comes in, the duplicate caller waits for the original to complete and /// receives the same results. /// + /// If the computation is cancelled or panics, another caller waiting for the same key may retry + /// it. + /// /// Once the function completes, the key, if not [`forgotten`], is removed from the group, /// allowing future calls with the same key to execute the function again. /// @@ -124,35 +185,18 @@ where where F: AsyncFnOnce() -> V, { - // 1. Get or create the OnceCell. - let cell = { - let mut map = self.map.lock(); - map.entry(key.clone()) - .or_insert_with(|| Arc::new(OnceCell::new())) - .clone() - }; - - // 2. Try to initialize the cell. - // OnceCell::get_or_init guarantees that only one task executes the closure. - let res = cell + let guard = WorkCleanupGuard::new(self, key); + let entry = guard.entry(); + let result = entry .get_or_init(async || { let result = func().await; - - // Cleanup: remove the key from the map. - // We must ensure we remove the entry corresponding to *this* cell. - let mut map = self.map.lock(); - if let Some(existing) = map.get(&key) { - // Check if the map still points to our cell. - if Arc::ptr_eq(&cell, existing) { - map.remove(&key); - } - } - + self.map.lock().remove_entry(entry); result }) - .await; - - res.clone() + .await + .clone(); + guard.dismiss(); + result } /// Executes and returns the results of the given function, making sure that only one execution @@ -161,8 +205,8 @@ where /// If a duplicate comes in, the duplicate caller waits for the original to complete and /// receives the same results. /// - /// If the computation fails, the error is returned for the caller. Other tasks waiting for the - /// result will retry the computation. + /// If the computation returns an error, it is returned to that caller. After an error, + /// cancellation, or panic, another caller may retry the computation. /// /// Once the function completes successfully, the key, if not [`forgotten`], is removed from /// the group, allowing future calls with the same key to execute the function again. @@ -205,35 +249,18 @@ where where F: AsyncFnOnce() -> Result, { - // 1. Get or create the OnceCell. - let cell = { - let mut map = self.map.lock(); - map.entry(key.clone()) - .or_insert_with(|| Arc::new(OnceCell::new())) - .clone() - }; - - // 2. Try to initialize the cell. - // OnceCell::get_or_try_init guarantees that only one task executes the closure. - let res = cell + let guard = WorkCleanupGuard::new(self, key); + let entry = guard.entry(); + let result = entry .get_or_try_init(async || { let result = func().await?; - - // Cleanup: remove the key from the map. - // We must ensure we remove the entry corresponding to *this* cell. - let mut map = self.map.lock(); - if let Some(existing) = map.get(&key) { - // Check if the map still points to our cell. - if Arc::ptr_eq(&cell, existing) { - map.remove(&key); - } - } - + self.map.lock().remove_entry(entry); Ok(result) }) - .await?; - - Ok(res.clone()) + .await? + .clone(); + guard.dismiss(); + Ok(result) } /// Forgets about the given key. diff --git a/mea/src/singleflight/tests.rs b/mea/src/singleflight/tests.rs index 1260231..65c0457 100644 --- a/mea/src/singleflight/tests.rs +++ b/mea/src/singleflight/tests.rs @@ -17,6 +17,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; +use crate::poll_once; use crate::singleflight::Group; #[tokio::test] @@ -26,6 +27,16 @@ async fn test_simple() { assert_eq!(res, "val"); } +#[tokio::test] +async fn test_non_clone_key() { + #[derive(Hash, PartialEq, Eq)] + struct Key(&'static str); + + let group = Group::new(); + let res = group.work(Key("key"), || async { "val" }).await; + assert_eq!(res, "val"); +} + #[tokio::test] async fn test_coalescing() { let group = Arc::new(Group::new()); @@ -93,7 +104,7 @@ async fn test_forget() { let g1 = group.clone(); let c1 = counter.clone(); let h1 = tokio::spawn(async move { - g1.work("key", || async move { + g1.work("key".to_owned(), || async move { tokio::time::sleep(Duration::from_millis(100)).await; c1.fetch_add(1, Ordering::SeqCst); "val1" @@ -103,12 +114,12 @@ async fn test_forget() { // Wait a bit to ensure the first call is established tokio::time::sleep(Duration::from_millis(10)).await; - group.forget(&"key"); + group.forget("key"); let g2 = group.clone(); let c2 = counter.clone(); let h2 = tokio::spawn(async move { - g2.work("key", || async move { + g2.work("key".to_owned(), || async move { tokio::time::sleep(Duration::from_millis(100)).await; c2.fetch_add(1, Ordering::SeqCst); "val2" @@ -137,12 +148,36 @@ async fn test_panic_safe() { // Wait for h1 to panic and exit let err = h1.await.unwrap_err(); assert!(err.is_panic()); + assert!(group.map.lock().is_empty()); // Next task should succeed (new attempt) let res = group.work("key", || async { "success".to_string() }).await; assert_eq!(res, "success"); } +#[tokio::test] +async fn test_cancelled_work_removes_empty_entry() { + let group = Arc::new(Group::<&str, &str>::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + + let group_clone = group.clone(); + let task = tokio::spawn(async move { + group_clone + .work("key", || async move { + started_tx.send(()).unwrap(); + std::future::pending().await + }) + .await + }); + + started_rx.await.unwrap(); + assert_eq!(group.map.lock().len(), 1); + + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(group.map.lock().is_empty()); +} + #[tokio::test] async fn test_try_work_simple() { let group = Group::new(); @@ -192,6 +227,7 @@ async fn test_try_work_failure() { .try_work("key", || async { Err::<&str, &str>("error") }) .await; assert_eq!(res, Err("error")); + assert!(group.map.lock().is_empty()); // Retry should work let res2 = group @@ -202,33 +238,25 @@ async fn test_try_work_failure() { #[tokio::test] async fn test_try_work_wait_and_retry() { - let group = Arc::new(Group::new()); - let counter = Arc::new(AtomicUsize::new(0)); + let group = Group::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); - let g1 = group.clone(); - let c1 = counter.clone(); - let h1 = tokio::spawn(async move { - g1.try_work("key", || async move { - c1.fetch_add(1, Ordering::SeqCst); - tokio::time::sleep(Duration::from_millis(100)).await; - Err::<&str, &str>("fail") - }) - .await + let first = group.try_work("key", || async move { + release_rx.await.unwrap(); + Err::<&str, &str>("fail") }); + tokio::pin!(first); + assert!(poll_once(first.as_mut()).is_pending()); - let g2 = group.clone(); - let c2 = counter.clone(); - let h2 = tokio::spawn(async move { - // Ensure h1 starts first - tokio::time::sleep(Duration::from_millis(10)).await; - g2.try_work("key", || async move { - c2.fetch_add(1, Ordering::SeqCst); - Ok::<&str, ()>("success") - }) - .await - }); + let retry = group.try_work("key", || async { Ok::<&str, &str>("success") }); + tokio::pin!(retry); + assert!(poll_once(retry.as_mut()).is_pending()); - assert_eq!(h1.await.unwrap(), Err("fail")); - assert_eq!(h2.await.unwrap(), Ok("success")); - assert_eq!(counter.load(Ordering::SeqCst), 2); + release_tx.send(()).unwrap(); + assert_eq!(first.await, Err("fail")); + + // The failed caller must not remove the cell while an existing waiter can still retry it. + assert_eq!(group.map.lock().len(), 1); + assert_eq!(retry.await, Ok("success")); + assert!(group.map.lock().is_empty()); }