Introduce stack snapshotting / cloning (#18741)

Add `StackSnapshot#` primitive type that represents a cloned stack (StgStack).
The cloning interface consists of two functions, that clone either the treads
own stack (cloneMyStack) or another threads stack (cloneThreadStack).

The stack snapshot is offline/cold, i.e. it isn't evaluated any further. This is
useful for analyses as it prevents concurrent modifications.

For technical details, please see Note [Stack Cloning].

Co-authored-by: Ben Gamari <bgamari.foss@gmail.com>
Co-authored-by: Matthew Pickering <matthewtpickering@gmail.com>

10 files changed, 165 insertions, 1 deletions

diff --git a/rts/CloneStack.c b/rts/CloneStack.c
new file mode 100644
index 0000000000..a8e826eec1
--- /dev/null
+++ b/rts/CloneStack.c
@@ -0,0 +1,103 @@
+/* ---------------------------------------------------------------------------
+ *
+ * (c) The GHC Team, 2001-2021
+ *
+ * Stack snapshotting.
+ */
+
+#include <string.h>
+
+#include "Rts.h"
+#include "rts/Messages.h"
+#include "Messages.h"
+#include "rts/storage/TSO.h"
+#include "stg/Types.h"
+#include "CloneStack.h"
+#include "StablePtr.h"
+#include "Threads.h"
+
+#if defined(DEBUG)
+#include "sm/Sanity.h"
+#endif
+
+static StgStack* cloneStackChunk(Capability* capability, const StgStack* stack)
+{
+  StgWord spOffset = stack->sp - stack->stack;
+  StgWord closureSizeBytes = sizeof(StgStack) + (stack->stack_size * sizeof(StgWord));
+
+  StgStack* newStackClosure = (StgStack*) allocate(capability, ROUNDUP_BYTES_TO_WDS(closureSizeBytes));
+
+  memcpy(newStackClosure, stack, closureSizeBytes);
+
+  newStackClosure->sp = newStackClosure->stack + spOffset;
+  // The new stack is not on the mutable list; clear the dirty flag such that
+  // we don't claim that it is.
+  newStackClosure->dirty = 0;
+
+#if defined(DEBUG)
+  checkClosure((StgClosure*) newStackClosure);
+#endif
+
+  return newStackClosure;
+}
+
+StgStack* cloneStack(Capability* capability, const StgStack* stack)
+{
+  StgStack *top_stack = cloneStackChunk(capability, stack);
+  StgStack *last_stack = top_stack;
+  while (true) {
+    // check whether the stack ends in an underflow frame
+    StgPtr top = last_stack->stack + last_stack->stack_size;
+    StgUnderflowFrame *underFlowFrame = ((StgUnderflowFrame *) top);
+    StgUnderflowFrame *frame = underFlowFrame--;
+    if (frame->info == &stg_stack_underflow_frame_info) {
+      StgStack *s = cloneStackChunk(capability, frame->next_chunk);
+      frame->next_chunk = s;
+      last_stack = s;
+    } else {
+      break;
+    }
+  }
+  return top_stack;
+}
+
+#if defined(THREADED_RTS)
+
+// ThreadId# in Haskell is a StgTSO* in RTS.
+void sendCloneStackMessage(StgTSO *tso, HsStablePtr mvar) {
+  Capability *srcCapability = rts_unsafeGetMyCapability();
+
+  MessageCloneStack *msg;
+  msg = (MessageCloneStack *)allocate(srcCapability, sizeofW(MessageCloneStack));
+  msg->tso = tso;
+  msg->result = (StgMVar*)deRefStablePtr(mvar);
+  SET_HDR(msg, &stg_MSG_CLONE_STACK_info, CCS_SYSTEM);
+  // Ensure that writes constructing Message are committed before sending.
+  write_barrier();
+
+  sendMessage(srcCapability, tso->cap, (Message *)msg);
+}
+
+void handleCloneStackMessage(MessageCloneStack *msg){
+  StgStack* newStackClosure = cloneStack(msg->tso->cap, msg->tso->stackobj);
+
+  // Lift StackSnapshot# to StackSnapshot by applying it's constructor.
+  // This is necessary because performTryPutMVar() puts the closure onto the
+  // stack for evaluation and stacks can not be evaluated (entered).
+  HaskellObj result = rts_apply(msg->tso->cap, StackSnapshot_constructor_closure, (HaskellObj) newStackClosure);
+
+  bool putMVarWasSuccessful = performTryPutMVar(msg->tso->cap, msg->result, result);
+
+  if(!putMVarWasSuccessful) {
+    barf("Can't put stack cloning result into MVar.");
+  }
+}
+
+#else // !defined(THREADED_RTS)
+
+GNU_ATTRIBUTE(__noreturn__)
+void sendCloneStackMessage(StgTSO *tso STG_UNUSED, HsStablePtr mvar STG_UNUSED) {
+  barf("Sending CloneStackMessages is only available in threaded RTS!");
+}
+
+#endif // end !defined(THREADED_RTS)
diff --git a/rts/CloneStack.h b/rts/CloneStack.h
new file mode 100644
index 0000000000..5f1c22039d
--- /dev/null
+++ b/rts/CloneStack.h
@@ -0,0 +1,23 @@
+/* ---------------------------------------------------------------------------
+ *
+ * (c) The GHC Team, 2001-2021
+ *
+ * Stack snapshotting.
+ */
+
+#pragma once
+
+extern StgClosure DLL_IMPORT_DATA_VARNAME(base_GHCziStackziCloneStack_StackSnapshot_closure);
+#define StackSnapshot_constructor_closure DLL_IMPORT_DATA_REF(base_GHCziStackziCloneStack_StackSnapshot_closure)
+
+StgStack* cloneStack(Capability* capability, const StgStack* stack);
+
+void sendCloneStackMessage(StgTSO *tso, HsStablePtr mvar);
+
+#include "BeginPrivate.h"
+
+#if defined(THREADED_RTS)
+void handleCloneStackMessage(MessageCloneStack *msg);
+#endif
+
+#include "EndPrivate.h"
diff --git a/rts/Messages.c b/rts/Messages.c
index 8cd0a5570f..2ec12da3ad 100644
--- a/rts/Messages.c
+++ b/rts/Messages.c
@@ -14,6 +14,7 @@
 #include "Threads.h"
 #include "RaiseAsync.h"
 #include "sm/Storage.h"
+#include "CloneStack.h"
 
 /* ----------------------------------------------------------------------------
    Send a message to another Capability
@@ -32,7 +33,8 @@ void sendMessage(Capability *from_cap, Capability *to_cap, Message *msg)
             i != &stg_MSG_BLACKHOLE_info &&
             i != &stg_MSG_TRY_WAKEUP_info &&
             i != &stg_IND_info && // can happen if a MSG_BLACKHOLE is revoked
-            i != &stg_WHITEHOLE_info) {
+            i != &stg_WHITEHOLE_info &&
+            i != &stg_MSG_CLONE_STACK_info) {
             barf("sendMessage: %p", i);
         }
     }
@@ -130,6 +132,10 @@ loop:
 #endif
         goto loop;
     }
+    else if(i == &stg_MSG_CLONE_STACK_info){
+        MessageCloneStack *cloneStackMessage = (MessageCloneStack*) m;
+        handleCloneStackMessage(cloneStackMessage);
+    }
     else
     {
         barf("executeMessage: %p", i);
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 8f99105b18..122bac4e08 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -2927,3 +2927,14 @@ stg_setThreadAllocationCounterzh ( I64 counter )
     StgTSO_alloc_limit(CurrentTSO) = counter + TO_I64(offset);
     return ();
 }
+
+stg_cloneMyStackzh () {
+    W_ stgStack;
+    W_ clonedStack;
+    stgStack = StgTSO_stackobj(CurrentTSO);
+    StgStack_sp(stgStack) = Sp;
+
+    ("ptr" clonedStack) = ccall cloneStack(MyCapability() "ptr", stgStack "ptr");
+
+    return (clonedStack);
+}
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index b17d017c31..56d94ee7d2 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -12,6 +12,7 @@
 #include "Rts.h"
 #include "TopHandler.h"
 #include "HsFFI.h"
+#include "CloneStack.h"
 
 #include "sm/Storage.h"
 #include "sm/NonMovingMark.h"
@@ -988,6 +989,7 @@
       SymI_HasProto(stg_traceBinaryEventzh)                             \
       SymI_HasProto(stg_getThreadAllocationCounterzh)                   \
       SymI_HasProto(stg_setThreadAllocationCounterzh)                   \
+      SymI_HasProto(stg_cloneMyStackzh)                                 \
       SymI_HasProto(getMonotonicNSec)                                   \
       SymI_HasProto(lockFile)                                           \
       SymI_HasProto(unlockFile)                                         \
@@ -1012,6 +1014,7 @@
       SymI_HasProto(keepCAFs)                                           \
       SymI_HasProto(registerInfoProvList)                               \
       SymI_HasProto(lookupIPE)                                          \
+      SymI_HasProto(sendCloneStackMessage)                              \
       RTS_USER_SIGNALS_SYMBOLS                                          \
       RTS_INTCHAR_SYMBOLS
 
diff --git a/rts/StgMiscClosures.cmm b/rts/StgMiscClosures.cmm
index 244f55d67a..b78fca74cc 100644
--- a/rts/StgMiscClosures.cmm
+++ b/rts/StgMiscClosures.cmm
@@ -775,6 +775,9 @@ INFO_TABLE_CONSTR(stg_MSG_BLACKHOLE,3,0,0,PRIM,"MSG_BLACKHOLE","MSG_BLACKHOLE")
 INFO_TABLE_CONSTR(stg_MSG_NULL,1,0,0,PRIM,"MSG_NULL","MSG_NULL")
 { foreign "C" barf("MSG_NULL object (%p) entered!", R1) never returns; }
 
+INFO_TABLE_CONSTR(stg_MSG_CLONE_STACK,3,0,0,PRIM,"MSG_CLONE_STACK","MSG_CLONE_STACK")
+{ foreign "C" barf("stg_MSG_CLONE_STACK object (%p) entered!", R1) never returns; }
+
 /* ----------------------------------------------------------------------------
    END_TSO_QUEUE
 
diff --git a/rts/include/rts/storage/Closures.h b/rts/include/rts/storage/Closures.h
index ebb836bca2..b28315f76a 100644
--- a/rts/include/rts/storage/Closures.h
+++ b/rts/include/rts/storage/Closures.h
@@ -427,6 +427,14 @@ typedef struct MessageBlackHole_ {
     StgClosure *bh;
 } MessageBlackHole;
 
+typedef struct MessageCloneStack_ {
+    StgHeader header;
+    Message   *link;
+    StgMVar   *result;
+    StgTSO    *tso;
+} MessageCloneStack;
+
+
 /* ----------------------------------------------------------------------------
    Compact Regions
    ------------------------------------------------------------------------- */
diff --git a/rts/include/stg/MiscClosures.h b/rts/include/stg/MiscClosures.h
index c4cb7d45a4..1eed441eae 100644
--- a/rts/include/stg/MiscClosures.h
+++ b/rts/include/stg/MiscClosures.h
@@ -200,6 +200,7 @@ RTS_ENTRY(stg_STM_AWOKEN);
 RTS_ENTRY(stg_MSG_TRY_WAKEUP);
 RTS_ENTRY(stg_MSG_THROWTO);
 RTS_ENTRY(stg_MSG_BLACKHOLE);
+RTS_ENTRY(stg_MSG_CLONE_STACK);
 RTS_ENTRY(stg_MSG_NULL);
 RTS_ENTRY(stg_MVAR_TSO_QUEUE);
 RTS_ENTRY(stg_catch);
@@ -570,6 +571,7 @@ RTS_FUN_DECL(stg_traceBinaryEventzh);
 RTS_FUN_DECL(stg_traceMarkerzh);
 RTS_FUN_DECL(stg_getThreadAllocationCounterzh);
 RTS_FUN_DECL(stg_setThreadAllocationCounterzh);
+RTS_FUN_DECL(stg_cloneMyStackzh);
 
 
 /* Other misc stuff */
diff --git a/rts/package.conf.in b/rts/package.conf.in
index bc24614085..8baf62bc7d 100644
--- a/rts/package.conf.in
+++ b/rts/package.conf.in
@@ -192,6 +192,7 @@ ld-options:
            * so we need to force it to be included in the binary. */
          , "-Wl,-u,_findPtr"
 #endif
+         , "-Wl,-u,_base_GHCziStackziCloneStack_StackSnapshot_closure"
 #else
            "-Wl,-u,base_GHCziTopHandler_runIO_closure"
          , "-Wl,-u,base_GHCziTopHandler_runNonIO_closure"
@@ -304,6 +305,7 @@ ld-options:
            * so we need to force it to be included in the binary. */
          , "-Wl,-u,findPtr"
 #endif
+         , "-Wl,-u,base_GHCziStackziCloneStack_StackSnapshot_closure"
 #endif
 
 /*  Pick up static libraries in preference over dynamic if in earlier search
diff --git a/rts/rts.cabal.in b/rts/rts.cabal.in
index 6e535a777c..300a3bcd46 100644
--- a/rts/rts.cabal.in
+++ b/rts/rts.cabal.in
@@ -309,6 +309,7 @@ library
          "-Wl,-u,_hs_atomicwrite8"
          "-Wl,-u,_hs_atomicwrite16"
          "-Wl,-u,_hs_atomicwrite32"
+         "-Wl,-u,_base_GHCziStackziCloneStack_StackSnapshot_closure"
 
       if flag(find-ptr)
         -- This symbol is useful in gdb, but not referred to anywhere,
@@ -391,6 +392,7 @@ library
          "-Wl,-u,hs_atomicwrite8"
          "-Wl,-u,hs_atomicwrite16"
          "-Wl,-u,hs_atomicwrite32"
+         "-Wl,-u,base_GHCziStackziCloneStack_StackSnapshot_closure"
 
       if flag(find-ptr)
         -- This symbol is useful in gdb, but not referred to anywhere,
@@ -454,6 +456,7 @@ library
                Arena.c
                Capability.c
                CheckUnload.c
+               CloneStack.c
                ClosureFlags.c
                Disassembler.c
                FileLock.c