1 files changed, 98 insertions, 0 deletions

diff --git a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
new file mode 100644
index 00000000..fc3b6369
--- /dev/null
+++ b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
@@ -0,0 +1,98 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Control.Concurrent.Lifted.Instrument
+    ( module Control.Concurrent.Lifted
+    , forkIO
+    , forkOS
+    , fork
+    , labelThread
+    , threadsInformation
+    , PerThread(..)
+    ) where
+
+import qualified Control.Concurrent.Lifted as Raw
+import Control.Concurrent.Lifted hiding (fork,forkOS)
+import Control.Exception (fromException)
+import Control.Monad.Trans.Control
+import System.IO.Unsafe
+import qualified Data.Map.Strict  as Map
+import Control.Exception.Lifted
+import Control.Monad.Base
+import qualified GHC.Conc as GHC
+import Data.Time()
+import Data.Time.Clock
+import DPut
+import DebugTag
+
+
+data PerThread = PerThread
+ { lbl :: String
+ , startTime :: UTCTime
+ }
+ deriving (Eq,Ord,Show)
+
+data GlobalState = GlobalState
+ { threads :: !(Map.Map ThreadId PerThread)
+ , reportException :: String -> IO ()
+ }
+
+globals :: MVar GlobalState
+globals = unsafePerformIO $ newMVar $ GlobalState
+    { threads = Map.empty
+    , reportException = dput XMisc
+    }
+{-# NOINLINE globals #-}
+
+
+forkIO :: IO () -> IO ThreadId
+forkIO = instrumented GHC.forkIO
+{-# INLINE forkIO #-}
+
+forkOS :: MonadBaseControl IO m => m () -> m ThreadId
+forkOS = instrumented Raw.forkOS
+{-# INLINE forkOS #-}
+
+fork :: MonadBaseControl IO m => m () -> m ThreadId
+fork = instrumented Raw.fork
+{-# INLINE fork #-}
+
+instrumented :: 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))
+                 (return ())
+            $ do catch action $ \e -> case fromException e of
+                    Just ThreadKilled -> return ()
+                    Nothing           -> liftBase $ 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 ()
+labelThread tid s = do
+    GHC.labelThread tid s
+    modifyThreads $! Map.adjust (\pt -> pt { lbl = s }) tid
+{-# INLINE labelThread #-}
+
+threadsInformation :: IO [(ThreadId,PerThread)]
+threadsInformation = do
+    m <- threads <$> readMVar globals
+    return $ Map.toList m
+
+
+modifyThreads :: MonadBase IO m => (Map.Map ThreadId PerThread -> Map.Map ThreadId PerThread) -> m ()
+modifyThreads f = do
+    g <- takeMVar globals
+    let f' st = st { threads = f (threads st) }
+    putMVar globals $! f' g