mmap adjacent chunks (#738)

When quarantining memory, call mmap once for adjacent chunks with the same
state.  This reduces the number of mmap calls and reduces start-up time.

Author: wks

src/util/heap/layout/fragmented_mapper.rs

@@ -1,5 +1,6 @@
 use super::mmapper::MapState;
 use super::Mmapper;
+use crate::util::constants::BYTES_IN_PAGE;
 use crate::util::conversions;
 use crate::util::heap::layout::vm_layout_constants::*;
 use crate::util::Address;
@@ -81,10 +82,19 @@ impl Mmapper for FragmentedMapper {
     }
 
     fn quarantine_address_range(&self, mut start: Address, pages: usize) -> Result<()> {
+        debug_assert!(start.is_aligned_to(BYTES_IN_PAGE));
+
         let end = start + conversions::pages_to_bytes(pages);
+
+        // Each `MapState` entry governs a chunk.
+        // Align down to the chunk start because we only mmap multiples of whole chunks.
+        let mmap_start = conversions::mmap_chunk_align_down(start);
+
+        // We collect the chunk states from slabs to process them in bulk.
+        let mut state_slices = vec![];
+
         // Iterate over the slabs covered
         while start < end {
-            let base = Self::slab_align_down(start);
             let high = if end > Self::slab_limit(start) && !Self::slab_limit(start).is_zero() {
                 Self::slab_limit(start)
             } else {
@@ -96,19 +106,26 @@ impl Mmapper for FragmentedMapper {
             let end_chunk = Self::chunk_index(slab, conversions::mmap_chunk_align_up(high));
 
             let mapped = self.get_or_allocate_slab_table(start);
+            state_slices.push(&mapped[start_chunk..end_chunk]);
 
-            /* Iterate over the chunks within the slab */
-            for (chunk, entry) in mapped.iter().enumerate().take(end_chunk).skip(start_chunk) {
-                if matches!(entry.load(Ordering::Relaxed), MapState::Quarantined) {
-                    continue;
-                }
-
-                let mmap_start = Self::chunk_index_to_address(base, chunk);
-                let _guard = self.lock.lock().unwrap();
-                MapState::transition_to_quarantined(entry, mmap_start).unwrap();
-            }
             start = high;
         }
+
+        #[cfg(debug_assertions)]
+        {
+            // Check if the number of entries are normal.
+            let mmap_end = conversions::mmap_chunk_align_up(end);
+            let num_slices = state_slices.iter().map(|s| s.len()).sum::<usize>();
+
+            debug_assert_eq!(mmap_start + BYTES_IN_CHUNK * num_slices, mmap_end);
+        }
+
+        // Transition the chunks in bulk.
+        {
+            let _guard = self.lock.lock().unwrap();
+            MapState::bulk_transition_to_quarantined(state_slices.as_slice(), mmap_start)?;
+        }
+
         Ok(())
     }
 

src/util/heap/layout/mmapper.rs

@@ -1,5 +1,6 @@
 use crate::util::heap::layout::vm_layout_constants::*;
 use crate::util::memory::*;
+use crate::util::rust_util::rev_group::RevisitableGroupByForIterator;
 use crate::util::Address;
 use atomic::{Atomic, Ordering};
 use std::io::Result;
@@ -129,6 +130,69 @@ impl MapState {
         res
     }
 
+    /// Equivalent to calling `transition_to_quarantined` on each element of `states`, but faster.
+    /// The caller should hold a lock before invoking this method.
+    ///
+    /// The memory region to transition starts from `mmap_start`. The size is the chunk size
+    /// multiplied by the total number of `Atomic<MapState>` in the nested slice of slices.
+    ///
+    /// This function is introduced to speed up initialization of MMTk.  MMTk tries to quarantine
+    /// very large amount of memory during start-up.  If we quarantine one chunk at a time, it will
+    /// require thousands of `mmap` calls to cover gigabytes of quarantined memory, introducing a
+    /// noticeable delay.
+    ///
+    /// Arguments:
+    ///
+    /// * `state_slices`: A slice of slices. Each inner slice is a part of a `Slab`.
+    /// * `mmap_start`: The start of the region to transition.
+    pub(super) fn bulk_transition_to_quarantined(
+        state_slices: &[&[Atomic<MapState>]],
+        mmap_start: Address,
+    ) -> Result<()> {
+        trace!(
+            "Trying to bulk-quarantine {} - {}",
+            mmap_start,
+            mmap_start + MMAP_CHUNK_BYTES * state_slices.iter().map(|s| s.len()).sum::<usize>(),
+        );
+
+        let mut start_index = 0;
+
+        for group in state_slices
+            .iter()
+            .copied()
+            .flatten()
+            .revisitable_group_by(|s| s.load(Ordering::Relaxed))
+        {
+            let end_index = start_index + group.len;
+            let start_addr = mmap_start + MMAP_CHUNK_BYTES * start_index;
+            let end_addr = mmap_start + MMAP_CHUNK_BYTES * end_index;
+
+            match group.key {
+                MapState::Unmapped => {
+                    trace!("Trying to quarantine {} - {}", start_addr, end_addr);
+                    mmap_noreserve(start_addr, end_addr - start_addr)?;
+
+                    for state in group {
+                        state.store(MapState::Quarantined, Ordering::Relaxed);
+                    }
+                }
+                MapState::Quarantined => {
+                    trace!("Already quarantine {} - {}", start_addr, end_addr);
+                }
+                MapState::Mapped => {
+                    trace!("Already mapped {} - {}", start_addr, end_addr);
+                }
+                MapState::Protected => {
+                    panic!("Cannot quarantine protected memory")
+                }
+            }
+
+            start_index = end_index;
+        }
+
+        Ok(())
+    }
+
     /// Check the current MapState of the chunk, and transition the chunk to MapState::Protected.
     /// The caller should hold a lock before invoking this method.
     pub(super) fn transition_to_protected(

src/util/rust_util.rs renamed to src/util/rust_util/mod.rs

@@ -2,6 +2,9 @@
 //! functionalities that we may expect the Rust programming language and its standard libraries
 //! to provide.
 
+pub mod rev_group;
+pub mod zeroed_alloc;
+
 /// Const function for min value of two usize numbers.
 pub const fn min_of_usize(a: usize, b: usize) -> usize {
     if a > b {
@@ -127,53 +130,3 @@ mod initialize_once_tests {
         assert_eq!(INITIALIZE_COUNT.load(Ordering::SeqCst), 1);
     }
 }
-
-/// This module is for allocating large arrays or vectors with initial zero values.
-///
-/// Note: The standard library uses the `IsZero` trait to specialize the intialization of `Vec<T>`
-/// if the initial element values are zero.  Primitive type, such as `i8`, `usize`, `f32`, as well
-/// as types with known representations such as `Option<NonZeroUsize>` implement the `IsZero`
-/// trait.  However, it has several limitations.
-///
-/// 1.  Composite types, such as `SpaceDescriptor(usize)`, doesn't implement the `IsZero` trait,
-///     even if it has the `#[repr(transparent)]` annotation.
-/// 2.  The `IsZero` trait is private to the `std` module, and we cannot use it.
-///
-/// Therefore, `vec![0usize; 33554432]` takes only 4 **microseconds**, while
-/// `vec![SpaceDescriptor(0); 33554432]` will take 22 **milliseconds** to execute on some machine.
-/// If such an allocation happens during start-up, the delay will be noticeable to light-weight
-/// scripting languages, such as Ruby.
-///
-/// We implement our own fast allocation of large zeroed vectors in this module.  If one day Rust
-/// provides a standard way to optimize for zeroed allocation of vectors of composite types, we
-/// can switch to the standard mechanism.
-pub mod zeroed_alloc {
-
-    use std::alloc::{alloc_zeroed, Layout};
-
-    /// Allocate a `Vec<T>` of all-zero values.
-    ///
-    /// This intends to be a faster alternative to `vec![T(0), size]`.  It will allocate pre-zeroed
-    /// buffer, and not store zero values to its elements as part of initialization.
-    ///
-    /// It is useful when creating large (hundreds of megabytes) Vecs when the execution time is
-    /// critical (such as during start-up, where a 100ms delay is obvious to small applications.)
-    /// However, because of its unsafe nature, it should only be used when necessary.
-    ///
-    /// Arguments:
-    ///
-    /// -   `T`: The element type.
-    /// -   `size`: The length and capacity of the created vector.
-    ///
-    /// Returns the created vector.
-    ///
-    /// # Unsafe
-    ///
-    /// This function is unsafe.  It will not call any constructor of `T`.  The user must ensure
-    /// that a value with all bits being zero is meaningful for type `T`.
-    pub(crate) unsafe fn new_zeroed_vec<T>(size: usize) -> Vec<T> {
-        let layout = Layout::array::<T>(size).unwrap();
-        let ptr = alloc_zeroed(layout) as *mut T;
-        Vec::from_raw_parts(ptr, size, size)
-    }
-}