Fix for CAF retention when dynamically loading & unloading code

In a situaion where we have some statically-linked code and we want to
load and unload a series of objects, we need the CAFs in the
statically-linked code to be retained indefinitely, while the CAFs in
the dynamically-linked code should be GC'd as normal, so that we can
detect when the code is unloadable.  This was wrong before - we GC'd
CAFs in the static code, leading to a crash in the rare case where we
use a CAF, GC it, and then load a new object that uses it again.

I also did some tidy up: RtsConfig now has a field keep_cafs to
indicate whether we want CAFs to be retained in static code.

6 files changed, 50 insertions, 13 deletions

diff --git a/includes/RtsAPI.h b/includes/RtsAPI.h
index 3b6de0f10c..4748060dee 100644
--- a/includes/RtsAPI.h
+++ b/includes/RtsAPI.h
@@ -73,6 +73,9 @@ typedef struct {
     // True if GHC was not passed -no-hs-main
     HsBool rts_hs_main;
 
+    // Whether to retain CAFs (default: false)
+    HsBool keep_cafs;
+
     // Called before processing command-line flags, so that default
     // settings for RtsFlags can be provided.
     void (* defaultsHook) (void);
diff --git a/includes/rts/storage/GC.h b/includes/rts/storage/GC.h
index 444ef8871b..f7da838739 100644
--- a/includes/rts/storage/GC.h
+++ b/includes/rts/storage/GC.h
@@ -200,11 +200,13 @@ void performMajorGC(void);
    The CAF table - used to let us revert CAFs in GHCi
    -------------------------------------------------------------------------- */
 
-StgInd *newCAF    (StgRegTable *reg, StgIndStatic *caf);
-StgInd *newDynCAF (StgRegTable *reg, StgIndStatic *caf);
+StgInd *newCAF         (StgRegTable *reg, StgIndStatic *caf);
+StgInd *newRetainedCAF (StgRegTable *reg, StgIndStatic *caf);
+StgInd *newGCdCAF      (StgRegTable *reg, StgIndStatic *caf);
 void revertCAFs (void);
 
 // Request that all CAFs are retained indefinitely.
+// (preferably use RtsConfig.keep_cafs instead)
 void setKeepCAFs (void);
 
 /* -----------------------------------------------------------------------------
diff --git a/rts/Linker.c b/rts/Linker.c
index 9d3ca1243e..bbf75bfd79 100644
--- a/rts/Linker.c
+++ b/rts/Linker.c
@@ -1512,7 +1512,7 @@ RTS_LIBFFI_SYMBOLS
 #define SymE_NeedsDataProto(vvv) SymE_HasDataProto(vvv)
 
 // SymI_HasProto_redirect allows us to redirect references to one symbol to
-// another symbol.  See newCAF/newDynCAF for an example.
+// another symbol.  See newCAF/newRetainedCAF/newGCdCAF for an example.
 #define SymI_HasProto_redirect(vvv,xxx)   \
     { MAYBE_LEADING_UNDERSCORE_STR(#vvv), \
       (void*)(&(xxx)) },
@@ -1692,10 +1692,10 @@ initLinker_ (int retain_cafs)
         barf("ghciInsertSymbolTable failed");
     }
 
-    // Redurect newCAF to newDynCAF if retain_cafs is true.
+    // Redurect newCAF to newRetainedCAF if retain_cafs is true.
     if (! ghciInsertSymbolTable(WSTR("(GHCi built-in symbols)"), symhash,
                                 MAYBE_LEADING_UNDERSCORE_STR("newCAF"),
-                                retain_cafs ? newDynCAF : newCAF,
+                                retain_cafs ? newRetainedCAF : newGCdCAF,
                                 HS_BOOL_FALSE, NULL)) {
         barf("ghciInsertSymbolTable failed");
     }
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index 94a6c0edeb..4e23eb8c20 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -62,6 +62,7 @@ const RtsConfig defaultRtsConfig  = {
     .rts_opts_suggestions = rtsTrue,
     .rts_opts = NULL,
     .rts_hs_main = rtsFalse,
+    .keep_cafs = rtsFalse,
     .defaultsHook = FlagDefaultsHook,
     .onExitHook = OnExitHook,
     .stackOverflowHook = StackOverflowHook,
diff --git a/rts/RtsStartup.c b/rts/RtsStartup.c
index c50bb07f75..f6544b6aba 100644
--- a/rts/RtsStartup.c
+++ b/rts/RtsStartup.c
@@ -138,12 +138,16 @@ hs_init_ghc(int *argc, char **argv[], RtsConfig rts_config)
     stat_startInit();
 
     /* Set the RTS flags to default values. */
-
     initRtsFlagsDefaults();
 
     /* Call the user hook to reset defaults, if present */
     rts_config.defaultsHook();
 
+    /* Whether to GC CAFs */
+    if (rts_config.keep_cafs) {
+        setKeepCAFs();
+    }
+
     /* Parse the flags, separating the RTS flags from the programs args */
     if (argc == NULL || argv == NULL) {
         // Use a default for argc & argv if either is not supplied
diff --git a/rts/sm/Storage.c b/rts/sm/Storage.c
index 85884fa12d..77796013ce 100644
--- a/rts/sm/Storage.c
+++ b/rts/sm/Storage.c
@@ -416,8 +416,8 @@ newCAF(StgRegTable *reg, StgIndStatic *caf)
     {
         // Note [dyn_caf_list]
         // If we are in GHCi _and_ we are using dynamic libraries,
-        // then we can't redirect newCAF calls to newDynCAF (see below),
-        // so we make newCAF behave almost like newDynCAF.
+        // then we can't redirect newCAF calls to newRetainedCAF (see below),
+        // so we make newCAF behave almost like newRetainedCAF.
         // The dynamic libraries might be used by both the interpreted
         // program and GHCi itself, so they must not be reverted.
         // This also means that in GHCi with dynamic libraries, CAFs are not
@@ -464,17 +464,17 @@ setKeepCAFs (void)
     keepCAFs = 1;
 }
 
-// An alternate version of newCaf which is used for dynamically loaded
+// An alternate version of newCAF which is used for dynamically loaded
 // object code in GHCi.  In this case we want to retain *all* CAFs in
 // the object code, because they might be demanded at any time from an
 // expression evaluated on the command line.
 // Also, GHCi might want to revert CAFs, so we add these to the
 // revertible_caf_list.
 //
-// The linker hackily arranges that references to newCaf from dynamic
-// code end up pointing to newDynCAF.
-StgInd *
-newDynCAF (StgRegTable *reg, StgIndStatic *caf)
+// The linker hackily arranges that references to newCAF from dynamic
+// code end up pointing to newRetainedCAF.
+//
+StgInd* newRetainedCAF (StgRegTable *reg, StgIndStatic *caf)
 {
     StgInd *bh;
 
@@ -491,6 +491,33 @@ newDynCAF (StgRegTable *reg, StgIndStatic *caf)
     return bh;
 }
 
+// If we are using loadObj/unloadObj in the linker, then we want to
+//
+//  - retain all CAFs in statically linked code (keepCAFs == rtsTrue),
+//    because we might link a new object that uses any of these CAFs.
+//
+//  - GC CAFs in dynamically-linked code, so that we can detect when
+//    a dynamically-linked object is unloadable.
+//
+// So for this case, we set keepCAFs to rtsTrue, and link newCAF to newGCdCAF
+// for dynamically-linked code.
+//
+StgInd* newGCdCAF (StgRegTable *reg, StgIndStatic *caf)
+{
+    StgInd *bh;
+
+    bh = lockCAF(reg, caf);
+    if (!bh) return NULL;
+
+    // Put this CAF on the mutable list for the old generation.
+    if (oldest_gen->no != 0) {
+        recordMutableCap((StgClosure*)caf,
+                         regTableToCapability(reg), oldest_gen->no);
+    }
+
+    return bh;
+}
+
 /* -----------------------------------------------------------------------------
    Nursery management.
    -------------------------------------------------------------------------- */