mmtk/util/rust_util/
atomic_box.rs

1use std::sync::atomic::{AtomicPtr, Ordering};
2
3use bytemuck::Zeroable;
4
5/// A lazily initialized box.  Similar to an `Option<Box<T>>`, but can be initialized atomically.
6///
7/// It is designed for implementing shared data.  Therefore, methods with `&self`, namely
8/// [`OnceOptionBox::get`] and the [`OnceOptionBox::get_or_init`] methods, only return shared
9/// references to the content (`&T`).  The user should use types that support multi-threaded
10/// accesses, such as mutexes or atomic types, if the inner type is supposed to be modified
11/// concurrently.
12///
13/// Once initialized, this object will own its content.  The content is allocated in the heap, and
14/// will be dropped and deallocated when this instance is dropped.
15///
16/// # Comparison to existing data structures
17///
18/// [`std::sync::OnceLock`] also provides thread-safe lazily-initialized cells.  But as its name
19/// suggests, it uses locks for synchronization, whereas `OnceOptionBox` is lock-free.  `OnceLock`
20/// also has a field of [`std::sync::Once`] which increases the space overhead.  `OnceOptionBox`
21/// only has one atomic pointer field and is more suitable for large arrays of lazily initialized
22/// elements.
23pub struct OnceOptionBox<T> {
24    inner: AtomicPtr<T>,
25}
26
27impl<T> OnceOptionBox<T> {
28    /// Create an empty `OnceOptionBox` instance.
29    pub fn new() -> OnceOptionBox<T> {
30        Self {
31            inner: AtomicPtr::new(std::ptr::null_mut()),
32        }
33    }
34
35    /// Get a reference to the content of this box, or `None` if not yet initialized.
36    pub fn get(&self, order: Ordering) -> Option<&T> {
37        let ptr = self.inner.load(order);
38        unsafe { ptr.as_ref() }
39    }
40
41    /// Get a reference to the content of this box.  If not initialized, it will call `init` to
42    /// initialize this box.
43    ///
44    /// When multiple threads attempt to initialize this box concurrently, all threads may call
45    /// their supplied `init` closure, but only one thread will successfully initialize this box to
46    /// the return value of `init`.  Other threads will drop their return values of `init`.  All
47    /// callers will return the reference to the value created by the successful thread.
48    pub fn get_or_init(
49        &self,
50        order_load: Ordering,
51        order_store: Ordering,
52        init: impl FnOnce() -> T,
53    ) -> &T {
54        if let Some(get_result) = self.get(order_load) {
55            return get_result;
56        }
57
58        let new_inner = Box::into_raw(Box::new(init()));
59        let cas_result = self.inner.compare_exchange(
60            std::ptr::null_mut(),
61            new_inner,
62            order_store,
63            Ordering::Relaxed,
64        );
65        match cas_result {
66            Ok(old_inner) => {
67                debug_assert_eq!(old_inner, std::ptr::null_mut());
68                unsafe { new_inner.as_ref().unwrap() }
69            }
70            Err(old_inner) => {
71                drop(unsafe { Box::from_raw(new_inner) });
72                unsafe { old_inner.as_ref().unwrap() }
73            }
74        }
75    }
76}
77
78impl<T> Drop for OnceOptionBox<T> {
79    fn drop(&mut self) {
80        let ptr = *self.inner.get_mut();
81        if !ptr.is_null() {
82            drop(unsafe { Box::from_raw(ptr) });
83        }
84    }
85}
86
87unsafe impl<T> Zeroable for OnceOptionBox<T> {}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn construct() {
95        let oob = OnceOptionBox::<usize>::new();
96        let elem = oob.get(Ordering::Relaxed);
97        assert_eq!(elem, None);
98    }
99
100    #[test]
101    fn init() {
102        let oob = OnceOptionBox::<usize>::new();
103        let elem = oob.get_or_init(Ordering::Relaxed, Ordering::Relaxed, || 42);
104        assert_eq!(*elem, 42);
105    }
106
107    #[test]
108    fn reinit() {
109        let oob = OnceOptionBox::<usize>::new();
110        let elem = oob.get_or_init(Ordering::Relaxed, Ordering::Relaxed, || 42);
111        assert_eq!(*elem, 42);
112        let elem2 = oob.get_or_init(Ordering::Relaxed, Ordering::Relaxed, || 43);
113        assert_eq!(*elem2, 42);
114    }
115}
mmtk/util/rust_util/atomic_box.rs

mmtk/util/rust_util/
atomic_box.rs
