asyncband/singleflight/mod.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Coalesce concurrent work that uses the same key.
19//!
20//! While work for a key is in flight, later callers wait and clone its successful value. A
21//! completed value is not cached: the entry is removed when the work succeeds, so a non-overlapping
22//! later call executes its own function.
23//!
24//! If the active call is cancelled or panics, a waiting caller may run its own function. `try_work`
25//! also lets a waiter retry after the active function returns an error. Use `OnceMap` instead when
26//! successful values should remain cached by key.
27
28use std::borrow::Borrow;
29use std::fmt;
30use std::hash::BuildHasher;
31use std::hash::Hash;
32use std::hash::RandomState;
33use std::sync::Arc;
34
35use hashbrown::HashTable;
36
37use crate::internal::mutex::Mutex;
38use crate::once::OnceCell;
39
40#[cfg(test)]
41mod tests;
42
43struct Entry<K, V> {
44 hash: u64,
45 key: K,
46 cell: OnceCell<V>,
47}
48
49/// A namespace that coalesces concurrent work by key.
50///
51/// Different keys run independently. A group neither caches completed results nor limits overall
52/// concurrency.
53pub struct Group<K, V, S = RandomState> {
54 // This lock protects only entry lookup, insertion, and removal.
55 // User work is always run after releasing it.
56 entries: Mutex<HashTable<Arc<Entry<K, V>>>>,
57 hasher: S,
58}
59
60impl<K, V, S> fmt::Debug for Group<K, V, S> {
61 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62 let in_flight = self.entries.lock().len();
63 f.debug_struct("Group")
64 .field("in_flight", &in_flight)
65 .finish()
66 }
67}
68
69impl<K, V, S> Group<K, V, S>
70where
71 K: Eq + Hash,
72 S: BuildHasher,
73{
74 fn get_or_insert(&self, key: K) -> Arc<Entry<K, V>> {
75 let hash = self.hasher.hash_one(&key);
76 let mut entries = self.entries.lock();
77 entries
78 .entry(hash, |entry| entry.key.eq(&key), |entry| entry.hash)
79 .or_insert_with(|| {
80 Arc::new(Entry {
81 hash,
82 key,
83 cell: OnceCell::new(),
84 })
85 })
86 .into_mut()
87 .clone()
88 }
89
90 fn remove<Q>(&self, key: &Q)
91 where
92 K: Borrow<Q>,
93 Q: Eq + Hash + ?Sized,
94 {
95 let hash = self.hasher.hash_one(key);
96 let removed = {
97 let mut entries = self.entries.lock();
98 let Ok(occupied) = entries.find_entry(hash, |entry| entry.key.borrow() == key) else {
99 return;
100 };
101 occupied.remove().0
102 };
103 drop(removed);
104 }
105
106 fn remove_if_current(&self, entry: &Arc<Entry<K, V>>) {
107 let removed = {
108 let mut entries = self.entries.lock();
109 let Ok(occupied) = entries.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, entry))
110 else {
111 return;
112 };
113 occupied.remove().0
114 };
115 drop(removed);
116 }
117
118 fn cleanup_abandoned_entry(&self, entry: Arc<Entry<K, V>>) {
119 let mut entries = self.entries.lock();
120 let Ok(occupied) = entries.find_entry(entry.hash, |stored| Arc::ptr_eq(stored, &entry))
121 else {
122 // `forget` detached the entry. It may be the final owner, so release it after
123 // unlocking rather than running user destructors under the table lock.
124 drop(entries);
125 drop(entry);
126 return;
127 };
128
129 // With group ownership confirmed and the table locked against new callers, two owners
130 // means the group and this cleanup guard are the only remaining references.
131 if Arc::strong_count(&entry) == 2 && !entry.cell.initialized() {
132 let (stored, _) = occupied.remove();
133 drop(entries);
134 drop(entry);
135 drop(stored);
136 } else {
137 // A waiting cleanup must observe this call's reference being released while no new
138 // caller can clone the group's reference.
139 drop(entry);
140 }
141 }
142}
143
144// Holds one call's entry so Drop can clean it up if the work is abandoned.
145struct WorkCleanupGuard<'a, K, V, S>
146where
147 K: Eq + Hash,
148 S: BuildHasher,
149{
150 group: &'a Group<K, V, S>,
151 entry: Option<Arc<Entry<K, V>>>,
152}
153
154impl<'a, K, V, S> WorkCleanupGuard<'a, K, V, S>
155where
156 K: Eq + Hash,
157 S: BuildHasher,
158{
159 fn new(group: &'a Group<K, V, S>, key: K) -> Self {
160 let entry = group.get_or_insert(key);
161
162 Self {
163 group,
164 entry: Some(entry),
165 }
166 }
167
168 fn entry(&self) -> &Arc<Entry<K, V>> {
169 self.entry.as_ref().unwrap()
170 }
171
172 fn dismiss(mut self) {
173 drop(self.entry.take());
174 }
175}
176
177impl<K, V, S> Drop for WorkCleanupGuard<'_, K, V, S>
178where
179 K: Eq + Hash,
180 S: BuildHasher,
181{
182 fn drop(&mut self) {
183 let Some(entry) = self.entry.take() else {
184 return;
185 };
186
187 self.group.cleanup_abandoned_entry(entry);
188 }
189}
190
191impl<K, V, S> Default for Group<K, V, S>
192where
193 K: Eq + Hash,
194 V: Clone,
195 S: BuildHasher + Default,
196{
197 fn default() -> Self {
198 Self::with_hasher(S::default())
199 }
200}
201
202impl<K, V> Group<K, V, RandomState>
203where
204 K: Eq + Hash,
205 V: Clone,
206{
207 /// Creates a new Group with the default hasher.
208 pub fn new() -> Self {
209 Self::with_hasher(RandomState::new())
210 }
211}
212
213impl<K, V, S> Group<K, V, S>
214where
215 K: Eq + Hash,
216 V: Clone,
217 S: BuildHasher,
218{
219 /// Creates a new Group with the given hasher.
220 pub fn with_hasher(hasher: S) -> Self {
221 Self {
222 entries: Mutex::new(HashTable::new()),
223 hasher,
224 }
225 }
226
227 /// Executes and returns the results of the given function, making sure that only one execution
228 /// is in-flight for a given key at a time.
229 ///
230 /// If a duplicate comes in, the duplicate caller waits for the original to complete and
231 /// receives the same results.
232 ///
233 /// If the computation is cancelled or panics, another caller waiting for the same key may retry
234 /// it.
235 ///
236 /// Once the function completes, the key, if not [`forgotten`], is removed from the group,
237 /// allowing future calls with the same key to execute the function again.
238 ///
239 /// [`forgotten`]: Self::forget
240 ///
241 /// # Deadlocks
242 ///
243 /// The function must not recursively call `work` or `try_work` for the same current key because
244 /// it would wait for its own result. Work for other keys remains independent.
245 ///
246 /// # Examples
247 ///
248 /// ```
249 /// use std::sync::Arc;
250 /// use std::sync::atomic::AtomicUsize;
251 /// use std::sync::atomic::Ordering;
252 /// use std::time::Duration;
253 ///
254 /// use asyncband::singleflight::Group;
255 ///
256 /// # #[tokio::main]
257 /// # async fn main() {
258 /// let group = Group::new();
259 /// let counter = Arc::new(AtomicUsize::new(0));
260 ///
261 /// let c1 = counter.clone();
262 /// let fut1 = group.work("key", || async move {
263 /// c1.fetch_add(1, Ordering::SeqCst);
264 /// // simulate heavy work to avoid immediate completion
265 /// tokio::time::sleep(Duration::from_millis(100)).await;
266 /// "result"
267 /// });
268 ///
269 /// let c2 = counter.clone();
270 /// let fut2 = group.work("key", || async move {
271 /// c2.fetch_add(1, Ordering::SeqCst);
272 /// // simulate heavy work to avoid immediate completion
273 /// tokio::time::sleep(Duration::from_millis(100)).await;
274 /// "result"
275 /// });
276 ///
277 /// let (r1, r2) = tokio::join!(fut1, fut2);
278 ///
279 /// assert_eq!(r1, "result");
280 /// assert_eq!(r2, "result");
281 /// assert_eq!(counter.load(Ordering::SeqCst), 1);
282 /// # }
283 /// ```
284 pub async fn work<F>(&self, key: K, func: F) -> V
285 where
286 F: AsyncFnOnce() -> V,
287 {
288 let guard = WorkCleanupGuard::new(self, key);
289 let entry = guard.entry();
290 let result = entry
291 .cell
292 .get_or_init(async || {
293 let result = func().await;
294 self.remove_if_current(entry);
295 result
296 })
297 .await
298 .clone();
299 guard.dismiss();
300 result
301 }
302
303 /// Executes and returns the results of the given function, making sure that only one execution
304 /// is in-flight for a given key at a time.
305 ///
306 /// If a duplicate comes in, the duplicate caller waits for the original to complete and
307 /// receives the same results.
308 ///
309 /// If the computation returns an error, it is returned to that caller. After an error,
310 /// cancellation, or panic, another caller may retry the computation.
311 ///
312 /// Once the function completes successfully, the key, if not [`forgotten`], is removed from
313 /// the group, allowing future calls with the same key to execute the function again.
314 ///
315 /// [`forgotten`]: Self::forget
316 ///
317 /// # Deadlocks
318 ///
319 /// The function must not recursively call `work` or `try_work` for the same current key because
320 /// it would wait for its own result. Work for other keys remains independent.
321 ///
322 /// # Examples
323 ///
324 /// ```
325 /// use std::sync::Arc;
326 /// use std::sync::atomic::AtomicUsize;
327 /// use std::sync::atomic::Ordering;
328 /// use std::time::Duration;
329 ///
330 /// use asyncband::singleflight::Group;
331 ///
332 /// # #[tokio::main]
333 /// # async fn main() {
334 /// let group = Group::new();
335 ///
336 /// let fut1 = group.try_work("key", || async move {
337 /// // simulate heavy work to avoid immediate completion
338 /// tokio::time::sleep(Duration::from_millis(100)).await;
339 /// Err::<_, &'static str>("fut1")
340 /// });
341 ///
342 /// let fut2 = group.try_work("key", || async move {
343 /// // simulate heavy work to avoid immediate completion
344 /// tokio::time::sleep(Duration::from_millis(200)).await;
345 /// Ok::<_, &'static str>("fut2")
346 /// });
347 ///
348 /// let (r1, r2) = tokio::join!(fut1, fut2);
349 ///
350 /// assert_eq!(r1, Err("fut1"));
351 /// assert_eq!(r2, Ok("fut2"));
352 /// # }
353 /// ```
354 pub async fn try_work<E, F>(&self, key: K, func: F) -> Result<V, E>
355 where
356 F: AsyncFnOnce() -> Result<V, E>,
357 {
358 let guard = WorkCleanupGuard::new(self, key);
359 let entry = guard.entry();
360 let result = entry
361 .cell
362 .get_or_try_init(async || {
363 let result = func().await?;
364 self.remove_if_current(entry);
365 Ok(result)
366 })
367 .await?
368 .clone();
369 guard.dismiss();
370 Ok(result)
371 }
372
373 /// Forgets about the given key.
374 ///
375 /// Future calls to `work` for this key will call the function rather than waiting for an
376 /// earlier call to complete. Existing calls to `work` for this key are not affected.
377 pub fn forget<Q>(&self, key: &Q)
378 where
379 K: Borrow<Q>,
380 Q: Hash + Eq + ?Sized,
381 {
382 self.remove(key);
383 }
384}