Added forKOS (and removed MonadIO constraint) from thread instrumentation.

1 files changed, 15 insertions, 5 deletions

diff --git a/src/Control/Concurrent/Lifted/Instrument.hs b/src/Control/Concurrent/Lifted/Instrument.hs
index 7e4a7356..5b3237cc 100644
--- a/src/Control/Concurrent/Lifted/Instrument.hs
+++ b/src/Control/Concurrent/Lifted/Instrument.hs
@@ -2,6 +2,7 @@
 module Control.Concurrent.Lifted.Instrument
     ( module Control.Concurrent.Lifted
     , forkIO
+    , forkOS
     , fork
     , labelThread
     , threadsInformation
@@ -9,7 +10,7 @@ module Control.Concurrent.Lifted.Instrument
     ) where
 
 import qualified Control.Concurrent.Lifted as Raw
-import Control.Concurrent.Lifted hiding (fork)
+import Control.Concurrent.Lifted hiding (fork,forkOS)
 import Control.Exception (fromException)
 import Control.Monad.Trans.Control
 import System.IO.Unsafe
@@ -46,16 +47,25 @@ forkIO :: IO () -> IO ThreadId
 forkIO = fork
 {-# INLINE forkIO #-}
 
-fork :: (MonadBaseControl IO m, MonadIO m) => m () -> m ThreadId
-fork action = do
-    t <- Raw.fork $ do
+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           -> liftIO $ do
+                    Nothing           -> liftBase $ do
                         g <- takeMVar globals
                         let l = concat [ show e
                                        , " ("