@@ -104,7 +104,7 @@ import Data.Foldable (forM_)
104104import Data.Typeable (Typeable)
105105import Control.Applicative ((<$>))
106106import Control.Exception (throw)
107−import Control.Monad (forever, forM, replicateM)
107+import Control.Monad (forever, forM, replicateM, when, replicateM_)
108108import Control.Monad.IO.Class (liftIO)
109109import Control.Concurrent (forkIO, threadDelay, ThreadId)
110110import Control.Concurrent.MVar (MVar, newMVar, readMVar, modifyMVar_)
@@ -112,26 +112,36 @@ import Control.Distributed.Process
112112 ( RemoteTable
113113 , NodeId
114114 , Process
115+ , ProcessId
115116 , WhereIsReply(..)
116117 , whereis
117118 , whereisRemoteAsync
119+ , registerRemoteAsync
118120 , reregisterRemoteAsync
119121 , getSelfPid
120122 , register
123+ , reregister
121124 , expect
122125 , nsendRemote
123126 , receiveWait
127+ , receiveTimeout
124128 , match
125129 , matchIf
126130 , processNodeId
127131 , monitorNode
132+ , monitor
128133 , unmonitor
129134 , NodeMonitorNotification(..)
135+ , ProcessRegistrationException
130136 , finally
131137 , newChan
132138 , receiveChan
133139 , nsend
134140 , SendPort
141+ , RegisterReply(..)
142+ , bracket
143+ , try
144+ , send
135145 )
136146import qualified Control.Distributed.Process.Node as Node
137147 ( LocalNode
@@ -157,7 +167,7 @@ data Backend = Backend {
157167 , findPeers :: Int -> IO [NodeId]
158168 -- | Make sure that all log messages are printed by the logger on the
159169 -- current node
160− , redirectLogsHere :: Process ()
170+ , redirectLogsHere :: [ProcessId] -> Process ()
161171 }
162172
163173data BackendState = BackendState {
@@ -244,12 +254,27 @@ peerDiscoveryDaemon backendState recv send = forever go
244254--------------------------------------------------------------------------------
245255
246256-- | Make sure that all log messages are printed by the logger on this node
247−apiRedirectLogsHere :: Backend -> Process ()
248−apiRedirectLogsHere backend = do
257+apiRedirectLogsHere :: Backend -> [ProcessId] -> Process ()
258+apiRedirectLogsHere _backend slavecontrollers = do
249259 mLogger <- whereis "logger"
260+ myPid <- getSelfPid
261+
250262 forM_ mLogger $ \logger -> do
251− nids <- liftIO $ findPeers backend 1000000
252− forM_ nids $ \nid -> reregisterRemoteAsync nid "logger" logger -- ignore async response
263+
264+ bracket
265+ (mapM monitor slavecontrollers)
266+ (mapM unmonitor)
267+ $ \_ -> do
268+
269+ -- fire off redirect requests
270+ forM_ slavecontrollers $ \pid -> send pid (RedirectLogsTo logger myPid)
271+
272+ -- Wait for the replies
273+ replicateM_ (length slavecontrollers) $ do
274+ receiveWait
275+ [ match (\(RedirectLogsReply from ok) -> return ())
276+ , match (\m@(NodeMonitorNotification {}) -> return ())
277+ ]
253278
254279--------------------------------------------------------------------------------
255280-- Slaves --
@@ -259,18 +284,31 @@ apiRedirectLogsHere backend = do
259284--
260285-- This datatype is not exposed; instead, we expose primitives for dealing
261286-- with slaves.
262−data SlaveControllerMsg =
263− SlaveTerminate
287+data SlaveControllerMsg
288+ = SlaveTerminate
289+ | RedirectLogsTo ProcessId ProcessId
264290 deriving (Typeable, Show)
265291
266292instance Binary SlaveControllerMsg where
267293 put SlaveTerminate = putWord8 0
294+ put (RedirectLogsTo a b) = do putWord8 1; put (a,b)
268295 get = do
269296 header <- getWord8
270297 case header of
271298 0 -> return SlaveTerminate
299+ 1 -> do (a,b) <- get; return (RedirectLogsTo a b)
272300 _ -> fail "SlaveControllerMsg.get: invalid"
273301
302+data RedirectLogsReply
303+ = RedirectLogsReply ProcessId Bool
304+ deriving (Typeable, Show)
305+
306+instance Binary RedirectLogsReply where
307+ put (RedirectLogsReply from ok) = put (from,ok)
308+ get = do
309+ (from,ok) <- get
310+ return (RedirectLogsReply from ok)
311+
274312-- | Calling 'slave' sets up a new local node and then waits. You start
275313-- processes on the slave by calling 'spawn' from other nodes.
276314--
@@ -292,41 +330,49 @@ slaveController = do
292330 msg <- expect
293331 case msg of
294332 SlaveTerminate -> return ()
333+ RedirectLogsTo loggerPid from -> do
334+ r <- try (reregister "logger" loggerPid)
335+ ok <- case (r :: Either ProcessRegistrationException ()) of
336+ Right _ -> return True
337+ Left _ -> do
338+ r <- try (register "logger" loggerPid)
339+ case (r :: Either ProcessRegistrationException ()) of
340+ Right _ -> return True
341+ Left _ -> return False
342+ pid <- getSelfPid
343+ send from (RedirectLogsReply pid ok)
344+ go
295345
296346-- | Terminate the slave at the given node ID
297347terminateSlave :: NodeId -> Process ()
298348terminateSlave nid = nsendRemote nid "slaveController" SlaveTerminate
299349
300350-- | Find slave nodes
301−findSlaves :: Backend -> Process [NodeId]
351+findSlaves :: Backend -> Process [ProcessId]
302352findSlaves backend = do
303353 nodes <- liftIO $ findPeers backend 1000000
304− -- Fire of asynchronous requests for the slave controller
305− refs <- forM nodes $ \nid -> do
306− whereisRemoteAsync nid "slaveController"
307− ref <- monitorNode nid
308− return (nid, ref)
309− -- Wait for the replies
310− catMaybes <$> replicateM (length nodes) (
311− receiveWait
312− [ matchIf (\(WhereIsReply label _) -> label == "slaveController")
313− (\(WhereIsReply _ mPid) ->
314− case mPid of
315− Nothing ->
316− return Nothing
317− Just pid -> do
318− let nid = processNodeId pid
319− Just ref = lookup nid refs
320− unmonitor ref
321− return (Just nid))
322− , match (\(NodeMonitorNotification {}) -> return Nothing)
323− ])
354+ -- Fire off asynchronous requests for the slave controller
355+
356+ bracket
357+ (mapM monitorNode nodes)
358+ (mapM unmonitor)
359+ $ \_ -> do
360+
361+ -- fire off whereis requests
362+ forM nodes $ \nid -> whereisRemoteAsync nid "slaveController"
363+
364+ -- Wait for the replies
365+ catMaybes <$> replicateM (length nodes) (
366+ receiveWait
367+ [ match (\(WhereIsReply "slaveController" mPid) -> return mPid)
368+ , match (\(NodeMonitorNotification {}) -> return Nothing)
369+ ])
324370
325371-- | Terminate all slaves
326372terminateAllSlaves :: Backend -> Process ()
327373terminateAllSlaves backend = do
328374 slaves <- findSlaves backend
329− forM_ slaves terminateSlave
375+ forM_ slaves $ \pid -> send pid SlaveTerminate
330376 liftIO $ threadDelay 1000000
331377
332378--------------------------------------------------------------------------------
@@ -353,8 +399,8 @@ startMaster backend proc = do
353399 node <- newLocalNode backend
354400 Node.runProcess node $ do
355401 slaves <- findSlaves backend
356− redirectLogsHere backend
357− proc slaves `finally` shutdownLogger
402+ redirectLogsHere backend slaves
403+ proc (map processNodeId slaves) `finally` shutdownLogger
358404
359405--
360406-- | shut down the logger process. This ensures that any pending