rts: Throw proper HeapOverflow exception on allocating large array

Test Plan: Validate, add tests

Reviewers: simonmar, austin, erikd

Reviewed By: simonmar

Subscribers: rwbarton, thomie

Differential Revision: https://phabricator.haskell.org/D4021

3 files changed, 79 insertions, 36 deletions

diff --git a/includes/rts/storage/GC.h b/includes/rts/storage/GC.h
index 387bd260d3..2aed7c57ee 100644
--- a/includes/rts/storage/GC.h
+++ b/includes/rts/storage/GC.h
@@ -184,8 +184,9 @@ extern generation * oldest_gen;
 
    -------------------------------------------------------------------------- */
 
-StgPtr  allocate        ( Capability *cap, W_ n );
-StgPtr  allocatePinned  ( Capability *cap, W_ n );
+StgPtr  allocate          ( Capability *cap, W_ n );
+StgPtr  allocateMightFail ( Capability *cap, W_ n );
+StgPtr  allocatePinned    ( Capability *cap, W_ n );
 
 /* memory allocator for executable memory */
 typedef void* AdjustorWritable;
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 4d54ecf6dc..b43dfbf554 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -62,7 +62,10 @@ stg_newByteArrayzh ( W_ n )
 
     payload_words = ROUNDUP_BYTES_TO_WDS(n);
     words = BYTES_TO_WDS(SIZEOF_StgArrBytes) + payload_words;
-    ("ptr" p) = ccall allocate(MyCapability() "ptr",words);
+    ("ptr" p) = ccall allocateMightFail(MyCapability() "ptr", words);
+    if (p == NULL) {
+        jump stg_raisezh(base_GHCziIOziException_heapOverflow_closure);
+    }
     TICK_ALLOC_PRIM(SIZEOF_StgArrBytes,WDS(payload_words),0);
     SET_HDR(p, stg_ARR_WORDS_info, CCCS);
     StgArrBytes_bytes(p) = n;
@@ -92,6 +95,9 @@ stg_newPinnedByteArrayzh ( W_ n )
     words = ROUNDUP_BYTES_TO_WDS(bytes);
 
     ("ptr" p) = ccall allocatePinned(MyCapability() "ptr", words);
+    if (p == NULL) {
+        jump stg_raisezh(base_GHCziIOziException_heapOverflow_closure);
+    }
     TICK_ALLOC_PRIM(SIZEOF_StgArrBytes,WDS(payload_words),0);
 
     /* Now we need to move p forward so that the payload is aligned
@@ -130,6 +136,9 @@ stg_newAlignedPinnedByteArrayzh ( W_ n, W_ alignment )
     words = ROUNDUP_BYTES_TO_WDS(bytes);
 
     ("ptr" p) = ccall allocatePinned(MyCapability() "ptr", words);
+    if (p == NULL) {
+        jump stg_raisezh(base_GHCziIOziException_heapOverflow_closure);
+    }
     TICK_ALLOC_PRIM(SIZEOF_StgArrBytes,WDS(payload_words),0);
 
     /* Now we need to move p forward so that the payload is aligned
@@ -240,7 +249,10 @@ stg_newArrayzh ( W_ n /* words */, gcptr init )
     // number of words.
     size = n + mutArrPtrsCardWords(n);
     words = BYTES_TO_WDS(SIZEOF_StgMutArrPtrs) + size;
-    ("ptr" arr) = ccall allocate(MyCapability() "ptr",words);
+    ("ptr" arr) = ccall allocateMightFail(MyCapability() "ptr",words);
+    if (arr == NULL) {
+        jump stg_raisezh(base_GHCziIOziException_heapOverflow_closure);
+    }
     TICK_ALLOC_PRIM(SIZEOF_StgMutArrPtrs, WDS(size), 0);
 
     SET_HDR(arr, stg_MUT_ARR_PTRS_DIRTY_info, CCCS);
@@ -366,7 +378,10 @@ stg_newArrayArrayzh ( W_ n /* words */ )
     // number of words.
     size = n + mutArrPtrsCardWords(n);
     words = BYTES_TO_WDS(SIZEOF_StgMutArrPtrs) + size;
-    ("ptr" arr) = ccall allocate(MyCapability() "ptr",words);
+    ("ptr" arr) = ccall allocateMightFail(MyCapability() "ptr",words);
+    if (arr == NULL) {
+        jump stg_raisezh(base_GHCziIOziException_heapOverflow_closure);
+    }
     TICK_ALLOC_PRIM(SIZEOF_StgMutArrPtrs, WDS(size), 0);
 
     SET_HDR(arr, stg_MUT_ARR_PTRS_DIRTY_info, W_[CCCS]);
@@ -398,7 +413,10 @@ stg_newSmallArrayzh ( W_ n /* words */, gcptr init )
     again: MAYBE_GC(again);
 
     words = BYTES_TO_WDS(SIZEOF_StgSmallMutArrPtrs) + n;
-    ("ptr" arr) = ccall allocate(MyCapability() "ptr",words);
+    ("ptr" arr) = ccall allocateMightFail(MyCapability() "ptr",words);
+    if (arr == NULL) {
+        jump stg_raisezh(base_GHCziIOziException_heapOverflow_closure);
+    }
     TICK_ALLOC_PRIM(SIZEOF_StgSmallMutArrPtrs, WDS(n), 0);
 
     SET_HDR(arr, stg_SMALL_MUT_ARR_PTRS_DIRTY_info, CCCS);
diff --git a/rts/sm/Storage.c b/rts/sm/Storage.c
index 6c5a73310c..e801c340f2 100644
--- a/rts/sm/Storage.c
+++ b/rts/sm/Storage.c
@@ -796,6 +796,20 @@ move_STACK (StgStack *src, StgStack *dest)
     dest->sp = (StgPtr)dest->sp + diff;
 }
 
