Add primop to list threads

A user came to #ghc yesterday wondering how best to check whether they
were leaking threads. We ended up using the eventlog but it seems to me
like it would be generally useful if Haskell programs could query their
own threads.

16 files changed, 135 insertions, 4 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index ac03c20dbd..d59e68626f 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -2938,6 +2938,26 @@ primop  NoDuplicateOp "noDuplicate#" GenPrimOp
 
 primop  ThreadStatusOp "threadStatus#" GenPrimOp
    ThreadId# -> State# RealWorld -> (# State# RealWorld, Int#, Int#, Int# #)
+   {Get the status of the given thread. Result is
+    @(ThreadStatus, Capability, Locked)@ where
+    @ThreadStatus@ is one of the status constants defined in
+    @rts/Constants.h@, @Capability@ is the number of
+    the capability which currently owns the thread, and
+    @Locked@ is a boolean indicating whether the
+    thread is bound to that capability.
+
+    @since 0.9}
+   with
+   out_of_line = True
+   has_side_effects = True
+
+primop ListThreadsOp "listThreads#" GenPrimOp
+   State# RealWorld -> (# State# RealWorld, Array# ThreadId# #)
+   { Returns an array of the threads started by the program. Note that this
+     threads which have finished execution may or may not be present in this
+     list, depending upon whether they have been collected by the garbage collector.
+
+     @since 0.10}
    with
    out_of_line = True
    has_side_effects = True
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index 5d459ba7ad..b0abef9f4e 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1624,6 +1624,7 @@ emitPrimOp cfg primop =
   MkApUpd0_Op -> alwaysExternal
   NewBCOOp -> alwaysExternal
   UnpackClosureOp -> alwaysExternal
+  ListThreadsOp -> alwaysExternal
   ClosureSizeOp -> alwaysExternal
   WhereFromOp   -> alwaysExternal
   GetApStackValOp -> alwaysExternal
diff --git a/libraries/base/GHC/Conc.hs b/libraries/base/GHC/Conc.hs
index 15397422a5..fbb7521860 100644
--- a/libraries/base/GHC/Conc.hs
+++ b/libraries/base/GHC/Conc.hs
@@ -45,6 +45,7 @@ module GHC.Conc
         , yield
         , labelThread
         , mkWeakThreadId
+        , listThreads
 
         , ThreadStatus(..), BlockReason(..)
         , threadStatus
