Throw an exception on heap overflow

This changes heap overflow to throw a HeapOverflow exception instead of
killing the process.

Test Plan: GHC CI

Reviewers: simonmar, austin, hvr, erikd, bgamari

Reviewed By: simonmar, bgamari

Subscribers: thomie

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

GHC Trac Issues: #1791

6 files changed, 164 insertions, 25 deletions

diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index 1368082730..c9da13bafc 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -129,7 +129,7 @@ void initRtsFlagsDefaults(void)
         maxStkSize = 8 * 1024 * 1024;
 
     RtsFlags.GcFlags.statsFile          = NULL;
-    RtsFlags.GcFlags.giveStats          = NO_GC_STATS;
+    RtsFlags.GcFlags.giveStats          = COLLECT_GC_STATS;
 
     RtsFlags.GcFlags.maxStkSize         = maxStkSize / sizeof(W_);
     RtsFlags.GcFlags.initialStkSize     = 1024 / sizeof(W_);
@@ -141,6 +141,7 @@ void initRtsFlagsDefaults(void)
     RtsFlags.GcFlags.nurseryChunkSize   = 0;
     RtsFlags.GcFlags.minOldGenSize      = (1024 * 1024)       / BLOCK_SIZE;
     RtsFlags.GcFlags.maxHeapSize        = 0;    /* off by default */
+    RtsFlags.GcFlags.heapLimitGrace     = (1024 * 1024);
     RtsFlags.GcFlags.heapSizeSuggestion = 0;    /* none */
     RtsFlags.GcFlags.heapSizeSuggestionAuto = false;
     RtsFlags.GcFlags.pcFreeHeap         = 3;    /* 3% */
@@ -428,6 +429,11 @@ usage_text[] = {
 "  -xq       The allocation limit given to a thread after it receives",
 "            an AllocationLimitExceeded exception. (default: 100k)",
 "",
+"  -Mgrace=<n>",
+"            The amount of allocation after the program receives a",
+"            HeapOverflow exception before the exception is thrown again, if",
+"            the program is still exceeding the heap limit.",
+"",
 "RTS options may also be specified using the GHCRTS environment variable.",
 "",
 "Other RTS options may be available for programs compiled a different way.",
@@ -905,11 +911,16 @@ error = true;
 
               case 'M':
                   OPTION_UNSAFE;
-                  RtsFlags.GcFlags.maxHeapSize =
-                      decodeSize(rts_argv[arg], 2, BLOCK_SIZE, HS_WORD_MAX)
-                      / BLOCK_SIZE;
-                  /* user give size in *bytes* but "maxHeapSize" is in
-                   * *blocks* */
+                  if (0 == strncmp("grace=", rts_argv[arg] + 2, 6)) {
+                      RtsFlags.GcFlags.heapLimitGrace =
+                          decodeSize(rts_argv[arg], 8, BLOCK_SIZE, HS_WORD_MAX);
+                  } else {
+                      RtsFlags.GcFlags.maxHeapSize =
+                          decodeSize(rts_argv[arg], 2, BLOCK_SIZE, HS_WORD_MAX)
+                          / BLOCK_SIZE;
+                      // user give size in *bytes* but "maxHeapSize" is in
+                      // *blocks*
+                  }
                   break;
 
               case 'm':
diff --git a/rts/RtsStartup.c b/rts/RtsStartup.c
index 955ad13b4b..98c1dd290c 100644
--- a/rts/RtsStartup.c
+++ b/rts/RtsStartup.c
@@ -36,6 +36,7 @@
 #include "LinkerInternals.h"
 #include "LibdwPool.h"
 #include "sm/CNF.h"
+#include "TopHandler.h"
 
 #if defined(PROFILING)
 # include "ProfHeap.h"
@@ -242,6 +243,9 @@ hs_init_ghc(int *argc, char **argv[], RtsConfig rts_config)
     getStablePtr((StgPtr)runHandlersPtr_closure);
 #endif
 
+    // Initialize the top-level handler system
+    initTopHandler();
+
     /* initialise the shared Typeable store */
     initGlobalStore();
 
@@ -414,6 +418,9 @@ hs_exit_(bool wait_foreign)
     /* free the Static Pointer Table */
     exitStaticPtrTable();
 
+    /* remove the top-level handler */
+    exitTopHandler();
+
     /* free the stable pointer table */
     exitStableTables();
 
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 28479fb508..4c21c2ac52 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -10,6 +10,7 @@
 #include "RtsSymbols.h"
 
 #include "Rts.h"
+#include "TopHandler.h"
 #include "HsFFI.h"
 
 #include "sm/Storage.h"
@@ -747,6 +748,7 @@
       SymI_HasProto(rts_setThreadAllocationCounter)                     \
       SymI_HasProto(rts_enableThreadAllocationLimit)                    \
       SymI_HasProto(rts_disableThreadAllocationLimit)                   \
+      SymI_HasProto(rts_setMainThread)                                  \
       SymI_HasProto(setProgArgv)                                        \
       SymI_HasProto(startupHaskell)                                     \
       SymI_HasProto(shutdownHaskell)                                    \
diff --git a/rts/Schedule.c b/rts/Schedule.c
index 49687b577a..02d8137a2c 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -42,6 +42,7 @@
 #include "ThreadPaused.h"
 #include "Messages.h"
 #include "Stable.h"
+#include "TopHandler.h"
 
 #ifdef HAVE_SYS_TYPES_H
 #include <sys/types.h>
@@ -72,9 +73,14 @@ StgTSO *blocked_queue_tl = NULL;
 StgTSO *sleeping_queue = NULL;    // perhaps replace with a hash table?
 #endif
 
-/* Set to true when the latest garbage collection failed to reclaim
- * enough space, and the runtime should proceed to shut itself down in
- * an orderly fashion (emitting profiling info etc.)
+// Bytes allocated since the last time a HeapOverflow exception was thrown by
+// the RTS
+uint64_t allocated_bytes_at_heapoverflow = 0;
+
+/* Set to true when the latest garbage collection failed to reclaim enough
+ * space, and the runtime should proceed to shut itself down in an orderly
+ * fashion (emitting profiling info etc.), OR throw an exception to the main
+ * thread, if it is still alive.
  */
 bool heap_overflow = false;
 
@@ -1888,24 +1894,46 @@ delete_threads_and_gc:
         releaseGCThreads(cap, idle_cap);
     }
 #endif
