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>

7 files changed, 283 insertions, 16 deletions

diff --git a/testsuite/tests/profiling/should_run/T7275.stdout b/testsuite/tests/profiling/should_run/T7275.stdout
index f99f019120..d0146366a7 100644
--- a/testsuite/tests/profiling/should_run/T7275.stdout
+++ b/testsuite/tests/profiling/should_run/T7275.stdout
@@ -3,19 +3,19 @@
 2
 3
 4
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
-(282)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
+(284)suzanne/robert
diff --git a/testsuite/tests/rts/all.T b/testsuite/tests/rts/all.T
index 1dfb20b791..2c73973680 100644
--- a/testsuite/tests/rts/all.T
+++ b/testsuite/tests/rts/all.T
@@ -487,3 +487,8 @@ 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'])
+test('cloneMyStack2', ignore_stdout, compile_and_run, [''])
+test('cloneMyStack_retBigStackFrame', [extra_files(['cloneStackLib.c']), ignore_stdout], compile_and_run, ['cloneStackLib.c'])
+
+test('cloneThreadStack', [only_ways(['threaded1']), extra_ways(['threaded1']), extra_files(['cloneStackLib.c'])], compile_and_run, ['cloneStackLib.c -threaded'])
diff --git a/testsuite/tests/rts/cloneMyStack.hs b/testsuite/tests/rts/cloneMyStack.hs
new file mode 100644
index 0000000000..cdc93e6004
--- /dev/null
+++ b/testsuite/tests/rts/cloneMyStack.hs
@@ -0,0 +1,29 @@
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnliftedFFITypes #-}
+
+import GHC.Prim (StackSnapshot#)
+import GHC.Stack.CloneStack
+import Foreign
+import Foreign.C.Types (CUInt)
+import System.Mem
+
+foreign import ccall "expectClosureTypes" expectClosureTypes:: StackSnapshot# -> Ptr CUInt -> Int -> IO ()
+
+-- | Clone the stack and check that all expected closures are on it in order.
+-- (The check is done by closure type.)
+-- In the meanwhile enforce a garbage collection to ensure that the stack
+-- snapshot is still valid afterwards (is not gc'ed while in use).
+main :: IO ()
+main = do
+    stackSnapshot <- cloneMyStack
+
+    performMajorGC
+
+    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/cloneMyStack2.hs b/testsuite/tests/rts/cloneMyStack2.hs
new file mode 100644
index 0000000000..068c816ce5
--- /dev/null
+++ b/testsuite/tests/rts/cloneMyStack2.hs
@@ -0,0 +1,15 @@
+{-# LANGUAGE BangPatterns #-}
+module Main where
+
+import GHC.Stack.CloneStack
+
+main = foo 100
+
+{-# NOINLINE foo #-}
+foo 0 = () <$ getStack
+foo n = print "x" >> foo (n - 1) >> print "x"
+
+-- This shouldn't segfault
+getStack = do
+  !s <- cloneMyStack
+  return ()
diff --git a/testsuite/tests/rts/cloneMyStack_retBigStackFrame.hs b/testsuite/tests/rts/cloneMyStack_retBigStackFrame.hs
new file mode 100644
index 0000000000..cccc8ec618
--- /dev/null
+++ b/testsuite/tests/rts/cloneMyStack_retBigStackFrame.hs
@@ -0,0 +1,51 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnliftedFFITypes #-}
+
+module Main where
+
+import Control.Concurrent
+import Data.IORef
+import GHC.IO.Unsafe
+import GHC.Prim (StackSnapshot#)
+import GHC.Stack.CloneStack
+import System.Mem
+
+foreign import ccall "expectSixtyFourOnesInRetBigFrame" expectSixtyFourOnesInRetBigFrame :: StackSnapshot# -> IO ()
+
+cloneStack_returnInt :: IORef (Maybe StackSnapshot) -> Int
+cloneStack_returnInt ioRef = unsafePerformIO $ do
+  stackSnapshot <- cloneMyStack
+
+  performMajorGC
+
+  writeIORef ioRef (Just stackSnapshot)
+
+  performMajorGC
+
+  return 42
+
+-- | Clone a stack with a RET_BIG closure and check it in snapshot.
+-- In the meanwhile enforce several garbage collections in different places to
+-- ensure that the stack snapshot is still valid afterwards (is not gc'ed while
+-- in use).
+main :: IO ()
+main = do
+  stackRef <- newIORef Nothing
+
+  bigFun (cloneStack_returnInt stackRef) 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
+
+  Just (StackSnapshot stackSnapshot) <- readIORef stackRef
+
+  performMajorGC
+
+  expectSixtyFourOnesInRetBigFrame stackSnapshot
+
+  return ()
+
+bigFun !a1 a2 a3 a4 a5 a6 a7 a8 a9 a10 a11 a12 a13 a14 a15 a16 a17 a18 a19 a20 a21 a22 a23 a24 a25 a26 a27 a28 a29 a30 a31 a32 a33 a34 a35 a36 a37 a38 a39 a40 a41 a42 a43 a44 a45 a46 a47 a48 a49 a50 a51 a52 a53 a54 a55 a56 a57 a58 a59 a60 a61 a62 a63 a64 a65 =
+  do
+    print $ a1 + a2 + a3 + a4 + a5 + a6 + a7 + a8 + a9 + a10 + a11 + a12 + a13 + a14 + a15 + a16 + a17 + a18 + a19 + a20 + a21 + a22 + a23 + a24 + a25 + a26 + a27 + a28 + a29 + a30 + a31 + a32 + a33 + a34 + a35 + a36 + a37 + a38 + a39 + a40 + a41 + a42 + a43 + a44 + a45 + a46 + a47 + a48 + a49 + a50 + a51 + a52 + a53 + a54 + a55 + a56 + a57 + a58 + a59 + a60 + a61 + a62 + a63 + a64 + a65
+
+    return ()
diff --git a/testsuite/tests/rts/cloneStackLib.c b/testsuite/tests/rts/cloneStackLib.c
new file mode 100644
index 0000000000..a89a069886
--- /dev/null
+++ b/testsuite/tests/rts/cloneStackLib.c
@@ -0,0 +1,111 @@
+#include "Rts.h"
+#include "RtsAPI.h"
+#include "rts/Messages.h"
+#include <string.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 %lu", types[i], info->type, i);
+        }
+    }
+}
+
+// Count all (#I 1) closures of the RET_BIG closure's payload.
+static int countOnes(StgPtr spBottom, StgPtr payload,
+                     StgLargeBitmap *large_bitmap, uint32_t size) {
+  StgWord bmp;
+  uint32_t i, j;
+  int ones = 0;
+
+  i = 0;
+  for (bmp = 0; i < size; bmp++) {
+    StgWord bitmap = large_bitmap->bitmap[bmp];
+    j = 0;
+    for (; i < size && j < BITS_IN(W_); j++, i++, bitmap >>= 1) {
+      if ((bitmap & 1) == 0) {
+        const StgClosure *closure = UNTAG_CLOSURE((StgClosure *)payload[i]);
+        const StgInfoTable *info = get_itbl(closure);
+
+        switch (info->type) {
+        case CONSTR_0_1: {
+          const StgConInfoTable *con_info = get_con_itbl(closure);
+          if (strcmp(GET_CON_DESC(con_info), "ghc-prim:GHC.Types.I#") == 0 &&
+              closure->payload[0] == 1) {
+            ones++;
+          }
+          break;
+        }
+        default: {
+          break;
+        }
+        }
+      }
+    }
+  }
+
+  return ones;
+}
+
+void expectSixtyFourOnesInRetBigFrame(StgStack *stack) {
+  StgPtr sp = stack->sp;
+  StgPtr spBottom = stack->stack + stack->stack_size;
+
+  for (; sp < spBottom; sp += stack_frame_sizeW((StgClosure *)sp)) {
+    const StgInfoTable *info = get_itbl((StgClosure *)sp);
+
+    if (info->type == RET_BIG) {
+      StgLargeBitmap *bitmap = GET_LARGE_BITMAP(info);
+      int ones = countOnes(spBottom, (StgPtr)((StgClosure *)sp)->payload,
+                           bitmap, bitmap->size);
+
+      if (ones != 64) {
+        barf("Expected 64 ones, got %i!", ones);
+      }
+    }
+  }
+}
diff --git a/testsuite/tests/rts/cloneThreadStack.hs b/testsuite/tests/rts/cloneThreadStack.hs
new file mode 100644
index 0000000000..11b37d3577
--- /dev/null
+++ b/testsuite/tests/rts/cloneThreadStack.hs
@@ -0,0 +1,56 @@
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnliftedFFITypes #-}
+
+import GHC.Prim (StackSnapshot#, ThreadId#)
+import GHC.Conc.Sync (ThreadId(..))
+import GHC.Stack.CloneStack
+import Control.Concurrent
+import GHC.Conc
+import System.Mem
+
+foreign import ccall "expectStacksToBeEqual" expectStacksToBeEqual:: StackSnapshot# -> ThreadId# -> IO ()
+
+foreign import ccall "expectStackToBeNotDirty" expectStackToBeNotDirty:: StackSnapshot# -> IO ()
+
+-- | Clone the stack of another thread and check it's snapshot for being equal
+-- with the live stack.
+-- In the meanwhile enforce a garbage collection to ensure that the stack
+-- snapshot is still valid afterwards (is not gc'ed while in use).
+main :: IO ()
+main = do
+    mVarToBeBlockedOn <- newEmptyMVar
+    threadId <- forkIO $ immediatelyBlocking mVarToBeBlockedOn
+
+    waitUntilBlocked threadId
+
+    stackSnapshot <- cloneThreadStack threadId
+
+    performMajorGC
+
+    let (StackSnapshot stack) = stackSnapshot
+    let (ThreadId tid#) = threadId
+    expectStacksToBeEqual stack tid#
+    expectStackToBeNotDirty stack
+
+immediatelyBlocking :: MVar Int -> IO ()
+immediatelyBlocking mVarToBeBlockedOn = do
+    takeMVar mVarToBeBlockedOn
+    return ()
+
+waitUntilBlocked :: ThreadId -> IO ()
+waitUntilBlocked tid = do
+    blocked <- isBlocked tid
+    if blocked then
+        return ()
+    else
+        do
+            threadDelay 100000
+            waitUntilBlocked tid
+
+isBlocked:: ThreadId -> IO Bool
+isBlocked = fmap isThreadStatusBlocked . threadStatus
+
+isThreadStatusBlocked :: ThreadStatus -> Bool
+isThreadStatusBlocked (ThreadBlocked _) = True
+isThreadStatusBlocked _ = False