1 files changed, 21 insertions, 12 deletions

diff --git a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
index 070745b6..2c35419e 100644
--- a/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
+++ b/lifted-concurrent/src/Control/Concurrent/Lifted/Instrument.hs
@@ -81,15 +81,22 @@ instrumented rawFork action = do
             $ 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
+                        bracket (takeMVar globals)
+                                (\g -> do
+                                    let l = concat [ show e
+                                                   , " ("
+                                                   , maybe "" lbl $ Map.lookup tid (threads g)
+                                                   , ")"
+                                                   ]
+                                    putMVar globals $! g { threads = Map.insert tid (PerThread l tm) $ threads g }
+                                    throwIO e)
+                                (\g -> do
+                                    let l = concat [ show e
+                                                   , " ("
+                                                   , maybe "" lbl $ Map.lookup tid (threads g)
+                                                   , ")"
+                                                   ]
+                                    reportException g l)
                  -- Remove the thread only if it terminated normally or was killed.
                  modifyThreads $! Map.delete tid
     return t
@@ -106,8 +113,10 @@ threadsInformation = do
     return $ Map.toList m
 
 
-modifyThreads :: MonadBase IO m => (Map.Map ThreadId PerThread -> Map.Map ThreadId PerThread) -> m ()
+modifyThreads :: MonadBaseControl 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
+    bracket (takeMVar globals)
+            (\g -> putMVar globals $! f' g)
+            (\g -> return ())