diff --git a/libraries/base/GHC/Conc/Sync.hs b/libraries/base/GHC/Conc/Sync.hs
index 6ae26678aa..bb8d6592b6 100644
--- a/libraries/base/GHC/Conc/Sync.hs
+++ b/libraries/base/GHC/Conc/Sync.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UnliftedFFITypes #-}
+{-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE Unsafe #-}
 
 {-# OPTIONS_HADDOCK not-home #-}
@@ -39,7 +40,7 @@ module GHC.Conc.Sync
         , yield
         , labelThread
         , mkWeakThreadId
-
+        , listThreads
         , ThreadStatus(..), BlockReason(..)
         , threadStatus
         , threadCapability
@@ -536,6 +537,25 @@ runSparks = IO loop
                       then (# s', () #)
                       else p `seq` loop s'
 
+-- | List the Haskell threads of the current process.
+--
+-- @since 4.18
+listThreads :: IO [ThreadId]
+listThreads = IO $ \s ->
+    case listThreads# s of
+      (# s', arr #) ->
+        (# s', mapListArrayUnlifted ThreadId arr #)
+
+mapListArrayUnlifted :: forall (a :: TYPE UnliftedRep) b. (a -> b) -> Array# a -> [b]
+mapListArrayUnlifted f arr = go 0#
+  where
+    sz = sizeofArray# arr
+    go i#
+      | isTrue# (i# ==# sz) = []
+      | otherwise = case indexArray# arr i# of
+                      (# x #) -> f x : go (i# +# 1#)
+{-# NOINLINE mapListArrayUnlifted #-}
+
 data BlockReason
   = BlockedOnMVar
         -- ^blocked on 'MVar'
@@ -575,6 +595,7 @@ data ThreadStatus
            , Show -- ^ @since 4.3.0.0
            )
 
+-- | Query the current execution status of a thread.
 threadStatus :: ThreadId -> IO ThreadStatus
 threadStatus (ThreadId t) = IO $ \s ->
    case threadStatus# t s of
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index 06570b5371..8991915db9 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -13,6 +13,8 @@
     [Core Libraries proposal](https://github.com/haskell/core-libraries-committee/issues/65).
   * Add default implementation of `(<>)` in terms of `sconcat` and `mempty` in
     terms of `mconcat`.
+  * `GHC.Conc.Sync.listThreads` was added, allowing the user to list the threads
+    (both running and blocked) of the program.
 
 ## 4.17.0.0 *TBA*
 
diff --git a/libraries/base/tests/all.T b/libraries/base/tests/all.T
index 69c990e0a2..7f0580b84a 100644
--- a/libraries/base/tests/all.T
+++ b/libraries/base/tests/all.T
@@ -273,3 +273,4 @@ test('T19288', exit_code(1), compile_and_run, [''])
 test('T19719', normal, compile_and_run, [''])
 test('T20107', extra_run_opts('+RTS -M50M'), compile_and_run, ['-package bytestring'])
 test('trace', normal, compile_and_run, [''])
+test('listThreads', normal, compile_and_run, [''])
diff --git a/libraries/base/tests/listThreads.hs b/libraries/base/tests/listThreads.hs
new file mode 100644
index 0000000000..00e2178297
--- /dev/null
+++ b/libraries/base/tests/listThreads.hs
@@ -0,0 +1,22 @@
+import Control.Concurrent
+import GHC.Conc.Sync
+
+dummyThread :: MVar () -> Int -> IO ()
+dummyThread mvar n = do
+  tid <- myThreadId
+  labelThread tid ("thread-"++show n)
+  readMVar mvar
+
+main :: IO ()
+main = do
+  mvar <- newEmptyMVar
+  let mkThread n = do
+        tid <- forkIO $ readMVar mvar
+        labelThread tid ("thread-"++show n)
+
+  mapM_ mkThread [0..100]
+  threads <- listThreads
+  -- TODO: Check labels
+  print $ length threads
+  putMVar mvar ()
+
diff --git a/libraries/base/tests/listThreads.stdout b/libraries/base/tests/listThreads.stdout
new file mode 100644
index 0000000000..257e563266
--- /dev/null
+++ b/libraries/base/tests/listThreads.stdout
@@ -0,0 +1 @@
+102
diff --git a/libraries/ghc-prim/changelog.md b/libraries/ghc-prim/changelog.md
index 6f665f3fcb..2c991a4a1e 100644
--- a/libraries/ghc-prim/changelog.md
+++ b/libraries/ghc-prim/changelog.md
@@ -1,3 +1,13 @@
+## 0.10.0
+
+- Shipped with GHC 9.6.1
+
+- The `listThreads#` primop was added, allowing the user to enumerate all
+  threads (running and blocked) in the program:
+  ```haskell
+  listThreads# :: State# RealWorld -> (# State# RealWorld, Array# ThreadId# #)
+  ```
+
 ## 0.9.0
 
 - Shipped with GHC 9.4.1
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 7b760e5702..503797f556 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -1081,6 +1081,14 @@ stg_labelThreadzh ( gcptr threadid, W_ addr )
     return ();
 }
 
+stg_listThreadszh ()
+{
+  P_ arr;
+
+  ("ptr" arr) = ccall listThreads(MyCapability() "ptr");
+  return (arr);
+}
+
 stg_isCurrentThreadBoundzh (/* no args */)
 {
     W_ r;
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 9731f4febf..40193a25a2 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -721,6 +721,7 @@ extern char **environ;
       SymI_HasDataProto(stg_isCurrentThreadBoundzh)                         \
       SymI_HasDataProto(stg_isEmptyMVarzh)                                  \
       SymI_HasDataProto(stg_killThreadzh)                                   \
+      SymI_HasDataProto(stg_listThreadszh)                                  \
       SymI_HasProto(loadArchive)                                        \
       SymI_HasProto(loadObj)                                            \
       SymI_HasProto(purgeObj)                                           \
diff --git a/rts/Threads.c b/rts/Threads.c
index af1758f2ef..6b478d046f 100644
--- a/rts/Threads.c
+++ b/rts/Threads.c
@@ -855,6 +855,43 @@ loop:
     return true;
 }
 
+StgMutArrPtrs *listThreads(Capability *cap)
+{
+    ACQUIRE_LOCK(&sched_mutex);
+
+    // First count how many threads we have...
+    StgWord n_threads = 0;
+    for (unsigned g = 0; g < RtsFlags.GcFlags.generations; g++) {
+        for (StgTSO *t = generations[g].threads; t != END_TSO_QUEUE; t = t->global_link) {
+            n_threads++;
+        }
+    }
+
+    // Allocate a suitably-sized array...
+    const StgWord size = n_threads + mutArrPtrsCardTableSize(n_threads);
+    StgMutArrPtrs *arr =
+        (StgMutArrPtrs *)allocate(cap, sizeofW(StgMutArrPtrs) + size);
+    TICK_ALLOC_PRIM(sizeofW(StgMutArrPtrs), n, 0);
+    arr->ptrs = n_threads;
+    arr->size = size;
+
+    // Populate it...
+    StgWord i = 0;
+    for (unsigned g = 0; g < RtsFlags.GcFlags.generations; g++) {
+        for (StgTSO *t = generations[g].threads; t != END_TSO_QUEUE; t = t->global_link) {
+            // It's possible that new threads have been created since we counted.
+            // Ignore them.
+            if (i == n_threads)
+                break;
+            arr->payload[i] = (StgClosure *) t;
+            i++;
+        }
+    }
+    CHECKM(i == n_threads, "listThreads: Found too few threads");
+    RELEASE_LOCK(&sched_mutex);
+    return arr;
+}
+
 /* ----------------------------------------------------------------------------
  * Debugging: why is a thread blocked
  * ------------------------------------------------------------------------- */
diff --git a/rts/include/rts/Threads.h b/rts/include/rts/Threads.h
index 46a20089a7..83ffce9238 100644
--- a/rts/include/rts/Threads.h
+++ b/rts/include/rts/Threads.h
@@ -53,6 +53,10 @@ StgThreadID rts_getThreadId                  (StgPtr tso);
 void        rts_enableThreadAllocationLimit  (StgPtr tso);
 void        rts_disableThreadAllocationLimit (StgPtr tso);
 
+// Forward declarations, defined in Closures.h
+struct _StgMutArrPtrs;
+struct _StgMutArrPtrs *listThreads               (Capability *cap);
+
 #if !defined(mingw32_HOST_OS)
 pid_t  forkProcess     (HsStablePtr *entry);
 #else
diff --git a/rts/include/rts/storage/Closures.h b/rts/include/rts/storage/Closures.h
index 84971d92a6..c1de80aa4d 100644
--- a/rts/include/rts/storage/Closures.h
+++ b/rts/include/rts/storage/Closures.h
@@ -206,7 +206,7 @@ typedef struct {
 // Closure types: MUT_ARR_PTRS_CLEAN, MUT_ARR_PTRS_DIRTY,
 // MUT_ARR_PTRS_FROZEN_DIRTY, MUT_ARR_PTRS_FROZEN_CLEAN, MUT_VAR_CLEAN,
 // MUT_VAR_DIRTY
-typedef struct {
+typedef struct _StgMutArrPtrs {
     StgHeader   header;
     StgWord     ptrs;
     StgWord     size; // ptrs plus card table
diff --git a/rts/include/rts/storage/GC.h b/rts/include/rts/storage/GC.h
index 478503aaee..77f7f38d9a 100644
--- a/rts/include/rts/storage/GC.h
+++ b/rts/include/rts/storage/GC.h
@@ -116,8 +116,9 @@ typedef struct generation_ {
     //
     memcount       max_blocks;
 
-    StgTSO *       threads;             // threads in this gen
-                                        // linked via global_link
+    StgTSO *       threads;             // threads in this generation.
+                                        // linked via global_link.
+                                        // protected by sched_mutex except during GC.
     StgWeak *      weak_ptr_list;       // weak pointers in this gen
 
     struct generation_ *to;             // destination gen for live objects
diff --git a/rts/include/stg/MiscClosures.h b/rts/include/stg/MiscClosures.h
index c1f9d3e94d..6b784306fb 100644
--- a/rts/include/stg/MiscClosures.h
+++ b/rts/include/stg/MiscClosures.h
@@ -533,6 +533,7 @@ RTS_FUN_DECL(stg_myThreadIdzh);
 RTS_FUN_DECL(stg_labelThreadzh);
 RTS_FUN_DECL(stg_isCurrentThreadBoundzh);
 RTS_FUN_DECL(stg_threadStatuszh);
+RTS_FUN_DECL(stg_listThreadszh);
 
 RTS_FUN_DECL(stg_mkWeakzh);
 RTS_FUN_DECL(stg_mkWeakNoFinalizzerzh);