greentic_flow/cache/
singleflight.rs

1use std::sync::Arc;
2
3use dashmap::DashMap;
4use tokio::sync::{Mutex, OwnedMutexGuard};
5
6use crate::cache::keys::ArtifactKey;
7
8#[derive(Clone, Debug, Default)]
9pub struct Singleflight {
10    locks: Arc<DashMap<ArtifactKey, Arc<Mutex<()>>>>,
11}
12
13impl Singleflight {
14    pub fn new() -> Self {
15        Self {
16            locks: Arc::new(DashMap::new()),
17        }
18    }
19
20    pub async fn acquire(&self, key: ArtifactKey) -> SingleflightGuard {
21        let lock = self
22            .locks
23            .entry(key.clone())
24            .or_insert_with(|| Arc::new(Mutex::new(())))
25            .clone();
26        let guard = lock.clone().lock_owned().await;
27        SingleflightGuard {
28            key,
29            lock,
30            guard: Some(guard),
31            locks: Arc::clone(&self.locks),
32        }
33    }
34}
35
36pub struct SingleflightGuard {
37    key: ArtifactKey,
38    lock: Arc<Mutex<()>>,
39    guard: Option<OwnedMutexGuard<()>>,
40    locks: Arc<DashMap<ArtifactKey, Arc<Mutex<()>>>>,
41}
42
43impl Drop for SingleflightGuard {
44    fn drop(&mut self) {
45        self.guard = None;
46        if Arc::strong_count(&self.lock) == 1 {
47            self.locks.remove(&self.key);
48        }
49    }
50}
greentic_flow/cache/singleflight.rs

greentic_flow/cache/
singleflight.rs
