@@ -47,6 +47,7 @@ module Control.Distributed.Process.Management
4747 , MxAgentId(..)
4848 , MxAgent()
4949 , mxAgent
50+ , mxAgentWithFinalize
5051 , MxSink
5152 , mxSink
5253 , mxGetId
@@ -77,6 +78,7 @@ import Control.Distributed.Process.Internal.Primitives
7778 , receiveWait
7879 , matchChan
7980 , unwrapMessage
81+ , onException
8082 )
8183import Control.Distributed.Process.Internal.Types
8284 ( Process
@@ -197,17 +199,27 @@ mxSink act msg = do
197199data MxPipeline s =
198200 MxPipeline
199201 {
200− current :: !(MxSink s)
201− , next :: !(MxPipeline s)
202+ current :: !(MxSink s)
203+ , next :: !(MxPipeline s)
202204 } | MxStop
203205
204206-- | Activates a new agent.
205207--
206−mxAgent :: MxAgentId
208+mxAgent :: MxAgentId -> s -> [MxSink s] -> Process ProcessId
209+mxAgent mxId st hs = mxAgentWithFinalize mxId st hs $ return ()
210+
211+-- | Activates a new agent. This variant takes a /finalizer/ expression,
212+-- that is run once the agent shuts down (even in case of failure/exceptions).
213+-- The /finalizer/ expression runs in the mx monad - @MxAgent s ()@ - such
214+-- that the agent's internal state remains accessible to the shutdown/cleanup
215+-- code.
216+--
217+mxAgentWithFinalize :: MxAgentId
207218 -> s
208219 -> [MxSink s]
220+ -> MxAgent s ()
209221 -> Process ProcessId
210−mxAgent mxId initState handlers = do
222+mxAgentWithFinalize mxId initState handlers dtor = do
211223 node <- processNode <$> ask
212224 pid <- liftIO $ mxNew (localEventBus node) $ start
213225 return pid
@@ -220,15 +232,21 @@ mxAgent mxId initState handlers = do
220232 -- liftIO $ putStrLn $ "registration == " ++ (show p)
221233 tablePid <- receiveWait [ matchChan rp (\(p :: ProcessId) -> return p) ]
222234 -- liftIO $ putStrLn "starting agent listener..."
223− runAgent handlers recvTChan $ MxAgentState mxId sendTChan tablePid initState
224−
225− runAgent hs c s = do
226− msg <- (liftIO $ atomically $ readTChan c)
227− (action, state) <- runPipeline msg s $ pipeline hs
228− case action of
229− MxAgentReady -> runAgent hs c state
230− MxAgentDeactivate _ -> {- TODO: log r -} return ()
231−-- MxAgentBecome h' -> runAgent h' c state
235+ runAgent dtor handlers recvTChan $ MxAgentState mxId sendTChan tablePid initState
236+
237+ runAgent eh hs c s =
238+ runAgentWithFinalizer eh hs c s `onException` runAgentFinalizer eh s
239+
240+ runAgentWithFinalizer eh' hs' c' s' = do
241+ msg <- (liftIO $ atomically $ readTChan c')
242+ (action, state) <- runPipeline msg s' $ pipeline hs'
243+ case action of
244+ MxAgentReady -> runAgent eh' hs' c' state
245+ MxAgentDeactivate _ -> runAgentFinalizer eh' state
246+-- MxAgentBecome h' -> runAgent h' c state
247+
248+ runAgentFinalizer :: MxAgent s () -> MxAgentState s -> Process ()
249+ runAgentFinalizer f s = ST.runStateT (unAgent f) s >>= return . fst
232250
233251 pipeline :: forall s . [MxSink s] -> MxPipeline s
234252 pipeline [] = MxStop