Add a primop to query the label of a thread

10 files changed, 47 insertions, 1 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index 13c53b493b..d46e5401fb 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -2938,6 +2938,16 @@ primop  NoDuplicateOp "noDuplicate#" GenPrimOp
    out_of_line = True
    has_side_effects = True
 
+primop GetThreadLabelOp "threadLabel#" GenPrimOp
+   ThreadId# -> State# RealWorld -> (# State# RealWorld, Int#, ByteArray# #)
+   {Get the label of the given thread.
+    Morally of type @ThreadId# -> IO (Maybe ByteArray#)@, with a @1#@ tag
+    denoting @Just@.
+
+    @since 0.10}
+   with
+   out_of_line      = True
+
 primop  ThreadStatusOp "threadStatus#" GenPrimOp
    ThreadId# -> State# RealWorld -> (# State# RealWorld, Int#, Int#, Int# #)
    {Get the status of the given thread. Result is
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index b0abef9f4e..1d482b0143 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1597,6 +1597,7 @@ emitPrimOp cfg primop =
   LabelThreadOp -> alwaysExternal
   IsCurrentThreadBoundOp -> alwaysExternal
   NoDuplicateOp -> alwaysExternal
+  GetThreadLabelOp -> alwaysExternal
   ThreadStatusOp -> alwaysExternal
   MkWeakOp -> alwaysExternal
   MkWeakNoFinalizerOp -> alwaysExternal
diff --git a/libraries/base/GHC/Conc/Sync.hs b/libraries/base/GHC/Conc/Sync.hs
index d2fd02de69..aabe0b159e 100644
--- a/libraries/base/GHC/Conc/Sync.hs
+++ b/libraries/base/GHC/Conc/Sync.hs
@@ -41,7 +41,9 @@ module GHC.Conc.Sync
         , labelThread
         , labelThreadByteArray#
         , mkWeakThreadId
+        -- ** Queries
         , listThreads
+        , threadLabel
         , ThreadStatus(..), BlockReason(..)
         , threadStatus
         , threadCapability
@@ -626,6 +628,19 @@ threadCapability (ThreadId t) = IO $ \s ->
    case threadStatus# t s of
      (# s', _, cap#, locked# #) -> (# s', (I# cap#, isTrue# (locked# /=# 0#)) #)
 
+-- | Query the label of thread, returning 'Nothing' if the
+-- thread's label has not been set.
+--
+-- @since 4.18
+threadLabel :: ThreadId -> IO (Maybe String)
+threadLabel (ThreadId t) = IO $ \s ->
+    case threadLabel# t s of
+      (# s', 1#, lbl #) ->
+          let lbl' = utf8DecodeByteArray# lbl
+          in (# s', Just lbl' #)
+      (# s', 0#, _ #) -> (# s', Nothing #)
+      _ -> error "threadLabel: impossible"
+
 -- | Make a weak pointer to a 'ThreadId'.  It can be important to do
 -- this if you want to hold a reference to a 'ThreadId' while still
 -- allowing the thread to receive the @BlockedIndefinitely@ family of
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index 01981a0f82..11f6d94151 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -19,6 +19,8 @@
     a thread label by way of a `ByteArray#` containing a UTF-8-encoded string.
     The old `GHC.Conc.Sync.labelThread` is now implemented in terms of this
     function.
+  * `GHC.Conc.Sync.threadLabel` was added, allowing the user to query the label
+    of a given `ThreadId`.
 
 ## 4.17.0.0 *TBA*
 
diff --git a/libraries/base/tests/listThreads.hs b/libraries/base/tests/listThreads.hs
index 00e2178297..398afd0cc7 100644
--- a/libraries/base/tests/listThreads.hs
+++ b/libraries/base/tests/listThreads.hs
@@ -1,4 +1,5 @@
 import Control.Concurrent
+import Data.List (sort)
 import GHC.Conc.Sync
 
 dummyThread :: MVar () -> Int -> IO ()
@@ -16,7 +17,7 @@ main = do
 
   mapM_ mkThread [0..100]
   threads <- listThreads
-  -- TODO: Check labels
   print $ length threads
+  print . sort =<< mapM threadLabel threads
   putMVar mvar ()
 
diff --git a/libraries/base/tests/listThreads.stdout b/libraries/base/tests/listThreads.stdout
index 257e563266..e7ed0d1d6c 100644
--- a/libraries/base/tests/listThreads.stdout
+++ b/libraries/base/tests/listThreads.stdout
@@ -1 +1,2 @@
 102
+[Nothing,Just "thread-0",Just "thread-1",Just "thread-10",Just "thread-100",Just "thread-11",Just "thread-12",Just "thread-13",Just "thread-14",Just "thread-15",Just "thread-16",Just "thread-17",Just "thread-18",Just "thread-19",Just "thread-2",Just "thread-20",Just "thread-21",Just "thread-22",Just "thread-23",Just "thread-24",Just "thread-25",Just "thread-26",Just "thread-27",Just "thread-28",Just "thread-29",Just "thread-3",Just "thread-30",Just "thread-31",Just "thread-32",Just "thread-33",Just "thread-34",Just "thread-35",Just "thread-36",Just "thread-37",Just "thread-38",Just "thread-39",Just "thread-4",Just "thread-40",Just "thread-41",Just "thread-42",Just "thread-43",Just "thread-44",Just "thread-45",Just "thread-46",Just "thread-47",Just "thread-48",Just "thread-49",Just "thread-5",Just "thread-50",Just "thread-51",Just "thread-52",Just "thread-53",Just "thread-54",Just "thread-55",Just "thread-56",Just "thread-57",Just "thread-58",Just "thread-59",Just "thread-6",Just "thread-60",Just "thread-61",Just "thread-62",Just "thread-63",Just "thread-64",Just "thread-65",Just "thread-66",Just "thread-67",Just "thread-68",Just "thread-69",Just "thread-7",Just "thread-70",Just "thread-71",Just "thread-72",Just "thread-73",Just "thread-74",Just "thread-75",Just "thread-76",Just "thread-77",Just "thread-78",Just "thread-79",Just "thread-8",Just "thread-80",Just "thread-81",Just "thread-82",Just "thread-83",Just "thread-84",Just "thread-85",Just "thread-86",Just "thread-87",Just "thread-88",Just "thread-89",Just "thread-9",Just "thread-90",Just "thread-91",Just "thread-92",Just "thread-93",Just "thread-94",Just "thread-95",Just "thread-96",Just "thread-97",Just "thread-98",Just "thread-99"]
diff --git a/libraries/ghc-prim/changelog.md b/libraries/ghc-prim/changelog.md
index 0648050834..ced48ff938 100644
--- a/libraries/ghc-prim/changelog.md
+++ b/libraries/ghc-prim/changelog.md
@@ -18,6 +18,9 @@
   ```
   Where the `ByteArray#` must contain a UTF-8-encoded string.
 
+- The `threadLabel#` primop was added, allowing the user to query the label of
+  a given `ThreadId#`.
+
 ## 0.9.0
 
 - Shipped with GHC 9.4.1
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 503797f556..96dd02e359 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -1096,6 +1096,17 @@ stg_isCurrentThreadBoundzh (/* no args */)
     return (r);
 }
 
+stg_threadLabelzh ( gcptr tso )
+{
+    W_ r;
+    r = StgTSO_label(tso);
+    if (r == 0) {
+        return (0, 0);
+    } else {
+        return (1, r);
+    }
+}
+
 stg_threadStatuszh ( gcptr tso )
 {
     W_ why_blocked;
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 40193a25a2..d5efc0cbad 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -722,6 +722,7 @@ extern char **environ;
       SymI_HasDataProto(stg_isEmptyMVarzh)                                  \
       SymI_HasDataProto(stg_killThreadzh)                                   \
       SymI_HasDataProto(stg_listThreadszh)                                  \
+      SymI_HasDataProto(stg_threadLabelzh)                                  \
       SymI_HasProto(loadArchive)                                        \
       SymI_HasProto(loadObj)                                            \
       SymI_HasProto(purgeObj)                                           \
diff --git a/rts/include/stg/MiscClosures.h b/rts/include/stg/MiscClosures.h
index 6b784306fb..d439730c72 100644
--- a/rts/include/stg/MiscClosures.h
+++ b/rts/include/stg/MiscClosures.h
@@ -532,6 +532,7 @@ RTS_FUN_DECL(stg_unmaskAsyncExceptionszh);
 RTS_FUN_DECL(stg_myThreadIdzh);
 RTS_FUN_DECL(stg_labelThreadzh);
 RTS_FUN_DECL(stg_isCurrentThreadBoundzh);
+RTS_FUN_DECL(stg_threadLabelzh);
 RTS_FUN_DECL(stg_threadStatuszh);
 RTS_FUN_DECL(stg_listThreadszh);