Introduce snapshotting of thread's own stack (#18741)

Introduce `StackSnapshot#` type and the `cloneMyStack#` primop, allowing
the user to reify the state of the calling thread's stack for later
inspection.

The stack snapshot is offline/cold, i.e. it isn't evaluated any further.

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>

17 files changed, 250 insertions, 2 deletions

diff --git a/compiler/GHC/Builtin/Names.hs b/compiler/GHC/Builtin/Names.hs
index 913b8a4ed9..f251856c75 100644
--- a/compiler/GHC/Builtin/Names.hs
+++ b/compiler/GHC/Builtin/Names.hs
@@ -1884,7 +1884,7 @@ statePrimTyConKey, stableNamePrimTyConKey, stableNameTyConKey,
     typeConKey, threadIdPrimTyConKey, bcoPrimTyConKey, ptrTyConKey,
     funPtrTyConKey, tVarPrimTyConKey, eqPrimTyConKey,
     eqReprPrimTyConKey, eqPhantPrimTyConKey,
-    compactPrimTyConKey :: Unique
+    compactPrimTyConKey, stackSnapshotPrimTyConKey :: Unique
 statePrimTyConKey                       = mkPreludeTyConUnique 50
 stableNamePrimTyConKey                  = mkPreludeTyConUnique 51
 stableNameTyConKey                      = mkPreludeTyConUnique 52
@@ -1912,6 +1912,7 @@ ptrTyConKey                             = mkPreludeTyConUnique 77
 funPtrTyConKey                          = mkPreludeTyConUnique 78
 tVarPrimTyConKey                        = mkPreludeTyConUnique 79
 compactPrimTyConKey                     = mkPreludeTyConUnique 80
+stackSnapshotPrimTyConKey               = mkPreludeTyConUnique 81
 
 eitherTyConKey :: Unique
 eitherTyConKey                          = mkPreludeTyConUnique 84
