Call Collection::out_of_memory if the allocation size is larger than max heap size (#896)

This closes #867.

Author: qinsoon

src/plan/global.rs

@@ -278,7 +278,10 @@ pub trait Plan: 'static + Sync + Downcast {
 
     /// Get the total number of pages for the heap.
     fn get_total_pages(&self) -> usize {
-        self.base().gc_trigger.policy.get_heap_size_in_pages()
+        self.base()
+            .gc_trigger
+            .policy
+            .get_current_heap_size_in_pages()
     }
 
     /// Get the number of pages that are still available for use. The available pages

src/policy/space.rs

@@ -13,7 +13,7 @@ use crate::util::heap::{PageResource, VMRequest};
 use crate::util::options::Options;
 use crate::vm::{ActivePlan, Collection};
 
-use crate::util::constants::LOG_BYTES_IN_MBYTE;
+use crate::util::constants::{LOG_BYTES_IN_MBYTE, LOG_BYTES_IN_PAGE};
 use crate::util::conversions;
 use crate::util::opaque_pointer::*;
 
@@ -46,8 +46,42 @@ pub trait Space<VM: VMBinding>: 'static + SFT + Sync + Downcast {
     /// Currently after we create a boxed plan, spaces in the plan have a non-moving address.
     fn initialize_sft(&self);
 
+    /// A check for the obvious out-of-memory case: if the requested size is larger than
+    /// the heap size, it is definitely an OOM. We would like to identify that, and
+    /// allows the binding to deal with OOM. Without this check, we will attempt
+    /// to allocate from the page resource. If the requested size is unrealistically large
+    /// (such as `usize::MAX`), it breaks the assumptions of our implementation of
+    /// page resource, vm map, etc. This check prevents that, and allows us to
+    /// handle the OOM case.
+    /// Each allocator that may request an arbitrary size should call this method before
+    /// acquring memory from the space. For example, bump pointer allocator and large object
+    /// allocator need to call this method. On the other hand, allocators that only allocate
+    /// memory in fixed size blocks do not need to call this method.
+    /// An allocator should call this method before doing any computation on the size to
+    /// avoid arithmatic overflow. If we have to do computation in the allocation fastpath and
+    /// overflow happens there, there is nothing we can do about it.
+    /// Return a boolean to indicate if we will be out of memory, determined by the check.
+    fn will_oom_on_acquire(&self, tls: VMThread, size: usize) -> bool {
+        let max_pages = self.get_gc_trigger().policy.get_max_heap_size_in_pages();
+        let requested_pages = size >> LOG_BYTES_IN_PAGE;
+        if requested_pages > max_pages {
+            VM::VMCollection::out_of_memory(
+                tls,
+                crate::util::alloc::AllocationError::HeapOutOfMemory,
+            );
+            return true;
+        }
+        false
+    }
+
     fn acquire(&self, tls: VMThread, pages: usize) -> Address {
         trace!("Space.acquire, tls={:?}", tls);
+
+        debug_assert!(
+            !self.will_oom_on_acquire(tls, pages << LOG_BYTES_IN_PAGE),
+            "The requested pages is larger than the max heap size. Is will_go_oom_on_acquire used before acquring memory?"
+        );
+
         // Should we poll to attempt to GC?
         // - If tls is collector, we cannot attempt a GC.
         // - If gc is disabled, we cannot attempt a GC.

src/util/alloc/bumpallocator.rs

@@ -158,6 +158,10 @@ impl<VM: VMBinding> BumpAllocator<VM> {
         offset: usize,
         stress_test: bool,
     ) -> Address {
+        if self.space.will_oom_on_acquire(self.tls, size) {
+            return Address::ZERO;
+        }
+
         let block_size = (size + BLOCK_MASK) & (!BLOCK_MASK);
         let acquired_start = self.space.acquire(self.tls, bytes_to_pages(block_size));
         if acquired_start.is_zero() {

src/util/alloc/large_object_allocator.rs

@@ -45,15 +45,13 @@ impl<VM: VMBinding> Allocator<VM> for LargeObjectAllocator<VM> {
     }
 
     fn alloc_slow_once(&mut self, size: usize, align: usize, _offset: usize) -> Address {
-        let header = 0; // HashSet is used instead of DoublyLinkedList
-        let maxbytes = allocator::get_maximum_aligned_size::<VM>(size + header, align);
-        let pages = crate::util::conversions::bytes_to_pages_up(maxbytes);
-        let sp = self.space.allocate_pages(self.tls, pages);
-        if sp.is_zero() {
-            sp
-        } else {
-            sp + header
+        if self.space.will_oom_on_acquire(self.tls, size) {
+            return Address::ZERO;
         }
+
+        let maxbytes = allocator::get_maximum_aligned_size::<VM>(size, align);
+        let pages = crate::util::conversions::bytes_to_pages_up(maxbytes);
+        self.space.allocate_pages(self.tls, pages)
     }
 }
 

src/util/heap/gc_trigger.rs

@@ -105,7 +105,9 @@ pub trait GCTriggerPolicy<VM: VMBinding>: Sync + Send {
     /// Is current heap full?
     fn is_heap_full(&self, plan: &'static dyn Plan<VM = VM>) -> bool;
     /// Return the current heap size (in pages)
-    fn get_heap_size_in_pages(&self) -> usize;
+    fn get_current_heap_size_in_pages(&self) -> usize;
+    /// Return the upper bound of heap size
+    fn get_max_heap_size_in_pages(&self) -> usize;
     /// Can the heap size grow?
     fn can_heap_size_grow(&self) -> bool;
 }
@@ -130,7 +132,11 @@ impl<VM: VMBinding> GCTriggerPolicy<VM> for FixedHeapSizeTrigger {
         plan.get_reserved_pages() > self.total_pages
     }
 
-    fn get_heap_size_in_pages(&self) -> usize {
+    fn get_current_heap_size_in_pages(&self) -> usize {
+        self.total_pages
+    }
+
+    fn get_max_heap_size_in_pages(&self) -> usize {
         self.total_pages
     }
 
@@ -394,10 +400,14 @@ impl<VM: VMBinding> GCTriggerPolicy<VM> for MemBalancerTrigger {
         plan.get_reserved_pages() > self.current_heap_pages.load(Ordering::Relaxed)
     }
 
-    fn get_heap_size_in_pages(&self) -> usize {
+    fn get_current_heap_size_in_pages(&self) -> usize {
         self.current_heap_pages.load(Ordering::Relaxed)
     }
 
+    fn get_max_heap_size_in_pages(&self) -> usize {
+        self.max_heap_pages
+    }
+
     fn can_heap_size_grow(&self) -> bool {
         self.current_heap_pages.load(Ordering::Relaxed) < self.max_heap_pages
     }

vmbindings/dummyvm/src/tests/allocate_with_initialize_collection.rs

@@ -12,7 +12,9 @@ pub fn allocate_with_initialize_collection() {
     mmtk_init(MB);
     mmtk_initialize_collection(VMThread::UNINITIALIZED);
     let handle = mmtk_bind_mutator(VMMutatorThread(VMThread::UNINITIALIZED));
-    // Attempt to allocate 2MB. This will trigger GC.
-    let addr = mmtk_alloc(handle, 2 * MB, 8, 0, AllocationSemantics::Default);
+    // Fill up the heap
+    let _ = mmtk_alloc(handle, MB, 8, 0, AllocationSemantics::Default);
+    // Attempt another allocation. This will trigger GC.
+    let addr = mmtk_alloc(handle, MB, 8, 0, AllocationSemantics::Default);
     assert!(!addr.is_zero());
 }

vmbindings/dummyvm/src/tests/allocate_without_initialize_collection.rs

@@ -11,7 +11,9 @@ pub fn allocate_without_initialize_collection() {
     // 1MB heap
     mmtk_init(MB);
     let handle = mmtk_bind_mutator(VMMutatorThread(VMThread::UNINITIALIZED));
-    // Attempt to allocate 2MB memory. This should trigger a GC, but as we never call initialize_collection(), we cannot do GC.
-    let addr = mmtk_alloc(handle, 2 * MB, 8, 0, AllocationSemantics::Default);
+    // Fill up the heap.
+    let _ = mmtk_alloc(handle, MB, 8, 0, AllocationSemantics::Default);
+    // Attempt another memory. This should trigger a GC, but as we never call initialize_collection(), we cannot do GC.
+    let addr = mmtk_alloc(handle, MB, 8, 0, AllocationSemantics::Default);
     assert!(!addr.is_zero());
 }

vmbindings/dummyvm/src/tests/fixtures/mod.rs

@@ -62,6 +62,27 @@ impl<T: FixtureContent> SerialFixture<T> {
         }
         func(c.as_ref().unwrap())
     }
+
+    pub fn with_fixture_expect_benign_panic<
+        F: Fn(&T) + std::panic::UnwindSafe + std::panic::RefUnwindSafe,
+    >(
+        &self,
+        func: F,
+    ) {
+        let res = {
+            // The lock will be dropped at the end of the block. So panic won't poison the lock.
+            let mut c = self.content.lock().unwrap();
+            if c.is_none() {
+                *c = Some(Box::new(T::create()));
+            }
+
+            std::panic::catch_unwind(|| func(c.as_ref().unwrap()))
+        };
+        // We do not hold the lock now. It is safe to resume now.
+        if let Err(e) = res {
+            std::panic::resume_unwind(e);
+        }
+    }
 }
 
 pub struct SingleObject {

vmbindings/dummyvm/src/tests/issue139.rs renamed to vmbindings/dummyvm/src/tests/issue139_allocate_unaligned_object_size.rs

@@ -0,0 +1,82 @@
+// GITHUB-CI: MMTK_PLAN=all
+
+use crate::api;
+use crate::tests::fixtures::{MutatorFixture, SerialFixture};
+use mmtk::plan::AllocationSemantics;
+
+lazy_static! {
+    static ref MUTATOR: SerialFixture<MutatorFixture> = SerialFixture::new();
+}
+
+#[test]
+#[should_panic(expected = "Out of memory with HeapOutOfMemory!")]
+pub fn allocate_max_size_object() {
+    let (size, align) = (usize::MAX, 8);
+
+    MUTATOR.with_fixture_expect_benign_panic(|fixture| {
+        api::mmtk_alloc(
+            fixture.mutator,
+            size,
+            align,
+            0,
+            AllocationSemantics::Default,
+        );
+    })
+}
+
+#[test]
+// This test panics with 'attempt to add with overflow', as we do computation with the size
+// in the fastpath. I don't think we want to do any extra check in the fastpath. There is
+// nothing we can do with it without sacrificing performance.
+#[should_panic(expected = "Out of memory with HeapOutOfMemory!")]
+#[ignore]
+pub fn allocate_max_size_object_after_succeed() {
+    MUTATOR.with_fixture_expect_benign_panic(|fixture| {
+        // Allocate something so we have a thread local allocation buffer
+        api::mmtk_alloc(fixture.mutator, 8, 8, 0, AllocationSemantics::Default);
+        // Allocate an unrealistically large object
+        api::mmtk_alloc(
+            fixture.mutator,
+            usize::MAX,
+            8,
+            0,
+            AllocationSemantics::Default,
+        );
+    })
+}
+
+#[test]
+#[should_panic(expected = "Out of memory with HeapOutOfMemory!")]
+pub fn allocate_unrealistically_large_object() {
+    const CHUNK: usize = 4 * 1024 * 1024; // 4MB
+                                          // Leave some room, so we won't have arithmetic overflow when we compute size and do alignment.
+    let (size, align) = (
+        mmtk::util::conversions::raw_align_down(usize::MAX - CHUNK, 4096),
+        8,
+    );
+
+    MUTATOR.with_fixture_expect_benign_panic(|fixture| {
+        api::mmtk_alloc(
+            fixture.mutator,
+            size,
+            align,
+            0,
+            AllocationSemantics::Default,
+        );
+    })
+}
+
+#[test]
+#[should_panic(expected = "Out of memory with HeapOutOfMemory!")]
+pub fn allocate_more_than_heap_size() {
+    // The heap has 1 MB. Allocating with 2MB will cause OOM.
+    MUTATOR.with_fixture_expect_benign_panic(|fixture| {
+        api::mmtk_alloc(
+            fixture.mutator,
+            2 * 1024 * 1024,
+            8,
+            0,
+            AllocationSemantics::Default,
+        );
+    })
+}

vmbindings/dummyvm/src/tests/issue867_allocate_unrealistically_large_object.rs

@@ -18,7 +18,8 @@ mod fixtures;
 mod handle_mmap_conflict;
 mod handle_mmap_oom;
 mod is_in_mmtk_spaces;
-mod issue139;
+mod issue139_allocate_unaligned_object_size;
+mod issue867_allocate_unrealistically_large_object;
 #[cfg(not(feature = "malloc_counted_size"))]
 mod malloc_api;
 #[cfg(feature = "malloc_counted_size")]