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}