Add alloc_with_options (#1218)

This PR adds an alternative allocation functions `alloc_with_options`.
The attached `AllocationOptions` change the allocation behavior, such as
avoiding GCs and allowing overcommit.

Our current `alloc` functions assume they are GC safe points and will
trigger GCs internally. A runtime may have different assumptions. We see
GHC has `allocateMightFail`
(https://gitlab.haskell.org/ghc/ghc/-/blob/90746a591919fc51a0ec9dec58d8f1c8397040e3/rts/sm/Storage.c?page=2#L1089).
Also Julia assumes perm alloc will not trigger a GC
(mmtk/mmtk-julia#172).
Having a variant of `alloc` that is not GC safepoints could be generally
useful.

---------

Co-authored-by: Kunshan Wang <wks1986@gmail.com>
Co-authored-by: Ben Gamari <ben@smart-cactus.org>

Author: 3 people

src/memory_manager.rs

@@ -17,8 +17,9 @@ use crate::plan::AllocationSemantics;
 use crate::plan::{Mutator, MutatorContext};
 use crate::scheduler::WorkBucketStage;
 use crate::scheduler::{GCWork, GCWorker};
+use crate::util::alloc::allocator::AllocationOptions;
 use crate::util::alloc::allocators::AllocatorSelector;
-use crate::util::constants::{LOG_BYTES_IN_PAGE, MIN_OBJECT_SIZE};
+use crate::util::constants::LOG_BYTES_IN_PAGE;
 use crate::util::heap::layout::vm_layout::vm_layout;
 use crate::util::opaque_pointer::*;
 use crate::util::{Address, ObjectReference};
@@ -139,11 +140,28 @@ pub fn flush_mutator<VM: VMBinding>(mutator: &mut Mutator<VM>) {
     mutator.flush()
 }
 
-/// Allocate memory for an object. For performance reasons, a VM should
-/// implement the allocation fast-path on their side rather than just calling this function.
+/// Allocate memory for an object.
 ///
-/// If the VM provides a non-zero `offset` parameter, then the returned address will be
-/// such that the `RETURNED_ADDRESS + offset` is aligned to the `align` parameter.
+/// When the allocation is successful, it returns the starting address of the new object.  The
+/// memory range for the new object is `size` bytes starting from the returned address, and
+/// `RETURNED_ADDRESS + offset` is guaranteed to be aligned to the `align` parameter.  The returned
+/// address of a successful allocation will never be zero.
+///
+/// If MMTk fails to allocate memory, it will attempt a GC to free up some memory and retry the
+/// allocation.  After triggering GC, it will call [`crate::vm::Collection::block_for_gc`] to suspend
+/// the current thread that is allocating. Callers of `alloc` must be aware of this behavior.
+/// For example, JIT compilers that support
+/// precise stack scanning need to make the call site of `alloc` a GC-safe point by generating stack maps. See
+/// [`alloc_with_options`] if it is undesirable to trigger GC at this allocation site.
+///
+/// If MMTk has attempted at least one GC, and still cannot free up enough memory, it will call
+/// [`crate::vm::Collection::out_of_memory`] to inform the binding. The VM binding
+/// can implement that method to handle the out-of-memory event in a VM-specific way, including but
+/// not limited to throwing exceptions or errors. If [`crate::vm::Collection::out_of_memory`] returns
+/// normally without panicking or throwing exceptions, this function will return zero.
+///
+/// For performance reasons, a VM should implement the allocation fast-path on their side rather
+/// than just calling this function.
 ///
 /// Arguments:
 /// * `mutator`: The mutator to perform this allocation request.
@@ -158,24 +176,46 @@ pub fn alloc<VM: VMBinding>(
     offset: usize,
     semantics: AllocationSemantics,
 ) -> Address {
-    // MMTk has assumptions about minimal object size.
-    // We need to make sure that all allocations comply with the min object size.
-    // Ideally, we check the allocation size, and if it is smaller, we transparently allocate the min
-    // object size (the VM does not need to know this). However, for the VM bindings we support at the moment,
-    // their object sizes are all larger than MMTk's min object size, so we simply put an assertion here.
-    // If you plan to use MMTk with a VM with its object size smaller than MMTk's min object size, you should
-    // meet the min object size in the fastpath.
-    debug_assert!(size >= MIN_OBJECT_SIZE);
-    // Assert alignment
-    debug_assert!(align >= VM::MIN_ALIGNMENT);
-    debug_assert!(align <= VM::MAX_ALIGNMENT);
-    // Assert offset
-    debug_assert!(VM::USE_ALLOCATION_OFFSET || offset == 0);
+    #[cfg(debug_assertions)]
+    crate::util::alloc::allocator::assert_allocation_args::<VM>(size, align, offset);
 
     mutator.alloc(size, align, offset, semantics)
 }
 
-/// Invoke the allocation slow path. This is only intended for use when a binding implements the fastpath on
+/// Allocate memory for an object.
+///
+/// This allocation function allows alternation to the allocation behaviors, specified by the
+/// [`crate::util::alloc::AllocationOptions`]. For example, one can allow
+/// overcommit the memory to go beyond the heap size without triggering a GC. This function can be
+/// used in certain cases where the runtime needs a different allocation behavior other than
+/// what the default [`alloc`] provides.
+///
+/// Arguments:
+/// * `mutator`: The mutator to perform this allocation request.
+/// * `size`: The number of bytes required for the object.
+/// * `align`: Required alignment for the object.
+/// * `offset`: Offset associated with the alignment.
+/// * `semantics`: The allocation semantic required for the allocation.
+/// * `options`: the allocation options to change the default allocation behavior for this request.
+pub fn alloc_with_options<VM: VMBinding>(
+    mutator: &mut Mutator<VM>,
+    size: usize,
+    align: usize,
+    offset: usize,
+    semantics: AllocationSemantics,
+    options: crate::util::alloc::allocator::AllocationOptions,
+) -> Address {
+    #[cfg(debug_assertions)]
+    crate::util::alloc::allocator::assert_allocation_args::<VM>(size, align, offset);
+
+    mutator.alloc_with_options(size, align, offset, semantics, options)
+}
+
+/// Invoke the allocation slow path of [`alloc`].
+/// Like [`alloc`], this function may trigger GC and call [`crate::vm::Collection::block_for_gc`] or
+/// [`crate::vm::Collection::out_of_memory`].  The caller needs to be aware of that.
+///
+/// *Notes*: This is only intended for use when a binding implements the fastpath on
 /// the binding side. When the binding handles fast path allocation and the fast path fails, it can use this
 /// method for slow path allocation. Calling before exhausting fast path allocaiton buffer will lead to bad
 /// performance.
@@ -196,6 +236,34 @@ pub fn alloc_slow<VM: VMBinding>(
     mutator.alloc_slow(size, align, offset, semantics)
 }
 
+/// Invoke the allocation slow path of [`alloc_with_options`].
+///
+/// Like [`alloc_with_options`], This allocation function allows alternation to the allocation behaviors, specified by the
+/// [`crate::util::alloc::AllocationOptions`]. For example, one can allow
+/// overcommit the memory to go beyond the heap size without triggering a GC. This function can be
+/// used in certain cases where the runtime needs a different allocation behavior other than
+/// what the default [`alloc`] provides.
+///
+/// Like [`alloc_slow`], this function is also only intended for use when a binding implements the
+/// fastpath on the binding side.
+///
+/// Arguments:
+/// * `mutator`: The mutator to perform this allocation request.
+/// * `size`: The number of bytes required for the object.
+/// * `align`: Required alignment for the object.
+/// * `offset`: Offset associated with the alignment.
+/// * `semantics`: The allocation semantic required for the allocation.
+pub fn alloc_slow_with_options<VM: VMBinding>(
+    mutator: &mut Mutator<VM>,
+    size: usize,
+    align: usize,
+    offset: usize,
+    semantics: AllocationSemantics,
+    options: AllocationOptions,
+) -> Address {
+    mutator.alloc_slow_with_options(size, align, offset, semantics, options)
+}
+
 /// Perform post-allocation actions, usually initializing object metadata. For many allocators none are
 /// required. For performance reasons, a VM should implement the post alloc fast-path on their side
 /// rather than just calling this function.

src/plan/mutator_context.rs

@@ -4,6 +4,7 @@ use crate::plan::barriers::Barrier;
 use crate::plan::global::Plan;
 use crate::plan::AllocationSemantics;
 use crate::policy::space::Space;
+use crate::util::alloc::allocator::AllocationOptions;
 use crate::util::alloc::allocators::{AllocatorSelector, Allocators};
 use crate::util::alloc::Allocator;
 use crate::util::{Address, ObjectReference};
@@ -158,11 +159,30 @@ impl<VM: VMBinding> MutatorContext<VM> for Mutator<VM> {
         offset: usize,
         allocator: AllocationSemantics,
     ) -> Address {
-        unsafe {
+        let allocator = unsafe {
             self.allocators
                 .get_allocator_mut(self.config.allocator_mapping[allocator])
-        }
-        .alloc(size, align, offset)
+        };
+        // The value should be default/unset at the beginning of an allocation request.
+        debug_assert!(allocator.get_context().get_alloc_options().is_default());
+        allocator.alloc(size, align, offset)
+    }
+
+    fn alloc_with_options(
+        &mut self,
+        size: usize,
+        align: usize,
+        offset: usize,
+        allocator: AllocationSemantics,
+        options: AllocationOptions,
+    ) -> Address {
+        let allocator = unsafe {
+            self.allocators
+                .get_allocator_mut(self.config.allocator_mapping[allocator])
+        };
+        // The value should be default/unset at the beginning of an allocation request.
+        debug_assert!(allocator.get_context().get_alloc_options().is_default());
+        allocator.alloc_with_options(size, align, offset, options)
     }
 
     fn alloc_slow(
@@ -172,11 +192,30 @@ impl<VM: VMBinding> MutatorContext<VM> for Mutator<VM> {
         offset: usize,
         allocator: AllocationSemantics,
     ) -> Address {
-        unsafe {
+        let allocator = unsafe {
             self.allocators
                 .get_allocator_mut(self.config.allocator_mapping[allocator])
-        }
-        .alloc_slow(size, align, offset)
+        };
+        // The value should be default/unset at the beginning of an allocation request.
+        debug_assert!(allocator.get_context().get_alloc_options().is_default());
+        allocator.alloc_slow(size, align, offset)
+    }
+
+    fn alloc_slow_with_options(
+        &mut self,
+        size: usize,
+        align: usize,
+        offset: usize,
+        allocator: AllocationSemantics,
+        options: AllocationOptions,
+    ) -> Address {
+        let allocator = unsafe {
+            self.allocators
+                .get_allocator_mut(self.config.allocator_mapping[allocator])
+        };
+        // The value should be default/unset at the beginning of an allocation request.
+        debug_assert!(allocator.get_context().get_alloc_options().is_default());
+        allocator.alloc_slow_with_options(size, align, offset, options)
     }
 
     // Note that this method is slow, and we expect VM bindings that care about performance to implement allocation fastpath sequence in their bindings.
@@ -308,7 +347,7 @@ pub trait MutatorContext<VM: VMBinding>: Send + 'static {
     fn prepare(&mut self, tls: VMWorkerThread);
     /// Do the release work for this mutator.
     fn release(&mut self, tls: VMWorkerThread);
-    /// Allocate memory for an object.
+    /// Allocate memory for an object. This function will trigger a GC on failed allocation.
     ///
     /// Arguments:
     /// * `size`: the number of bytes required for the object.
@@ -322,7 +361,25 @@ pub trait MutatorContext<VM: VMBinding>: Send + 'static {
         offset: usize,
         allocator: AllocationSemantics,
     ) -> Address;
-    /// The slow path allocation. This is only useful when the binding
+    /// Allocate memory for an object with more options to control this allocation request, e.g. not triggering a GC on fail.
+    ///
+    /// Arguments:
+    /// * `size`: the number of bytes required for the object.
+    /// * `align`: required alignment for the object.
+    /// * `offset`: offset associated with the alignment. The result plus the offset will be aligned to the given alignment.
+    /// * `allocator`: the allocation semantic used for this object.
+    /// * `options`: the allocation options to change the default allocation behavior for this request.
+    fn alloc_with_options(
+        &mut self,
+        size: usize,
+        align: usize,
+        offset: usize,
+        allocator: AllocationSemantics,
+        options: AllocationOptions,
+    ) -> Address;
+    /// The slow path allocation for [`MutatorContext::alloc`]. This function will trigger a GC on failed allocation.
+    ///
+    ///  This is only useful when the binding
     /// implements the fast path allocation, and would like to explicitly
     /// call the slow path after the fast path allocation fails.
     fn alloc_slow(
@@ -332,6 +389,19 @@ pub trait MutatorContext<VM: VMBinding>: Send + 'static {
         offset: usize,
         allocator: AllocationSemantics,
     ) -> Address;
+    /// The slow path allocation for [`MutatorContext::alloc_with_options`].
+    ///
+    /// This is only useful when the binding
+    /// implements the fast path allocation, and would like to explicitly
+    /// call the slow path after the fast path allocation fails.
+    fn alloc_slow_with_options(
+        &mut self,
+        size: usize,
+        align: usize,
+        offset: usize,
+        allocator: AllocationSemantics,
+        options: AllocationOptions,
+    ) -> Address;
     /// Perform post-allocation actions.  For many allocators none are
     /// required.
     ///

src/policy/immix/immixspace.rs

@@ -7,6 +7,7 @@ use crate::policy::sft::GCWorkerMutRef;
 use crate::policy::sft::SFT;
 use crate::policy::sft_map::SFTMap;
 use crate::policy::space::{CommonSpace, Space};
+use crate::util::alloc::allocator::AllocationOptions;
 use crate::util::alloc::allocator::AllocatorContext;
 use crate::util::constants::LOG_BYTES_IN_PAGE;
 use crate::util::heap::chunk_map::*;
@@ -532,8 +533,13 @@ impl<VM: VMBinding> ImmixSpace<VM> {
     }
 
     /// Allocate a clean block.
-    pub fn get_clean_block(&self, tls: VMThread, copy: bool) -> Option<Block> {
-        let block_address = self.acquire(tls, Block::PAGES);
+    pub fn get_clean_block(
+        &self,
+        tls: VMThread,
+        copy: bool,
+        alloc_options: AllocationOptions,
+    ) -> Option<Block> {
+        let block_address = self.acquire(tls, Block::PAGES, alloc_options);
         if block_address.is_zero() {
             return None;
         }

src/policy/largeobjectspace.rs

@@ -5,6 +5,7 @@ use crate::plan::VectorObjectQueue;
 use crate::policy::sft::GCWorkerMutRef;
 use crate::policy::sft::SFT;
 use crate::policy::space::{CommonSpace, Space};
+use crate::util::alloc::allocator::AllocationOptions;
 use crate::util::constants::BYTES_IN_PAGE;
 use crate::util::heap::{FreeListPageResource, PageResource};
 use crate::util::metadata;
@@ -309,8 +310,13 @@ impl<VM: VMBinding> LargeObjectSpace<VM> {
     }
 
     /// Allocate an object
-    pub fn allocate_pages(&self, tls: VMThread, pages: usize) -> Address {
-        self.acquire(tls, pages)
+    pub fn allocate_pages(
+        &self,
+        tls: VMThread,
+        pages: usize,
+        alloc_options: AllocationOptions,
+    ) -> Address {
+        self.acquire(tls, pages, alloc_options)
     }
 
     /// Test if the object's mark bit is the same as the given value. If it is not the same,

src/policy/lockfreeimmortalspace.rs

@@ -8,6 +8,7 @@ use crate::policy::sft::SFT;
 use crate::policy::space::{CommonSpace, Space};
 use crate::util::address::Address;
 
+use crate::util::alloc::allocator::AllocationOptions;
 use crate::util::conversions;
 use crate::util::heap::gc_trigger::GCTrigger;
 use crate::util::heap::layout::vm_layout::vm_layout;
@@ -140,7 +141,7 @@ impl<VM: VMBinding> Space<VM> for LockFreeImmortalSpace<VM> {
         data_pages + meta_pages
     }
 
-    fn acquire(&self, _tls: VMThread, pages: usize) -> Address {
+    fn acquire(&self, _tls: VMThread, pages: usize, alloc_options: AllocationOptions) -> Address {
         trace!("LockFreeImmortalSpace::acquire");
         let bytes = conversions::pages_to_bytes(pages);
         let start = self
@@ -150,7 +151,11 @@ impl<VM: VMBinding> Space<VM> for LockFreeImmortalSpace<VM> {
             })
             .expect("update cursor failed");
         if start + bytes > self.limit {
-            panic!("OutOfMemory")
+            if alloc_options.on_fail.allow_oom_call() {
+                panic!("OutOfMemory");
+            } else {
+                return Address::ZERO;
+            }
         }
         if self.slow_path_zeroing {
             crate::util::memory::zero(start, bytes);

src/policy/marksweepspace/native_ms/global.rs

@@ -24,6 +24,7 @@ use crate::plan::ObjectQueue;
 use crate::plan::VectorObjectQueue;
 use crate::policy::sft::SFT;
 use crate::policy::space::{CommonSpace, Space};
+use crate::util::alloc::allocator::AllocationOptions;
 use crate::util::constants::LOG_BYTES_IN_PAGE;
 use crate::util::heap::chunk_map::*;
 use crate::util::linear_scan::Region;
@@ -462,7 +463,13 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
         crate::util::metadata::vo_bit::bzero_vo_bit(block.start(), Block::BYTES);
     }
 
-    pub fn acquire_block(&self, tls: VMThread, size: usize, align: usize) -> BlockAcquireResult {
+    pub fn acquire_block(
+        &self,
+        tls: VMThread,
+        size: usize,
+        align: usize,
+        alloc_options: AllocationOptions,
+    ) -> BlockAcquireResult {
         {
             let mut abandoned = self.abandoned.lock().unwrap();
             let bin = mi_bin::<VM>(size, align);
@@ -484,7 +491,7 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
             }
         }
 
-        let acquired = self.acquire(tls, Block::BYTES >> LOG_BYTES_IN_PAGE);
+        let acquired = self.acquire(tls, Block::BYTES >> LOG_BYTES_IN_PAGE, alloc_options);
         if acquired.is_zero() {
             BlockAcquireResult::Exhausted
         } else {