+STATIC_INLINE void
+accountAllocation(Capability *cap, W_ n)
+{
+    TICK_ALLOC_HEAP_NOCTR(WDS(n));
+    CCS_ALLOC(cap->r.rCCCS,n);
+    if (cap->r.rCurrentTSO != NULL) {
+        // cap->r.rCurrentTSO->alloc_limit -= n*sizeof(W_)
+        ASSIGN_Int64((W_*)&(cap->r.rCurrentTSO->alloc_limit),
+                     (PK_Int64((W_*)&(cap->r.rCurrentTSO->alloc_limit))
+                      - n*sizeof(W_)));
+    }
+
+}
+
 /* -----------------------------------------------------------------------------
    StgPtr allocate (Capability *cap, W_ n)
 
@@ -812,21 +826,37 @@ move_STACK (StgStack *src, StgStack *dest)
    that operation fails, then the whole process will be killed.
    -------------------------------------------------------------------------- */
 
+/*
+ * Allocate some n words of heap memory; terminating
+ * on heap overflow
+ */
 StgPtr
 allocate (Capability *cap, W_ n)
 {
+    StgPtr p = allocateMightFail(cap, n);
+    if (p == NULL) {
+        reportHeapOverflow();
+        // heapOverflow() doesn't exit (see #2592), but we aren't
+        // in a position to do a clean shutdown here: we
+        // either have to allocate the memory or exit now.
+        // Allocating the memory would be bad, because the user
+        // has requested that we not exceed maxHeapSize, so we
+        // just exit.
+        stg_exit(EXIT_HEAPOVERFLOW);
+    }
+    return p;
+}
+
+/*
+ * Allocate some n words of heap memory; returning NULL
+ * on heap overflow
+ */
+StgPtr
+allocateMightFail (Capability *cap, W_ n)
+{
     bdescr *bd;
     StgPtr p;
 
-    TICK_ALLOC_HEAP_NOCTR(WDS(n));
-    CCS_ALLOC(cap->r.rCCCS,n);
-    if (cap->r.rCurrentTSO != NULL) {
-        // cap->r.rCurrentTSO->alloc_limit -= n*sizeof(W_)
-        ASSIGN_Int64((W_*)&(cap->r.rCurrentTSO->alloc_limit),
-                     (PK_Int64((W_*)&(cap->r.rCurrentTSO->alloc_limit))
-                      - n*sizeof(W_)));
-    }
-
     if (n >= LARGE_OBJECT_THRESHOLD/sizeof(W_)) {
         // The largest number of words such that
         // the computation of req_blocks will not overflow.
@@ -845,16 +875,12 @@ allocate (Capability *cap, W_ n)
             req_blocks >= HS_INT32_MAX)   // avoid overflow when
                                           // calling allocGroup() below
         {
-            reportHeapOverflow();
-            // heapOverflow() doesn't exit (see #2592), but we aren't
-            // in a position to do a clean shutdown here: we
-            // either have to allocate the memory or exit now.
-            // Allocating the memory would be bad, because the user
-            // has requested that we not exceed maxHeapSize, so we
-            // just exit.
-            stg_exit(EXIT_HEAPOVERFLOW);
+            return NULL;
         }
 
+        // Only credit allocation after we've passed the size check above
+        accountAllocation(cap, n);
+
         ACQUIRE_SM_LOCK
         bd = allocGroupOnNode(cap->node,req_blocks);
         dbl_link_onto(bd, &g0->large_objects);
@@ -870,6 +896,7 @@ allocate (Capability *cap, W_ n)
 
     /* small allocation (<LARGE_OBJECT_THRESHOLD) */
 
+    accountAllocation(cap, n);
     bd = cap->r.rCurrentAlloc;
     if (bd == NULL || bd->free + n > bd->start + BLOCK_SIZE_W) {
 
@@ -955,7 +982,8 @@ allocate (Capability *cap, W_ n)
    to pinned ByteArrays, not scavenging is ok.
 
    This function is called by newPinnedByteArray# which immediately
-   fills the allocated memory with a MutableByteArray#.
+   fills the allocated memory with a MutableByteArray#. Note that
+   this returns NULL on heap overflow.
    ------------------------------------------------------------------------- */
 
 StgPtr
@@ -967,20 +995,16 @@ allocatePinned (Capability *cap, W_ n)
     // If the request is for a large object, then allocate()
     // will give us a pinned object anyway.
     if (n >= LARGE_OBJECT_THRESHOLD/sizeof(W_)) {
-        p = allocate(cap, n);
-        Bdescr(p)->flags |= BF_PINNED;
-        return p;
-    }
-
-    TICK_ALLOC_HEAP_NOCTR(WDS(n));
-    CCS_ALLOC(cap->r.rCCCS,n);
-    if (cap->r.rCurrentTSO != NULL) {
-        // cap->r.rCurrentTSO->alloc_limit -= n*sizeof(W_);
-        ASSIGN_Int64((W_*)&(cap->r.rCurrentTSO->alloc_limit),
-                     (PK_Int64((W_*)&(cap->r.rCurrentTSO->alloc_limit))
-                      - n*sizeof(W_)));
+        p = allocateMightFail(cap, n);
+        if (p == NULL) {
+            return NULL;
+        } else {
+            Bdescr(p)->flags |= BF_PINNED;
+            return p;
+        }
     }
 
+    accountAllocation(cap, n);
     bd = cap->pinned_object_block;
 
     // If we don't have a block of pinned objects yet, or the current