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

commonware_utils/sync/
mod.rs
