Updates to thread instrumentation.

1 files changed, 24 insertions, 5 deletions

diff --git a/src/Control/Concurrent/Lifted/Instrument.hs b/src/Control/Concurrent/Lifted/Instrument.hs
index 9a409569..7e4a7356 100644
--- a/src/Control/Concurrent/Lifted/Instrument.hs
+++ b/src/Control/Concurrent/Lifted/Instrument.hs
@@ -10,6 +10,7 @@ module Control.Concurrent.Lifted.Instrument
 
 import qualified Control.Concurrent.Lifted as Raw
 import Control.Concurrent.Lifted hiding (fork)
+import Control.Exception (fromException)
 import Control.Monad.Trans.Control
 import System.IO.Unsafe
 import qualified Data.Map.Strict  as Map
@@ -18,6 +19,9 @@ import Control.Monad.Base
 import qualified GHC.Conc as GHC
 import Data.Time()
 import Data.Time.Clock
+import System.IO
+import Control.Monad.IO.Class
+
 
 data PerThread = PerThread
  { lbl :: String
@@ -25,13 +29,15 @@ data PerThread = PerThread
  }
  deriving (Eq,Ord,Show)
 
-newtype GlobalState = GlobalState
- { threads :: Map.Map ThreadId PerThread
+data GlobalState = GlobalState
+ { threads :: !(Map.Map ThreadId PerThread)
+ , reportException :: String -> IO ()
  }
 
 globals :: MVar GlobalState
 globals = unsafePerformIO $ newMVar $ GlobalState
     { threads = Map.empty
+    , reportException = hPutStrLn stderr
     }
 {-# NOINLINE globals #-}
 
@@ -40,14 +46,27 @@ forkIO :: IO () -> IO ThreadId
 forkIO = fork
 {-# INLINE forkIO #-}
 
-fork :: MonadBaseControl IO m => m () -> m ThreadId
+fork :: (MonadBaseControl IO m, MonadIO m) => m () -> m ThreadId
 fork action = do
     t <- Raw.fork $ do
         tid <- myThreadId
         tm <- liftBase getCurrentTime
         bracket_ (modifyThreads $! Map.insert tid (PerThread "" tm))
-                 (modifyThreads $! Map.delete tid)
-                 action
+                 (return ())
+            $ do catch action $ \e -> case fromException e of
+                    Just ThreadKilled -> return ()
+                    Nothing           -> liftIO $ do
+                        g <- takeMVar globals
+                        let l = concat [ show e
+                                       , " ("
+                                       , maybe "" lbl $ Map.lookup tid (threads g)
+                                       , ")"
+                                       ]
+                        reportException g l
+                        putMVar globals $! g { threads = Map.insert tid (PerThread l tm) $ threads g }
+                        throwIO e
+                 -- Remove the thread only if it terminated normally or was killed.
+                 modifyThreads $! Map.delete tid
     return t
 
 labelThread :: ThreadId -> String -> IO ()