during shutdown, only free the heap if we waited for foreign calls to exit

3 files changed, 8 insertions, 5 deletions

diff --git a/rts/RtsStartup.c b/rts/RtsStartup.c
index b0cddbd18f..d7a8d954a2 100644
--- a/rts/RtsStartup.c
+++ b/rts/RtsStartup.c
@@ -450,8 +450,11 @@ hs_exit_(rtsBool wait_foreign)
     /* free hash table storage */
     exitHashTable();
 
-    // Finally, free all our storage
-    freeStorage();
+    // Finally, free all our storage.  However, we only free the heap
+    // memory if we have waited for foreign calls to complete;
+    // otherwise a foreign call in progress may still be referencing
+    // heap memory (e.g. by being passed a ByteArray#).
+    freeStorage(wait_foreign);
 
 #if defined(DEBUG)
     /* and shut down the allocator debugging */
diff --git a/rts/sm/Storage.c b/rts/sm/Storage.c
index c9422e60b7..10a0a38d95 100644
--- a/rts/sm/Storage.c
+++ b/rts/sm/Storage.c
@@ -213,10 +213,10 @@ exitStorage (void)
 }
 
 void
-freeStorage (void)
+freeStorage (rtsBool free_heap)
 {
     stgFree(generations);
-    freeAllMBlocks();
+    if (free_heap) freeAllMBlocks();
 #if defined(THREADED_RTS)
     closeMutex(&sm_mutex);
 #endif
diff --git a/rts/sm/Storage.h b/rts/sm/Storage.h
index f6c50ad2f5..e541193573 100644
--- a/rts/sm/Storage.h
+++ b/rts/sm/Storage.h
@@ -19,7 +19,7 @@
 
 void initStorage(void);
 void exitStorage(void);
-void freeStorage(void);
+void freeStorage(rtsBool free_heap);
 
 /* -----------------------------------------------------------------------------
    Storage manager state