mmtk/util/rust_util/
mod.rs

1//! This module works around limitations of the Rust programming language, and provides missing
2//! functionalities that we may expect the Rust programming language and its standard libraries
3//! to provide.
4
5pub mod atomic_box;
6pub mod rev_group;
7pub mod zeroed_alloc;
8
9/// Const function for min value of two usize numbers.
10pub const fn min_of_usize(a: usize, b: usize) -> usize {
11    if a > b {
12        b
13    } else {
14        a
15    }
16}
17
18#[cfg(feature = "nightly")]
19pub use core::intrinsics::{likely, unlikely};
20
21// likely() and unlikely() compiler hints in stable Rust
22// [1]: https://github.com/rust-lang/hashbrown/blob/a41bd76de0a53838725b997c6085e024c47a0455/src/raw/mod.rs#L48-L70
23// [2]: https://users.rust-lang.org/t/compiler-hint-for-unlikely-likely-for-if-branches/62102/3
24#[cfg(not(feature = "nightly"))]
25#[inline]
26#[cold]
27fn cold() {}
28
29#[cfg(not(feature = "nightly"))]
30#[inline]
31pub fn likely(b: bool) -> bool {
32    if !b {
33        cold();
34    }
35    b
36}
37#[cfg(not(feature = "nightly"))]
38#[inline]
39pub fn unlikely(b: bool) -> bool {
40    if b {
41        cold();
42    }
43    b
44}
45
46use std::cell::UnsafeCell;
47use std::mem::MaybeUninit;
48use std::sync::Once;
49
50/// InitializeOnce creates an uninitialized value that needs to be manually initialized later. InitializeOnce
51/// guarantees the value is only initialized once. This type is used to allow more efficient reads.
52/// Unlike the `lazy_static!` which checks whether the static is initialized
53/// in every read, InitializeOnce has no extra check for reads.
54pub struct InitializeOnce<T: 'static> {
55    v: UnsafeCell<MaybeUninit<T>>,
56    /// This is used to guarantee `init_fn` is only called once.
57    once: Once,
58}
59
60impl<T> InitializeOnce<T> {
61    pub const fn new() -> Self {
62        InitializeOnce {
63            v: UnsafeCell::new(MaybeUninit::uninit()),
64            once: Once::new(),
65        }
66    }
67
68    /// Initialize the value. This should be called before ever using the struct.
69    /// If this method is called by multiple threads, the first thread will
70    /// initialize the value, and the other threads will be blocked until the
71    /// initialization is done (`Once` returns).
72    pub fn initialize_once(&self, init_fn: &'static dyn Fn() -> T) {
73        self.once.call_once(|| {
74            unsafe { &mut *self.v.get() }.write(init_fn());
75        });
76        debug_assert!(self.once.is_completed());
77    }
78
79    /// Get the value. This should only be used after initialize_once()
80    pub fn get_ref(&self) -> &T {
81        // We only assert in debug builds.
82        debug_assert!(self.once.is_completed());
83        unsafe { (*self.v.get()).assume_init_ref() }
84    }
85
86    /// Get a mutable reference to the value.
87    /// This is currently only used for SFTMap during plan creation (single threaded),
88    /// and before the plan creation is done, the binding cannot use MMTK at all.
89    ///
90    /// # Safety
91    /// The caller needs to make sure there is no race when mutating the value.
92    #[allow(clippy::mut_from_ref)]
93    pub unsafe fn get_mut(&self) -> &mut T {
94        // We only assert in debug builds.
95        debug_assert!(self.once.is_completed());
96        unsafe { (*self.v.get()).assume_init_mut() }
97    }
98}
99
100impl<T> std::ops::Deref for InitializeOnce<T> {
101    type Target = T;
102    fn deref(&self) -> &Self::Target {
103        self.get_ref()
104    }
105}
106
107unsafe impl<T> Sync for InitializeOnce<T> {}
108
109/// Create a formatted string that makes the best effort idenfying the current process and thread.
110pub fn debug_process_thread_id() -> String {
111    let pid = unsafe { libc::getpid() };
112    #[cfg(target_os = "linux")]
113    {
114        // `gettid()` is Linux-specific.
115        let tid = unsafe { libc::gettid() };
116        format!("PID: {}, TID: {}", pid, tid)
117    }
118    #[cfg(not(target_os = "linux"))]
119    {
120        // TODO: When we support other platforms, use platform-specific methods to get thread
121        // identifiers.
122        format!("PID: {}", pid)
123    }
124}
125
126#[cfg(test)]
127mod initialize_once_tests {
128    use super::*;
129
130    #[test]
131    fn test_threads_compete_initialize() {
132        use std::sync::atomic::AtomicUsize;
133        use std::sync::atomic::Ordering;
134        use std::thread;
135
136        // Create multiple threads to initialize the same `InitializeOnce` value
137        const N_THREADS: usize = 1000;
138        // The test value
139        static I: InitializeOnce<usize> = InitializeOnce::new();
140        // Count how many times the function is called
141        static INITIALIZE_COUNT: AtomicUsize = AtomicUsize::new(0);
142        // The function to create initial value
143        fn initialize_usize() -> usize {
144            INITIALIZE_COUNT.fetch_add(1, Ordering::SeqCst);
145            42
146        }
147
148        let mut threads = vec![];
149        for _ in 1..N_THREADS {
150            threads.push(thread::spawn(|| {
151                I.initialize_once(&initialize_usize);
152                // Every thread should see the value correctly initialized.
153                assert_eq!(*I, 42);
154            }));
155        }
156        threads.into_iter().for_each(|t| t.join().unwrap());
157
158        // The initialize_usize should only be called once
159        assert_eq!(INITIALIZE_COUNT.load(Ordering::SeqCst), 1);
160    }
161}
mmtk/util/rust_util/mod.rs

mmtk/util/rust_util/
mod.rs
