nestrs_locks/
store.rs

1use async_trait::async_trait;
2use std::collections::HashMap;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5use tokio::sync::Mutex;
6use uuid::Uuid;
7
8#[async_trait]
9pub trait LockStore: Send + Sync {
10    async fn try_acquire(&self, name: &str, ttl: Duration) -> Option<LockGuard>;
11    async fn release(&self, name: &str, token: &str);
12}
13
14pub type SharedLockStore = Arc<dyn LockStore>;
15
16pub struct LockGuard {
17    store: SharedLockStore,
18    name: String,
19    token: String,
20    released: bool,
21}
22
23impl LockGuard {
24    pub fn name(&self) -> &str {
25        &self.name
26    }
27
28    pub async fn release(mut self) {
29        if !self.released {
30            self.store.release(&self.name, &self.token).await;
31            self.released = true;
32        }
33    }
34}
35
36impl Drop for LockGuard {
37    fn drop(&mut self) {
38        if self.released {
39            return;
40        }
41        let store = self.store.clone();
42        let name = self.name.clone();
43        let token = self.token.clone();
44        if let Ok(handle) = tokio::runtime::Handle::try_current() {
45            handle.spawn(async move {
46                store.release(&name, &token).await;
47            });
48        }
49    }
50}
51
52struct Entry {
53    token: String,
54    until: Instant,
55}
56
57pub struct InMemoryLockStore {
58    inner: Mutex<HashMap<String, Entry>>,
59}
60
61impl InMemoryLockStore {
62    pub fn new() -> Self {
63        Self {
64            inner: Mutex::new(HashMap::new()),
65        }
66    }
67}
68
69impl Default for InMemoryLockStore {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75#[async_trait]
76impl LockStore for InMemoryLockStore {
77    async fn try_acquire(&self, name: &str, ttl: Duration) -> Option<LockGuard> {
78        // Cannot construct LockGuard with Arc<Self> from &self without Arc wrapping.
79        // Use InMemoryLockStore::shared() factory from options instead.
80        let _ = (name, ttl);
81        None
82    }
83
84    async fn release(&self, name: &str, token: &str) {
85        let mut guard = self.inner.lock().await;
86        if let Some(e) = guard.get(name) {
87            if e.token == token {
88                guard.remove(name);
89            }
90        }
91    }
92}
93
94/// Arc-backed in-memory store (preferred).
95pub struct SharedInMemoryLockStore {
96    inner: Arc<Mutex<HashMap<String, Entry>>>,
97}
98
99impl SharedInMemoryLockStore {
100    pub fn new() -> Arc<Self> {
101        Arc::new(Self {
102            inner: Arc::new(Mutex::new(HashMap::new())),
103        })
104    }
105}
106
107#[async_trait]
108impl LockStore for SharedInMemoryLockStore {
109    async fn try_acquire(&self, name: &str, ttl: Duration) -> Option<LockGuard> {
110        let mut guard = self.inner.lock().await;
111        let now = Instant::now();
112        if let Some(e) = guard.get(name) {
113            if e.until > now {
114                return None;
115            }
116        }
117        let token = Uuid::new_v4().to_string();
118        guard.insert(
119            name.to_string(),
120            Entry {
121                token: token.clone(),
122                until: now + ttl,
123            },
124        );
125        let store: SharedLockStore = Arc::new(SharedInMemoryLockStore {
126            inner: self.inner.clone(),
127        });
128        Some(LockGuard {
129            store,
130            name: name.to_string(),
131            token,
132            released: false,
133        })
134    }
135
136    async fn release(&self, name: &str, token: &str) {
137        let mut guard = self.inner.lock().await;
138        if let Some(e) = guard.get(name) {
139            if e.token == token {
140                guard.remove(name);
141            }
142        }
143    }
144}
145
146#[cfg(feature = "cache-redis")]
147pub struct RedisLockStore {
148    client: redis::Client,
149    prefix: String,
150}
151
152#[cfg(feature = "cache-redis")]
153impl RedisLockStore {
154    pub fn new(url: &str) -> redis::RedisResult<Arc<Self>> {
155        Ok(Arc::new(Self {
156            client: redis::Client::open(url)?,
157            prefix: "nestrs:lock:".into(),
158        }))
159    }
160}
161
162#[cfg(feature = "cache-redis")]
163#[async_trait]
164impl LockStore for RedisLockStore {
165    async fn try_acquire(&self, name: &str, ttl: Duration) -> Option<LockGuard> {
166        let mut conn = self.client.get_multiplexed_async_connection().await.ok()?;
167        let key = format!("{}{}", self.prefix, name);
168        let token = Uuid::new_v4().to_string();
169        let set: bool = redis::cmd("SET")
170            .arg(&key)
171            .arg(&token)
172            .arg("NX")
173            .arg("PX")
174            .arg(ttl.as_millis() as u64)
175            .query_async(&mut conn)
176            .await
177            .ok()?;
178        if !set {
179            return None;
180        }
181        let store: SharedLockStore = Arc::new(RedisLockStore {
182            client: self.client.clone(),
183            prefix: self.prefix.clone(),
184        });
185        Some(LockGuard {
186            store,
187            name: name.to_string(),
188            token,
189            released: false,
190        })
191    }
192
193    async fn release(&self, name: &str, token: &str) {
194        use redis::AsyncCommands;
195        if let Ok(mut conn) = self.client.get_multiplexed_async_connection().await {
196            let key = format!("{}{}", self.prefix, name);
197            let current: Option<String> = conn.get(&key).await.ok().flatten();
198            if current.as_deref() == Some(token) {
199                let _: Result<(), _> = conn.del(key).await;
200            }
201        }
202    }
203}
nestrs_locks/store.rs

nestrs_locks/
store.rs
