commonware_utils/sync/
mod.rs

1//! Utilities for working with synchronization primitives.
2//!
3//! # Choosing A Lock
4//!
5//! Prefer blocking locks for shared data:
6//! - [Mutex]
7//! - [RwLock]
8//!
9//! Use async locks only when you must hold a lock guard across an `.await` point:
10//! - [AsyncMutex]
11//! - [AsyncRwLock]
12//! - [TracedAsyncMutex] and [TracedAsyncRwLock] for coordination-point locks whose
13//!   acquisition wait should be attributable in traces.
14//! - [UpgradableAsyncRwLock] when you need to read first and then conditionally upgrade to write
15//!   without allowing another writer to slip in between.
16//!
17//! Async locks are more expensive and should generally be reserved for coordination around
18//! asynchronous I/O resources. For plain in-memory data, blocking locks are usually the right
19//! default.
20//!
21//! Do not hold blocking lock guards across `.await`.
22//!
23//! Async lock guards may span `.await` when needed, but keep those critical sections as small as
24//! possible because long-held guards increase contention and deadlock risk.
25
26use core::ops::{Deref, DerefMut};
27pub use parking_lot::{
28    Condvar, MappedMutexGuard, Mutex, MutexGuard, Once, RwLock, RwLockReadGuard, RwLockWriteGuard,
29};
30pub use tokio::sync::{
31    Barrier, Mutex as AsyncMutex, MutexGuard as AsyncMutexGuard, Notify,
32    OwnedMutexGuard as OwnedAsyncMutexGuard, RwLock as AsyncRwLock,
33    RwLockReadGuard as AsyncRwLockReadGuard, RwLockWriteGuard as AsyncRwLockWriteGuard,
34};
35
36mod exclusive;
37pub use exclusive::Exclusive;
38
39/// A Tokio-based async mutex whose acquisitions are recorded as tracing spans.
40///
41/// Each lock is identified by a static name, recorded as the `lock` field on the
42/// `utils.mutex.lock` span so wait time is attributable to a specific lock.
43pub struct TracedAsyncMutex<T> {
44    name: &'static str,
45    inner: tokio::sync::Mutex<T>,
46}
47
48impl<T> TracedAsyncMutex<T> {
49    /// Create a new mutex wrapping `value`, identified by `name` in traces.
50    pub fn new(name: &'static str, value: T) -> Self {
51        Self {
52            name,
53            inner: tokio::sync::Mutex::new(value),
54        }
55    }
56
57    /// Acquire the mutex, recording lock-wait time.
58    #[tracing::instrument(name = "utils.mutex.lock", level = "info", skip_all, fields(lock = self.name))]
59    pub async fn lock(&self) -> AsyncMutexGuard<'_, T> {
60        self.inner.lock().await
61    }
62}
63
64/// A Tokio-based async rwlock whose acquisitions are recorded as tracing spans.
65///
66/// Each lock is identified by a static name, recorded as the `lock` field on the
67/// `utils.rwlock.read` and `utils.rwlock.write` spans so wait time is attributable to a
68/// specific lock.
69pub struct TracedAsyncRwLock<T> {
70    name: &'static str,
71    inner: tokio::sync::RwLock<T>,
72}
73
74impl<T> TracedAsyncRwLock<T> {
75    /// Create a new lock wrapping `value`, identified by `name` in traces.
76    pub fn new(name: &'static str, value: T) -> Self {
77        Self {
78            name,
79            inner: tokio::sync::RwLock::new(value),
80        }
81    }
82
83    /// Acquire a shared read guard, recording lock-wait time.
84    #[tracing::instrument(name = "utils.rwlock.read", level = "info", skip_all, fields(lock = self.name))]
85    pub async fn read(&self) -> AsyncRwLockReadGuard<'_, T> {
86        self.inner.read().await
87    }
88
89    /// Acquire an exclusive write guard, recording lock-wait time.
90    #[tracing::instrument(name = "utils.rwlock.write", level = "info", skip_all, fields(lock = self.name))]
91    pub async fn write(&self) -> AsyncRwLockWriteGuard<'_, T> {
92        self.inner.write().await
93    }
94}
95
96/// A Tokio-based async rwlock with an upgradable read mode.
97///
98/// All `write` and `upgradable_read` acquisitions take an internal async mutex ("gate") first.
99/// This ensures that upgrading from read to write does not allow another writer to slip in.
100pub struct UpgradableAsyncRwLock<T> {
101    rw: tokio::sync::RwLock<T>,
102    gate: tokio::sync::Mutex<()>,
103}
104
105impl<T> UpgradableAsyncRwLock<T> {
106    /// Create a new lock wrapping `value`.
107    pub fn new(value: T) -> Self {
108        Self {
109            rw: tokio::sync::RwLock::new(value),
110            gate: tokio::sync::Mutex::new(()),
111        }
112    }
113
114    /// Acquire a shared read guard.
115    pub async fn read(&self) -> tokio::sync::RwLockReadGuard<'_, T> {
116        self.rw.read().await
117    }
118
119    /// Acquire an exclusive write guard.
120    ///
121    /// Writers are serialized through the internal gate.
122    pub async fn write(&self) -> UpgradableAsyncRwLockWriteGuard<'_, T> {
123        let gate_guard = self.gate.lock().await;
124        let guard = self.rw.write().await;
125        UpgradableAsyncRwLockWriteGuard {
126            lock: self,
127            guard,
128            gate_guard,
129        }
130    }
131
132    /// Acquire an upgradable read guard.
133    ///
134    /// This allows shared reads, then a later [UpgradableAsyncRwLockUpgradableReadGuard::upgrade]
135    /// to exclusive write while holding the same gate token.
136    pub async fn upgradable_read(&self) -> UpgradableAsyncRwLockUpgradableReadGuard<'_, T> {
137        let gate_guard = self.gate.lock().await;
138        let guard = self.rw.read().await;
139        UpgradableAsyncRwLockUpgradableReadGuard {
140            lock: self,
141            guard,
142            gate_guard,
143        }
144    }
145
146    /// Consume the lock and return the wrapped value.
147    pub fn into_inner(self) -> T {
148        self.rw.into_inner()
149    }
150}
151
152/// Exclusive write guard for [UpgradableAsyncRwLock].
153pub struct UpgradableAsyncRwLockWriteGuard<'a, T> {
154    lock: &'a UpgradableAsyncRwLock<T>,
155    guard: tokio::sync::RwLockWriteGuard<'a, T>,
156    gate_guard: tokio::sync::MutexGuard<'a, ()>,
157}
158
159impl<'a, T> UpgradableAsyncRwLockWriteGuard<'a, T> {
160    /// Downgrade to an upgradable read guard while retaining the internal gate token.
161    pub fn downgrade_to_upgradable(self) -> UpgradableAsyncRwLockUpgradableReadGuard<'a, T> {
162        let Self {
163            lock,
164            guard,
165            gate_guard,
166        } = self;
167        let guard = tokio::sync::RwLockWriteGuard::downgrade(guard);
168        UpgradableAsyncRwLockUpgradableReadGuard {
169            lock,
170            guard,
171            gate_guard,
172        }
173    }
174}
175
176impl<T> Deref for UpgradableAsyncRwLockWriteGuard<'_, T> {
177    type Target = T;
178
179    fn deref(&self) -> &Self::Target {
180        &self.guard
181    }
182}
183
184impl<T> DerefMut for UpgradableAsyncRwLockWriteGuard<'_, T> {
185    fn deref_mut(&mut self) -> &mut Self::Target {
186        &mut self.guard
187    }
188}
189
190/// Upgradable read guard for [UpgradableAsyncRwLock].
191pub struct UpgradableAsyncRwLockUpgradableReadGuard<'a, T> {
192    lock: &'a UpgradableAsyncRwLock<T>,
193    guard: tokio::sync::RwLockReadGuard<'a, T>,
194    gate_guard: tokio::sync::MutexGuard<'a, ()>,
195}
196
197impl<'a, T> UpgradableAsyncRwLockUpgradableReadGuard<'a, T> {
198    /// Upgrade this guard to an exclusive writer.
199    pub async fn upgrade(self) -> UpgradableAsyncRwLockWriteGuard<'a, T> {
200        let Self {
201            lock,
202            guard,
203            gate_guard,
204        } = self;
205        drop(guard);
206        let guard = lock.rw.write().await;
207        UpgradableAsyncRwLockWriteGuard {
208            lock,
209            guard,
210            gate_guard,
211        }
212    }
213}
214
215impl<T> Deref for UpgradableAsyncRwLockUpgradableReadGuard<'_, T> {
216    type Target = T;
217
218    fn deref(&self) -> &Self::Target {
219        &self.guard
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::{AsyncRwLock, TracedAsyncMutex, TracedAsyncRwLock, UpgradableAsyncRwLock};
226    use futures::{FutureExt, pin_mut};
227
228    #[test]
229    fn test_traced_async_mutex() {
230        futures::executor::block_on(async {
231            let lock = TracedAsyncMutex::new("test", 100u64);
232
233            let mut guard = lock.lock().await;
234            *guard += 1;
235            drop(guard);
236
237            assert_eq!(*lock.lock().await, 101);
238        });
239    }
240
241    #[test]
242    fn test_traced_async_rwlock() {
243        futures::executor::block_on(async {
244            let lock = TracedAsyncRwLock::new("test", 100u64);
245
246            let r1 = lock.read().await;
247            let r2 = lock.read().await;
248            assert_eq!(*r1 + *r2, 200);
249
250            drop((r1, r2));
251            let mut writer = lock.write().await;
252            *writer += 1;
253
254            assert_eq!(*writer, 101);
255        });
256    }
257
258    #[test]
259    fn test_async_rwlock() {
260        futures::executor::block_on(async {
261            let lock = AsyncRwLock::new(100u64);
262
263            let r1 = lock.read().await;
264            let r2 = lock.read().await;
265            assert_eq!(*r1 + *r2, 200);
266
267            drop((r1, r2));
268            let mut writer = lock.write().await;
269            *writer += 1;
270
271            assert_eq!(*writer, 101);
272        });
273    }
274
275    #[test]
276    fn test_upgradable_read_blocks_write() {
277        futures::executor::block_on(async {
278            let lock = UpgradableAsyncRwLock::new(1u64);
279            let upgradable = lock.upgradable_read().await;
280
281            let write = lock.write();
282            pin_mut!(write);
283            assert!(write.as_mut().now_or_never().is_none());
284
285            drop(upgradable);
286
287            let mut write = write.await;
288            *write = 2;
289            drop(write);
290
291            assert_eq!(*lock.read().await, 2);
292        });
293    }
294
295    #[test]
296    fn test_read_allowed_during_upgradable_read() {
297        futures::executor::block_on(async {
298            let lock = UpgradableAsyncRwLock::new(5u64);
299            let upgradable = lock.upgradable_read().await;
300            let reader = lock.read().await;
301            assert_eq!(*upgradable, 5);
302            assert_eq!(*reader, 5);
303        });
304    }
305
306    #[test]
307    fn test_upgrade_prevents_writer_interleaving() {
308        futures::executor::block_on(async {
309            let lock = UpgradableAsyncRwLock::new(1u64);
310            let upgradable = lock.upgradable_read().await;
311
312            let writer = async {
313                let mut writer = lock.write().await;
314                let observed = *writer;
315                *writer = 7;
316                observed
317            };
318            pin_mut!(writer);
319            assert!(writer.as_mut().now_or_never().is_none());
320
321            let mut upgraded = upgradable.upgrade().await;
322            *upgraded = 5;
323            drop(upgraded);
324
325            assert_eq!(writer.await, 5);
326        });
327    }
328
329    #[test]
330    fn test_downgrade_to_upgradable() {
331        futures::executor::block_on(async {
332            let lock = UpgradableAsyncRwLock::new(10u64);
333            let mut writer = lock.write().await;
334            *writer = 11;
335
336            let upgradable = writer.downgrade_to_upgradable();
337            let writer = lock.write();
338            pin_mut!(writer);
339            assert!(writer.as_mut().now_or_never().is_none());
340            drop(upgradable);
341
342            let writer = writer.await;
343            assert_eq!(*writer, 11);
344        });
345    }
346}
commonware_utils/sync/mod.rs

commonware_utils/sync/
mod.rs