-
     if (heap_overflow && sched_state < SCHED_INTERRUPTING) {
-        // GC set the heap_overflow flag, so we should proceed with
-        // an orderly shutdown now.  Ultimately we want the main
-        // thread to return to its caller with HeapExhausted, at which
-        // point the caller should call hs_exit().  The first step is
-        // to delete all the threads.
-        //
-        // Another way to do this would be to raise an exception in
-        // the main thread, which we really should do because it gives
-        // the program a chance to clean up.  But how do we find the
-        // main thread?  It should presumably be the same one that
-        // gets ^C exceptions, but that's all done on the Haskell side
-        // (GHC.TopHandler).
-        sched_state = SCHED_INTERRUPTING;
-        goto delete_threads_and_gc;
-    }
+        // GC set the heap_overflow flag.  We should throw an exception if we
+        // can, or shut down otherwise.
+
+        // Get the thread to which Ctrl-C is thrown
+        StgTSO *main_thread = getTopHandlerThread();
+        if (main_thread == NULL) {
+            // GC set the heap_overflow flag, and there is no main thread to
+            // throw an exception to, so we should proceed with an orderly
+            // shutdown now.  Ultimately we want the main thread to return to
+            // its caller with HeapExhausted, at which point the caller should
+            // call hs_exit().  The first step is to delete all the threads.
+            sched_state = SCHED_INTERRUPTING;
+            goto delete_threads_and_gc;
+        }
 
