Fixes to thread instrumentation.

1 files changed, 24 insertions, 10 deletions

diff --git a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
index 2c35419e..afcb81ac 100644
--- a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
+++ b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
@@ -15,6 +15,7 @@ import qualified Control.Concurrent.Lifted as Raw
 import Control.Concurrent.Lifted hiding (fork,forkOS)
 import Control.Exception (fromException)
 import Control.Monad.Trans.Control
+import GHC.Stack
 import System.IO.Unsafe
 import qualified Data.Map.Strict  as Map
 import Control.Exception.Lifted
@@ -44,39 +45,48 @@ globals = unsafePerformIO $ newMVar $ GlobalState
     }
 {-# NOINLINE globals #-}
 
-forkLabeled :: String -> IO () -> IO ThreadId
+forkLabeled :: HasCallStack => String -> IO () -> IO ThreadId
 forkLabeled lbl action = do
-    t <- forkIO action
+    t <- instrumented GHC.forkIO action
     labelThread t lbl
     return t
 {-# INLINE forkLabeled #-}
 
-forkOSLabeled :: String -> IO () -> IO ThreadId
+forkOSLabeled :: HasCallStack => String -> IO () -> IO ThreadId
 forkOSLabeled lbl action = do
-    t <- forkOS action
+    t <- instrumented Raw.forkOS action
     labelThread t lbl
     return t
 {-# INLINE forkOSLabeled #-}
 
-forkIO :: IO () -> IO ThreadId
+forkIO :: HasCallStack => IO () -> IO ThreadId
 forkIO = instrumented GHC.forkIO
 {-# INLINE forkIO #-}
 
-forkOS :: MonadBaseControl IO m => m () -> m ThreadId
+forkOS :: ( HasCallStack, MonadBaseControl IO m ) => m () -> m ThreadId
 forkOS = instrumented Raw.forkOS
 {-# INLINE forkOS #-}
 
-fork :: MonadBaseControl IO m => m () -> m ThreadId
+fork :: ( HasCallStack, MonadBaseControl IO m ) => m () -> m ThreadId
 fork = instrumented Raw.fork
 {-# INLINE fork #-}
 
-instrumented :: MonadBaseControl IO m =>
+shortCallStack :: [([Char], SrcLoc)] -> String
+shortCallStack []          = ""
+shortCallStack ((_,loc):_) = (srcLocFile loc) ++ ":" ++ show (srcLocStartLine loc)
+
+defaultLabel :: CallStack -> String
+defaultLabel stack = case getCallStack stack of
+    _ : sites  -> shortCallStack sites
+    sites -> shortCallStack sites
+
+instrumented :: ( HasCallStack, MonadBaseControl IO m ) =>
                 (m () -> m ThreadId) -> m () -> m ThreadId
 instrumented rawFork action = do
     t <- rawFork $ do
         tid <- myThreadId
         tm <- liftBase getCurrentTime
-        bracket_ (modifyThreads $! Map.insert tid (PerThread "" tm))
+        bracket_ (modifyThreads $! \ts -> Map.union ts (Map.singleton tid (PerThread (defaultLabel callStack) tm)))
                  (return ())
             $ do catch action $ \e -> case fromException e of
                     Just ThreadKilled -> return ()
@@ -99,12 +109,16 @@ instrumented rawFork action = do
                                     reportException g l)
                  -- Remove the thread only if it terminated normally or was killed.
                  modifyThreads $! Map.delete tid
+    liftBase $ labelThread t (defaultLabel callStack)
     return t
 
 labelThread :: ThreadId -> String -> IO ()
 labelThread tid s = do
     GHC.labelThread tid s
-    modifyThreads $! Map.adjust (\pt -> pt { lbl = s }) tid
+    tm <- liftBase getCurrentTime
+    let updateIt (Just pt) = Just $ pt { lbl = s }
+        updateIt Nothing   = Just $ PerThread s tm
+    modifyThreads $! Map.alter updateIt tid
 {-# INLINE labelThread #-}
 
 threadsInformation :: IO [(ThreadId,PerThread)]