Fix #2592: do an orderly shutdown when the heap is exhausted

Really we should be raising an exception in this case, but that's
tricky (see comments).  At least now we shut down the runtime
correctly rather than just exiting.

6 files changed, 51 insertions, 10 deletions

diff --git a/includes/RtsAPI.h b/includes/RtsAPI.h
index 99aaa59ffb..41f0fc026c 100644
--- a/includes/RtsAPI.h
+++ b/includes/RtsAPI.h
@@ -22,7 +22,8 @@ typedef enum {
     NoStatus,    /* not finished yet */
     Success,	 /* completed successfully */
     Killed,	 /* uncaught exception */
-    Interrupted  /* stopped in response to a call to interruptStgRts */
+    Interrupted, /* stopped in response to a call to interruptStgRts */
+    HeapExhausted /* out of memory */
 } SchedulerStatus;
 
 typedef StgClosure *HaskellObj;
diff --git a/rts/Main.c b/rts/Main.c
index 434f79156b..aff3011e36 100644
--- a/rts/Main.c
+++ b/rts/Main.c
@@ -126,6 +126,9 @@ static void real_main(void)
       errorBelch("interrupted");
       exit_status = EXIT_INTERRUPTED;
       break;
+    case HeapExhausted:
+      exit_status = EXIT_HEAPOVERFLOW;
+      break;
     case Success:
       exit_status = EXIT_SUCCESS;
       break;
diff --git a/rts/RtsUtils.c b/rts/RtsUtils.c
index 0123531e27..7d6c4a557c 100644
--- a/rts/RtsUtils.c
+++ b/rts/RtsUtils.c
@@ -13,6 +13,7 @@
 #include "RtsFlags.h"
 #include "RtsUtils.h"
 #include "Ticky.h"
+#include "Schedule.h"
 
 #ifdef HAVE_TIME_H
 #include <time.h>
@@ -272,15 +273,14 @@ stackOverflow(void)
 void
 heapOverflow(void)
 {
-  /* don't fflush(stdout); WORKAROUND bug in Linux glibc */
-  OutOfHeapHook(0/*unknown request size*/, 
-		RtsFlags.GcFlags.maxHeapSize * BLOCK_SIZE);
-  
-#if defined(TICKY_TICKY)
-  if (RtsFlags.TickyFlags.showTickyStats) PrintTickyInfo();
-#endif
+    if (!heap_overflow)
+    {
+        /* don't fflush(stdout); WORKAROUND bug in Linux glibc */
+        OutOfHeapHook(0/*unknown request size*/,
+                      RtsFlags.GcFlags.maxHeapSize * BLOCK_SIZE);
 
-  stg_exit(EXIT_HEAPOVERFLOW);
+        heap_overflow = rtsTrue;
+    }
 }
 
 /* -----------------------------------------------------------------------------
diff --git a/rts/Schedule.c b/rts/Schedule.c
index 31a487515a..33715b1ecd 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -90,6 +90,12 @@ StgTSO *blackhole_queue = NULL;
  */
 rtsBool blackholes_need_checking = rtsFalse;
 
+/* 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.)
+ */
+rtsBool heap_overflow = rtsFalse;
+
 /* flag that tracks whether we have done any execution in this time slice.
  * LOCK: currently none, perhaps we should lock (but needs to be
  * updated in the fast path of the scheduler).
@@ -1436,7 +1442,11 @@ scheduleHandleThreadFinished (Capability *cap STG_UNUSED, Task *task, StgTSO *t)
 		  *(task->ret) = NULL;
 	      }
 	      if (sched_state >= SCHED_INTERRUPTING) {
-		  task->stat = Interrupted;
+                  if (heap_overflow) {
+                      task->stat = HeapExhausted;
+                  } else {
+                      task->stat = Interrupted;
+                  }
 	      } else {
 		  task->stat = Killed;
 	      }
@@ -1567,6 +1577,7 @@ scheduleDoGC (Capability *cap, Task *task USED_IF_THREADS, rtsBool force_major)
     
     IF_DEBUG(scheduler, printAllThreads());
 
+delete_threads_and_gc:
     /*
      * We now have all the capabilities; if we're in an interrupting
      * state, then we should take the opportunity to delete all the
@@ -1595,6 +1606,23 @@ scheduleDoGC (Capability *cap, Task *task USED_IF_THREADS, rtsBool force_major)
 	performHeapProfile = rtsFalse;
     }
 
+    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;
+    }
+
 #ifdef SPARKBALANCE
     /* JB 
        Once we are all together... this would be the place to balance all
diff --git a/rts/Schedule.h b/rts/Schedule.h
index c3334e6f69..d311801405 100644
--- a/rts/Schedule.h
+++ b/rts/Schedule.h
@@ -139,6 +139,8 @@ extern  StgTSO *RTS_VAR(sleeping_queue);
  */
 extern rtsBool blackholes_need_checking;
 
+extern rtsBool heap_overflow;
+
 #if defined(THREADED_RTS)
 extern Mutex RTS_VAR(sched_mutex);
 #endif
diff --git a/rts/sm/Storage.c b/rts/sm/Storage.c
index bf7c452d9b..6fa90cf8e9 100644
--- a/rts/sm/Storage.c
+++ b/rts/sm/Storage.c
@@ -617,6 +617,13 @@ allocateInGen (generation *g, lnat n)
         if (RtsFlags.GcFlags.maxHeapSize > 0 && 
             req_blocks >= RtsFlags.GcFlags.maxHeapSize) {
             heapOverflow();
+            // 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);
         }
 
 	bd = allocGroup(req_blocks);