Thread instrumentation: Avoid unlikely leak on race condition.

1 files changed, 14 insertions, 3 deletions

diff --git a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
index afcb81ac..a0bb7dc5 100644
--- a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
+++ b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
@@ -83,6 +83,7 @@ defaultLabel stack = case getCallStack stack of
 instrumented :: ( HasCallStack, MonadBaseControl IO m ) =>
                 (m () -> m ThreadId) -> m () -> m ThreadId
 instrumented rawFork action = do
+    mvar <- newEmptyMVar
     t <- rawFork $ do
         tid <- myThreadId
         tm <- liftBase getCurrentTime
@@ -98,6 +99,7 @@ instrumented rawFork action = do
                                                    , maybe "" lbl $ Map.lookup tid (threads g)
                                                    , ")"
                                                    ]
+                                    foldr seq (return ()) l
                                     putMVar globals $! g { threads = Map.insert tid (PerThread l tm) $ threads g }
                                     throwIO e)
                                 (\g -> do
@@ -108,17 +110,26 @@ instrumented rawFork action = do
                                                    ]
                                     reportException g l)
                  -- Remove the thread only if it terminated normally or was killed.
+                 takeMVar mvar
                  modifyThreads $! Map.delete tid
-    liftBase $ labelThread t (defaultLabel callStack)
+    liftBase $ labelThread_ t (defaultLabel callStack)
+    putMVar mvar ()
     return t
 
-labelThread :: ThreadId -> String -> IO ()
-labelThread tid s = do
+labelThread_ :: ThreadId -> String -> IO ()
+labelThread_ tid s = do
+    foldr seq (return ()) s
     GHC.labelThread tid s
     tm <- liftBase getCurrentTime
     let updateIt (Just pt) = Just $ pt { lbl = s }
         updateIt Nothing   = Just $ PerThread s tm
     modifyThreads $! Map.alter updateIt tid
+
+labelThread :: ThreadId -> String -> IO ()
+labelThread tid s = do
+    foldr seq (return ()) s
+    GHC.labelThread tid s
+    modifyThreads $! Map.adjust (\pt -> pt { lbl = s }) tid
 {-# INLINE labelThread #-}
 
 threadsInformation :: IO [(ThreadId,PerThread)]