+        heap_overflow = false;
+        const uint64_t allocation_count = getAllocations();
+        if (RtsFlags.GcFlags.heapLimitGrace <
+              allocation_count - allocated_bytes_at_heapoverflow ||
+              allocated_bytes_at_heapoverflow == 0) {
+            allocated_bytes_at_heapoverflow = allocation_count;
+            // We used to simply exit, but throwing an exception gives the
+            // program a chance to clean up.  It also lets the exception be
+            // caught.
+
+            // FIXME this is not a good way to tell a program to release
+            // resources.  It is neither reliable (the RTS crashes if it fails
+            // to allocate memory from the OS) nor very usable (it is always
+            // thrown to the main thread, which might not be able to do anything
+            // useful with it).  We really should have a more general way to
+            // release resources in low-memory conditions.  Nevertheless, this
+            // is still a big improvement over just exiting.
+
+            // FIXME again: perhaps we should throw a synchronous exception
+            // instead an asynchronous one, or have a way for the program to
+            // register a handler to be called when heap overflow happens.
+            throwToSelf(cap, main_thread, heapOverflow_closure);
+        }
+    }
 #ifdef SPARKBALANCE
     /* JB
        Once we are all together... this would be the place to balance all
@@ -2608,6 +2636,8 @@ initScheduler(void)
 
   ACQUIRE_LOCK(&sched_mutex);
 
+  allocated_bytes_at_heapoverflow = 0;
+
   /* A capability holds the state a native thread needs in
    * order to execute STG code. At least one capability is
    * floating around (only THREADED_RTS builds have more than one).
diff --git a/rts/TopHandler.c b/rts/TopHandler.c
new file mode 100644
index 0000000000..ff53b32338
--- /dev/null
+++ b/rts/TopHandler.c
@@ -0,0 +1,62 @@
+#include "Rts.h"
+#include "Stable.h"
+#include "TopHandler.h"
+
+#ifdef THREADED_RTS
+static Mutex m; // Protects the operations on topHandlerPtr,
+                // which aren't atomic
+#endif
+static StgStablePtr topHandlerPtr;
+
+void rts_setMainThread(StgWeak *weak) {
+    ACQUIRE_LOCK(&m);
+    if (topHandlerPtr != NULL) {
+        freeStablePtr(topHandlerPtr); // OK to do under the lock
+    }
+    topHandlerPtr = getStablePtr((StgPtr)weak);
+    // referent is a Weak#
+    ASSERT(weak->header.info == &stg_WEAK_info);
+
+    // See Note [rts_setMainThread has an unsound type] in
+    // libraries/base/GHC/TopHandler.hs.
+    ASSERT(weak->key->header.info == &stg_TSO_info);
+
+    RELEASE_LOCK(&m);
+}
+
+StgTSO *getTopHandlerThread(void) {
+    ACQUIRE_LOCK(&m);
+    StgWeak *weak = (StgWeak*)deRefStablePtr(topHandlerPtr);
+    RELEASE_LOCK(&m);
+    const StgInfoTable *info = weak->header.info;
+    if (info == &stg_WEAK_info) {
+        StgClosure *key = ((StgWeak*)weak)->key;
+
+        // See Note [rts_setMainThread has an unsound type] in
+        // libraries/base/GHC/TopHandler.hs.
+        ASSERT(key->header.info == &stg_TSO_info);
+
+        return (StgTSO *)key;
+    } else if (info == &stg_DEAD_WEAK_info) {
+        return NULL;
+    } else {
+        barf("getTopHandlerThread: neither a WEAK nor a DEAD_WEAK: %p %p %d",
+             weak, info, info->type);
+        return NULL;
+    }
+}
+
+void initTopHandler(void) {
+#ifdef THREADED_RTS
+    initMutex(&m);
+#endif
+    topHandlerPtr = NULL;
+}
+
+void exitTopHandler(void) {
+    freeStablePtr(topHandlerPtr);
+    topHandlerPtr = NULL;
+#ifdef THREADED_RTS
+    closeMutex(&m);
+#endif
+}
diff --git a/rts/TopHandler.h b/rts/TopHandler.h
new file mode 100644
index 0000000000..fddebb05ef
--- /dev/null
+++ b/rts/TopHandler.h
@@ -0,0 +1,27 @@
+/* -----------------------------------------------------------------------------
+ *
+ * (c) The GHC Team, 2016
+ *
+ * Top-level handler support
+ *
+ * ---------------------------------------------------------------------------*/
+
+#include <BeginPrivate.h>
+#include <rts/Types.h>
+#include <rts/storage/Closures.h>
+#include <stg/Types.h>
+#include <rts/Stable.h>
+// Initialize the top handler subsystem
+void initTopHandler(void);
+
+// Exit the top handler subsystem
+void exitTopHandler(void);
+
+// Get the thread that handles ctrl-c, etc
+// Returns NULL if there is no such thread
+StgTSO *getTopHandlerThread(void);
+
+#include <EndPrivate.h>
+
+// Called from Haskell
+void rts_setMainThread(StgWeak *ptr);