Add a primop to query the label of a thread

7 files changed, 33 insertions, 2 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index f564395342..4ba996342b 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -2999,6 +2999,7 @@ primop  MyThreadIdOp "myThreadId#" GenPrimOp
 
 primop LabelThreadOp "labelThread#" GenPrimOp
    ThreadId# -> Addr# -> State# RealWorld -> State# RealWorld
+   {Set the label of the given thread.}
    with
    has_side_effects = True
    out_of_line      = True
@@ -3015,6 +3016,12 @@ primop  NoDuplicateOp "noDuplicate#" GenPrimOp
    out_of_line = True
    has_side_effects = True
 
+primop GetThreadLabelOp "threadLabel#" GenPrimOp
+   ThreadId# -> State# RealWorld -> (# State# RealWorld, Addr# #)
+   {Get the label of the given thread. Returns NULL if not set.}
+   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 b7194c3d0f..84b324997a 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1634,6 +1634,7 @@ emitPrimOp dflags primop = case primop of
   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 3e150cfe9f..72ae0c3141 100644
--- a/libraries/base/GHC/Conc/Sync.hs
+++ b/libraries/base/GHC/Conc/Sync.hs
@@ -39,7 +39,9 @@ module GHC.Conc.Sync
         , yield
         , labelThread
         , mkWeakThreadId
+        -- ** Queries
         , listThreads
+        , threadLabel
         , ThreadStatus(..), BlockReason(..)
         , threadStatus
         , threadCapability
@@ -626,6 +628,17 @@ 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.
+threadLabel :: ThreadId -> IO (Maybe String)
+threadLabel (ThreadId t) = IO $ \s ->
+    case threadLabel# t s of
+      (# s', lbl #) ->
+          let lbl'
+                | Ptr lbl == nullPtr = Nothing
+                | otherwise          = Just $ unpackCStringUtf8# lbl
+          in (# s', lbl' #)
+
 -- | 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/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/rts/PrimOps.cmm b/rts/PrimOps.cmm
index febd536db3..392ef6074b 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -1122,6 +1122,13 @@ stg_isCurrentThreadBoundzh (/* no args */)
     return (r);
 }
 
+stg_threadLabelzh ( gcptr tso )
+{
+    W_ r;
+    (r) = ccall lookupThreadLabel (StgTSO_id(tso));
+    return (r);
+}
+
 stg_threadStatuszh ( gcptr tso )
 {
     W_ why_blocked;
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 518fbde28e..aad5669367 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -682,13 +682,14 @@
       SymI_HasProto(initLinker_)                                        \
       SymI_HasProto(stg_unpackClosurezh)                                \
       SymI_HasProto(stg_closureSizzezh)                                 \
-      SymI_HasProto(stg_whereFromzh)                                 \
+      SymI_HasProto(stg_whereFromzh)                                    \
       SymI_HasProto(stg_getApStackValzh)                                \
       SymI_HasProto(stg_getSparkzh)                                     \
       SymI_HasProto(stg_numSparkszh)                                    \
       SymI_HasProto(stg_isCurrentThreadBoundzh)                         \
       SymI_HasProto(stg_isEmptyMVarzh)                                  \
       SymI_HasProto(stg_killThreadzh)                                   \
+      SymI_NeedsProto(stg_threadLabelzh)                                \
       SymI_HasProto(stg_listThreadszh)                                  \
       SymI_HasProto(loadArchive)                                        \
       SymI_HasProto(loadObj)                                            \