diff --git a/compiler/GHC/Builtin/Types/Prim.hs b/compiler/GHC/Builtin/Types/Prim.hs
index ce4f1e5dc0..9ba6d87927 100644
--- a/compiler/GHC/Builtin/Types/Prim.hs
+++ b/compiler/GHC/Builtin/Types/Prim.hs
@@ -83,6 +83,7 @@ module GHC.Builtin.Types.Prim(
         bcoPrimTyCon,                   bcoPrimTy,
         weakPrimTyCon,                  mkWeakPrimTy,
         threadIdPrimTyCon,              threadIdPrimTy,
+        stackSnapshotPrimTyCon,         stackSnapshotPrimTy,
 
         int8PrimTyCon,          int8PrimTy, int8PrimTyConName,
         word8PrimTyCon,         word8PrimTy, word8PrimTyConName,
@@ -203,6 +204,7 @@ exposedPrimTyCons
     , word16PrimTyCon
     , word32PrimTyCon
     , word64PrimTyCon
+    , stackSnapshotPrimTyCon
 
     , tYPETyCon
     , funTyCon
@@ -225,7 +227,7 @@ mkBuiltInPrimTc fs unique tycon
                   BuiltInSyntax
 
 
-charPrimTyConName, intPrimTyConName, int8PrimTyConName, int16PrimTyConName, int32PrimTyConName, int64PrimTyConName, wordPrimTyConName, word32PrimTyConName, word8PrimTyConName, word16PrimTyConName, word64PrimTyConName, addrPrimTyConName, floatPrimTyConName, doublePrimTyConName, statePrimTyConName, proxyPrimTyConName, realWorldTyConName, arrayPrimTyConName, arrayArrayPrimTyConName, smallArrayPrimTyConName, byteArrayPrimTyConName, mutableArrayPrimTyConName, mutableByteArrayPrimTyConName, mutableArrayArrayPrimTyConName, smallMutableArrayPrimTyConName, mutVarPrimTyConName, mVarPrimTyConName, ioPortPrimTyConName, tVarPrimTyConName, stablePtrPrimTyConName, stableNamePrimTyConName, compactPrimTyConName, bcoPrimTyConName, weakPrimTyConName, threadIdPrimTyConName, eqPrimTyConName, eqReprPrimTyConName, eqPhantPrimTyConName :: Name
+charPrimTyConName, intPrimTyConName, int8PrimTyConName, int16PrimTyConName, int32PrimTyConName, int64PrimTyConName, wordPrimTyConName, word32PrimTyConName, word8PrimTyConName, word16PrimTyConName, word64PrimTyConName, addrPrimTyConName, floatPrimTyConName, doublePrimTyConName, statePrimTyConName, proxyPrimTyConName, realWorldTyConName, arrayPrimTyConName, arrayArrayPrimTyConName, smallArrayPrimTyConName, byteArrayPrimTyConName, mutableArrayPrimTyConName, mutableByteArrayPrimTyConName, mutableArrayArrayPrimTyConName, smallMutableArrayPrimTyConName, mutVarPrimTyConName, mVarPrimTyConName, ioPortPrimTyConName, tVarPrimTyConName, stablePtrPrimTyConName, stableNamePrimTyConName, compactPrimTyConName, bcoPrimTyConName, weakPrimTyConName, threadIdPrimTyConName, eqPrimTyConName, eqReprPrimTyConName, eqPhantPrimTyConName, stackSnapshotPrimTyConName :: Name
 charPrimTyConName             = mkPrimTc (fsLit "Char#") charPrimTyConKey charPrimTyCon
 intPrimTyConName              = mkPrimTc (fsLit "Int#") intPrimTyConKey  intPrimTyCon
 int8PrimTyConName             = mkPrimTc (fsLit "Int8#") int8PrimTyConKey int8PrimTyCon
@@ -261,6 +263,7 @@ tVarPrimTyConName             = mkPrimTc (fsLit "TVar#") tVarPrimTyConKey tVarPr
 stablePtrPrimTyConName        = mkPrimTc (fsLit "StablePtr#") stablePtrPrimTyConKey stablePtrPrimTyCon
 stableNamePrimTyConName       = mkPrimTc (fsLit "StableName#") stableNamePrimTyConKey stableNamePrimTyCon
 compactPrimTyConName          = mkPrimTc (fsLit "Compact#") compactPrimTyConKey compactPrimTyCon
+stackSnapshotPrimTyConName    = mkPrimTc (fsLit "StackSnapshot#") stackSnapshotPrimTyConKey stackSnapshotPrimTyCon
 bcoPrimTyConName              = mkPrimTc (fsLit "BCO") bcoPrimTyConKey bcoPrimTyCon
 weakPrimTyConName             = mkPrimTc (fsLit "Weak#") weakPrimTyConKey weakPrimTyCon
 threadIdPrimTyConName         = mkPrimTc (fsLit "ThreadId#") threadIdPrimTyConKey threadIdPrimTyCon
@@ -1159,6 +1162,21 @@ compactPrimTy = mkTyConTy compactPrimTyCon
 {-
 ************************************************************************
 *                                                                      *
+   The @StackSnapshot#@ type
+*                                                                      *
+************************************************************************
+-}
+
+stackSnapshotPrimTyCon :: TyCon
+stackSnapshotPrimTyCon = pcPrimTyCon0 stackSnapshotPrimTyConName UnliftedRep
+
+stackSnapshotPrimTy :: Type
+stackSnapshotPrimTy = mkTyConTy stackSnapshotPrimTyCon
+
+
+{-
+************************************************************************
+*                                                                      *
    The ``bytecode object'' type
 *                                                                      *
 ************************************************************************
diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index bf3b879449..6f5bc57292 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -3647,6 +3647,16 @@ primop  SetThreadAllocationCounter "setThreadAllocationCounter#" GenPrimOp
    has_side_effects = True
    out_of_line      = True
 
+primtype StackSnapshot#
+
+primop  CloneMyStack "cloneMyStack#" GenPrimOp
+   State# RealWorld -> (# State# RealWorld, StackSnapshot# #)
+   { Clones the stack of the current Haskell thread. }
+   with
+   has_side_effects = True
+   out_of_line      = True
+
+
 ------------------------------------------------------------------------
 section "Safe coercions"
 ------------------------------------------------------------------------
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index e268761cd7..d054eb8cfa 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1669,6 +1669,7 @@ emitPrimOp dflags primop = case primop of
   TraceEventBinaryOp -> alwaysExternal
   TraceMarkerOp -> alwaysExternal
   SetThreadAllocationCounter -> alwaysExternal
+  CloneMyStack -> alwaysExternal
 
   -- See Note [keepAlive# magic] in GHC.CoreToStg.Prep.
   KeepAliveOp -> panic "keepAlive# should have been eliminated in CorePrep"
diff --git a/libraries/base/GHC/Stack/CloneStack.hs b/libraries/base/GHC/Stack/CloneStack.hs
new file mode 100644
index 0000000000..ab217032f3
--- /dev/null
+++ b/libraries/base/GHC/Stack/CloneStack.hs
@@ -0,0 +1,47 @@
+{-# LANGUAGE MagicHash        #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+-- |
+-- This module exposes an interface for capturing the state of a thread's
+-- execution stack for diagnostics purposes.
+--
+-- @since 2.16.0.0
+module GHC.Stack.CloneStack (
+  StackSnapshot(..),
+  cloneMyStack
+  ) where
+
+import GHC.Prim (StackSnapshot#, cloneMyStack#)
+import GHC.IO (IO(..))
+
+-- | A frozen snapshot of the state of an execution stack.
+--
+-- @since 2.16.0.0
+data StackSnapshot = StackSnapshot !StackSnapshot#
+
+{-
+Note [Stack Cloning]
+~~~~~~~~~~~~~~~~~~~~
+"Cloning" a stack means that it's StgStack closure is copied including the
+stack memory (stack[]). The stack pointer (sp) of the clone is adjusted to be
+valid.
+
+The clone is "offline"/"cold", i.e. it won't be evaluated any further. This is
+useful for further analyses like stack unwinding or traversal.
+
+There are two different ways to clone a stack:
+1. By the corresponding thread via a primop call (cloneMyStack#).
+2. By sending a RTS message (Messages.c) with a MVar to the corresponding
+   thread and receiving the stack by taking it out of this MVar.
+
+A StackSnapshot# is really a pointer to an immutable StgStack closure with
+the invariant that stack->sp points to a valid frame.
+-}
+
+-- | Clone the stack of the executing thread
+--
+-- @since 2.16.0.0
+cloneMyStack :: IO StackSnapshot
+cloneMyStack = IO $ \s ->
+   case (cloneMyStack# s) of (# s1, stack #) -> (# s1, StackSnapshot stack #)
+
diff --git a/libraries/base/base.cabal b/libraries/base/base.cabal
index 7e09412839..3dc4b640a2 100644
--- a/libraries/base/base.cabal
+++ b/libraries/base/base.cabal
@@ -264,6 +264,7 @@ Library
         GHC.ResponseFile
         GHC.RTS.Flags
         GHC.ST
+        GHC.Stack.CloneStack
         GHC.StaticPtr
         GHC.STRef
         GHC.Show
diff --git a/rts/CloneStack.c b/rts/CloneStack.c
new file mode 100644
index 0000000000..84e0055465
--- /dev/null
+++ b/rts/CloneStack.c
@@ -0,0 +1,59 @@
+/* ---------------------------------------------------------------------------
+ *
+ * (c) The GHC Team, 2001-2021
+ *
+ * Stack snapshotting.
+ */
+
+#include <string.h>
+
+#include "Rts.h"
+#include "rts/storage/TSO.h"
+#include "stg/Types.h"
+#include "CloneStack.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;
+}
diff --git a/rts/CloneStack.h b/rts/CloneStack.h
new file mode 100644
index 0000000000..d9824c0c68
--- /dev/null
+++ b/rts/CloneStack.h
@@ -0,0 +1,13 @@
+/* ---------------------------------------------------------------------------
+ *
+ * (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);
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..1eb326cfb2 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)                                         \
diff --git a/rts/include/stg/MiscClosures.h b/rts/include/stg/MiscClosures.h
index c4cb7d45a4..988aae7345 100644
--- a/rts/include/stg/MiscClosures.h
+++ b/rts/include/stg/MiscClosures.h
@@ -570,6 +570,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
diff --git a/testsuite/tests/rts/all.T b/testsuite/tests/rts/all.T
index d2146a3f48..dae4aa26e7 100644
--- a/testsuite/tests/rts/all.T
+++ b/testsuite/tests/rts/all.T
@@ -486,3 +486,4 @@ test('T19381', extra_run_opts('+RTS -T -RTS'), compile_and_run, [''])
 test('T20199', normal, makefile_test, [])
 
 test('ipeMap', [c_src], compile_and_run, [''])
+test('cloneMyStack', [extra_files(['cloneStackLib.c'])], compile_and_run, ['cloneStackLib.c'])
diff --git a/testsuite/tests/rts/cloneMyStack.hs b/testsuite/tests/rts/cloneMyStack.hs
new file mode 100644
index 0000000000..0265205c28
--- /dev/null
+++ b/testsuite/tests/rts/cloneMyStack.hs
@@ -0,0 +1,22 @@
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnliftedFFITypes #-}
+
+import GHC.Prim (StackSnapshot#)
+import GHC.Stack.CloneStack
+import Foreign
+import Foreign.C.Types (CUInt)
+
+foreign import ccall "expectClosureTypes" expectClosureTypes:: StackSnapshot# -> Ptr CUInt -> Int -> IO ()
+
+main :: IO ()
+main = do
+    stackSnapshot <- cloneMyStack
+
+    let (StackSnapshot stack) = stackSnapshot
+    let expectedClosureTypes = [ 30 -- RET_SMALL
+                               , 30 -- RET_SMALL
+                               , 34 -- CATCH_FRAME
+                               ,36 -- STOP_FRAME
+                               ]
+    withArray expectedClosureTypes (\ptr -> expectClosureTypes stack ptr (length expectedClosureTypes))
diff --git a/testsuite/tests/rts/cloneStackLib.c b/testsuite/tests/rts/cloneStackLib.c
new file mode 100644
index 0000000000..113c2eb7d0
--- /dev/null
+++ b/testsuite/tests/rts/cloneStackLib.c
@@ -0,0 +1,55 @@
+#include "Rts.h"
+#include "RtsAPI.h"
+#include "rts/Messages.h"
+
+
+void expectStacksToBeEqual(StgStack *clonedStack, StgTSO *tso) {
+    StgStack *liveStack = tso->stackobj;
+
+    if(liveStack->header.info != clonedStack->header.info){
+        barf("info table pointer not equal! Expected same pointer address, but got %p and %p", liveStack->header.info, clonedStack->header.info);
+    }
+
+    StgInfoTable *info = INFO_PTR_TO_STRUCT(liveStack->header.info);
+
+    if (info->type != STACK) {
+        barf("Expected a closure of type STACK!");
+    }
+
+    if(liveStack->stack_size != clonedStack->stack_size){
+        barf("Expected same stack_size!");
+    }
+
+    if(liveStack->marking != clonedStack->marking){
+        barf("Expected same marking flags!");
+    }
+
+    for(StgWord i = liveStack->stack_size - 1; (liveStack->stack + i) >= liveStack->sp; i--){
+        if(liveStack->stack[i] != clonedStack->stack[i]){
+            barf("Expected stack word %lu to be equal on both stacks.", i);
+        }
+    }
+}
+
+void expectStackToBeNotDirty(StgStack *stack) {
+    if(stack->dirty != 0) {
+        barf("Expected stack to be not dirty. But dirty flag was set to %u", stack->dirty);
+    }
+}
+
+void expectClosureTypes(StgStack *stack, unsigned int types[], size_t typesSize){
+    StgPtr sp = stack->sp;
+    StgPtr spBottom = stack->stack + stack->stack_size;
+
+    for (StgWord i = 0; sp < spBottom; sp += stack_frame_sizeW((StgClosure *)sp), i++) {
+        const StgInfoTable *info = get_itbl((StgClosure *)sp);
+
+        if(i >= typesSize) {
+            barf("Stack size exceeds expectation!");
+        }
+
+        if(info->type != types[i]) {
+            barf("Wrong closure type on stack! Expected %u but got %u in position %i", types[i], info->type, i);
+        }
+    }
+}
diff --git a/utils/genprimopcode/Main.hs b/utils/genprimopcode/Main.hs
index 06a4922aa3..c13447e527 100644
--- a/utils/genprimopcode/Main.hs
+++ b/utils/genprimopcode/Main.hs
@@ -884,6 +884,7 @@ ppType (TyApp (TyCon "ThreadId#")   []) = "threadIdPrimTy"
 ppType (TyApp (TyCon "ForeignObj#") []) = "foreignObjPrimTy"
 ppType (TyApp (TyCon "BCO")         []) = "bcoPrimTy"
 ppType (TyApp (TyCon "Compact#")    []) = "compactPrimTy"
+ppType (TyApp (TyCon "StackSnapshot#") []) = "stackSnapshotPrimTy"
 ppType (TyApp (TyCon "()")          []) = "unitTy"      -- unitTy is GHC.Builtin.Types's name for ()
 
 ppType (TyVar "a")                      = "alphaTy"