egglog_concurrency/resettable_oncelock.rs
1//! An abstraction providing high scalability of shared reads and limited support for interior
2//! mutability and shared writes.
3//!
4//! The main thing we use this for is maintaining cached indexes in the `core-relations` crate.
5
6use std::{cell::UnsafeCell, sync::Once};
7
8/// `ResettableOnceLock` provides thread-safe access to a value of type `T` via a specific state
9/// machine.
10///
11/// * Fresh values of the lock start in a "to be updated" state.
12/// * When values are in "to be updated" then calls to `get` will return None.
13/// * In the to be updated state, the first call to `get_or_update` will run `update` on the object
14/// stored in the lock and transition the lock to the "updated" state.
15/// * Once in the updated state, calls to `get` will return a reference to the shared object,
16/// future calls to `get_or_update` will behave just like a call to `get` in this way.
17/// * A call to `reset` will transition the lock back to the "to be updated" state. Crucially, this
18/// requires mutable access to the lock object, making the patterns expressable via this lock less
19/// expressive than a regular mutex.
20pub struct ResettableOnceLock<T> {
21 data: UnsafeCell<T>,
22 update: Once,
23}
24
25unsafe impl<T: Send> Send for ResettableOnceLock<T> {}
26unsafe impl<T: Send> Sync for ResettableOnceLock<T> {}
27
28impl<T> ResettableOnceLock<T> {
29 pub fn new(elt: T) -> ResettableOnceLock<T> {
30 ResettableOnceLock {
31 data: UnsafeCell::new(elt),
32 update: Once::new(),
33 }
34 }
35
36 pub fn get(&self) -> Option<&T> {
37 if self.update.is_completed() {
38 // SAFETY: if `is_completed` is true, then no one will access this value mutably until
39 // a call to `reset`, which requires mutable access to the underlying lock.
40 unsafe { Some(&*self.data.get()) }
41 } else {
42 None
43 }
44 }
45
46 pub fn get_or_update(&self, update: impl FnOnce(&mut T)) -> &T {
47 if let Some(elt) = self.get() {
48 return elt;
49 }
50 self.update.call_once_force(|_| {
51 // SAFETY: We are calling this within `call_once_force` which guarantees a single
52 // thread will ever run at a given time.
53 //
54 // Further mutable accesses have to wait for a call to `reset`.
55 let mut_elt = unsafe { &mut *self.data.get() };
56
57 update(mut_elt);
58 });
59 // SAFETY: same as `get`. We know that is_completed() will return true if we get to this
60 // point.
61 unsafe { &*self.data.get() }
62 }
63
64 /// Reset the state of the lock to gate further accesses to the underlying value on another
65 /// call to `get_or_update`.
66 pub fn reset(&mut self) {
67 self.update = Once::new();
68 }
69}