fixup! Improve exception handling in callLocal.

Author: qnikst

src/Control/Distributed/Process.hs

@@ -160,16 +160,20 @@ import Prelude hiding (catch)
 #endif
 
 import Control.Monad.IO.Class (liftIO)
-import Control.Applicative ((<$>), (<*))
+import Control.Applicative ((<$>))
 import Control.Monad.Reader (ask)
-import Control.Concurrent (myThreadId)
 import Control.Concurrent.MVar
-  ( newMVar
-  , newEmptyMVar
-  , tryTakeMVar
+  ( newEmptyMVar
   , takeMVar
   , putMVar
   )
+import Control.Concurrent.STM (atomically)
+import Control.Concurrent.STM.TMVar
+  ( TMVar
+  , newEmptyTMVarIO
+  , putTMVar
+  , takeTMVar
+  )
 import Control.Distributed.Static
   ( Closure
   , closure
@@ -310,8 +314,7 @@ import Control.Distributed.Process.Internal.Spawn
   , spawnSupervised
   , call
   )
-import Control.Exception (SomeException, throwIO, throwTo)
-import qualified Control.Exception as Exception
+import Control.Exception (SomeException, throwIO)
 
 -- INTERNAL NOTES
 --
@@ -410,27 +413,29 @@ spawnChannelLocal proc = do
 callLocal ::
      Process a  -- ^ Process to run
   -> Process a  -- ^ Value returned
-callLocal proc = mask $ \release -> do
-  lock  <- liftIO $ newMVar ()
-  tidMV <- liftIO $ newEmptyMVar
-  mv    <- liftIO newEmptyMVar
-  _ <- spawnLocal $ mask $ \release' -> do
-         liftIO $ myThreadId >>= putMVar tidMV
-         ep <- try $ release' $ proc <* liftIO (takeMVar lock)
-         liftIO $ putMVar mv ep
-  tid <- liftIO $ takeMVar tidMV
-  let fetchResult =
-        (takeMVar mv >>= either (throwIO :: SomeException -> IO a) return)
-        `Exception.catch`
-            (\e -> Exception.mask $ \release' -> do
-                     ml <- tryTakeMVar lock
-                     case ml of
-                        Nothing -> -- lock already taked by worker so we need
-                                   -- to rethrow exception now
-                          throwIO e
-                        Just{}  -> do
-                          throwTo tid (e::SomeException)
-                          putMVar lock ()
-                     release' fetchResult)
-  release $ liftIO $ fetchResult
-
+callLocal proc = do 
+  parent <- getSelfPid
+  mv    <- liftIO newEmptyTMVarIO :: Process (TMVar (Either SomeException a))
+  child <- spawnLocal $ mask $ \release -> do
+             link parent
+             ep <- try $ release $ proc
+             liftIO $ atomically $ putTMVar mv ep
+  withMonitor child $ do
+    fetchResult child mv `catch` (\e -> do exit child (show (e::SomeException))
+                                           waitForExit child 
+                                           liftIO $ throwIO e)
+  where
+    waitForExit child =
+        receiveWait
+            [ matchIf (\(ProcessMonitorNotification _ ch _) -> child == ch)
+                      (\_ -> return ())
+            ]
+    fetchResult child mv = do
+        receiveWait
+            [ matchSTM (takeTMVar mv)
+                       (\rs -> do waitForExit child -- avoid monitor events leak
+                                  liftIO $ either throwIO return rs)
+            , matchIf  (\(ProcessMonitorNotification _ ch _) -> child == ch)
+                       (\(ProcessMonitorNotification _ _ reason) ->
+                            fail $ "callLocal: remote process died: " ++ show reason)
+            ]