rio-process-pool-1.0.1: src/RIO/ProcessPool/Broker.hs
{-# LANGUAGE Strict #-}
-- | A broker extracts a /key/ value from incoming messages
-- and creates, keeps and destroys a /resource/ for each key.
--
-- The demultiplexed messages and their resources are passed to
-- a custom 'MessageHandler'/
--
-- The user provides a 'Demultiplexer' is a pure function that
-- returns a key for the resource associated
-- to the message and potientially changes the
-- message.
--
-- The demultiplexer may also return a value indicating that
-- a new resource must be created, or that a message
-- shall be ignored.
--
-- The broker is run in a seperate process using 'async'.
-- The usual way to stop a broker is to 'cancel' it.
--
-- When cancelling a broker, the resource cleanup
-- actions for all resources will be called with
-- async exceptions masked.
--
-- In order to prevent the resource map filling up with
-- /dead/ resources, the user of this module has to ensure
-- that whenever a resource is not required anymore, a message
-- will be sent to the broker, that will cause the 'MessageHandler'
-- to be executed for the resource, which will in turn return,
-- return 'RemoveResource'.
module RIO.ProcessPool.Broker
( spawnBroker,
BrokerConfig (..),
BrokerResult (..),
ResourceCreator,
Demultiplexer,
ResourceCleaner,
MessageHandler,
Multiplexed (..),
ResourceUpdate (..),
)
where
import qualified Data.Map.Strict as Map
import RIO
import UnliftIO.MessageBox
( IsMessageBox (Input, newInput, receive),
IsMessageBoxArg (MessageBox, newMessageBox),
)
-- | Spawn a broker with a new 'MessageBox',
-- and return its message 'Input' channel as well as
-- the 'Async' handle of the spawned process, needed to
-- stop the broker process.
--
-- * @k@ is the /key/ for the resource associated to an incoming
-- message
-- * @w'@ is the type of incoming messages.
-- * @w@ is the type of the demultiplexed messages.
-- * @a@ specifies the resource type.
-- * @m@ is the base monad
spawnBroker ::
forall brokerBoxArg k w' w a m.
( HasLogFunc m,
Ord k,
Display k,
IsMessageBoxArg brokerBoxArg
) =>
brokerBoxArg ->
BrokerConfig k w' w a m ->
RIO
m
( Either
SomeException
( Input (MessageBox brokerBoxArg) w',
Async BrokerResult
)
)
spawnBroker brokerBoxArg config = do
brokerA <- async $ do
mBrokerBox <-
tryAny
( do
b <- newMessageBox brokerBoxArg
i <- newInput b
return (b, i)
)
case mBrokerBox of
Left er -> return (Left er)
Right (brokerBox, brokerInp) -> do
aInner <- mask_ $
asyncWithUnmask $ \unmaskInner ->
brokerLoop unmaskInner brokerBox config Map.empty
return (Right (brokerInp, aInner))
join <$> waitCatch brokerA
-- | This is just what the 'Async' returned from
-- 'spawnBroker' returns, it's current purpose is to
-- make code easier to read.
--
-- Instead of some @Async ()@ that could be anything,
-- there is @Async BrokerResult@.
data BrokerResult = MkBrokerResult
deriving stock (Show, Eq)
-- | The broker configuration, used by 'spawnBroker'.
--
-- * @k@ is the /key/ for the resource associated to an incoming
-- message
-- * @w'@ is the type of incoming messages.
-- * @w@ is the type of the demultiplexed messages.
-- * @a@ specifies the resource type.
-- * @m@ is the base monad
data BrokerConfig k w' w a m = MkBrokerConfig
{ demultiplexer :: !(Demultiplexer w' k w),
messageDispatcher :: !(MessageHandler k w a m),
resourceCreator :: !(ResourceCreator k w a m),
resourceCleaner :: !(ResourceCleaner k a m)
}
-- | User supplied callback to extract the key and the 'Multiplexed'
-- from a message.
-- (Sync-) Exceptions thrown from this function are caught and lead
-- to dropping of the incoming message, while the broker continues.
--
-- * @k@ is the /key/ for the resource associated to an incoming
-- message
-- * @w'@ is the type of incoming messages.
-- * @w@ is the type of the demultiplexed messages.
type Demultiplexer w' k w = w' -> Multiplexed k w
-- | User supplied callback to use the 'Multiplexed' message and
-- the associated resource.
-- (Sync-) Exceptions thrown from this function are caught and lead
-- to immediate cleanup of the resource but the broker continues.
--
-- * Type @k@ is the /key/ for the resource associated to an incoming
-- message
-- * Type @w@ is the type of incoming, demultiplexed, messages.
-- * Type @a@ specifies the resource type.
-- * Type @m@ is the base monad
type MessageHandler k w a m = k -> w -> a -> RIO m (ResourceUpdate a)
-- | This value indicates in what state a worker is in after the
-- 'MessageHandler' action was executed.
data ResourceUpdate a
= -- | The resources is still required.
KeepResource
| -- | The resource is still required but must be updated.
UpdateResource a
| -- | The resource is obsolete and can
-- be removed from the broker.
-- The broker will call 'ResourceCleaner' either
-- on the current, or an updated resource value.
RemoveResource !(Maybe a)
-- | The action that the broker has to take for in incoming message.
--
-- * @k@ is the /key/ for the resource associated to an incoming
-- message
-- * @w@ is the type of the demultiplexed messages.
data Multiplexed k w
= -- | The message is an initialization message, that requires the
-- creation of a new resouce for the given key.
-- When the resource is created, then /maybe/ additionally
-- a message will also be dispatched.
Initialize k !(Maybe w)
| -- | Dispatch a message using an existing resource.
-- Silently ignore if no resource for the key exists.
Dispatch k w
-- deriving stock (Show)
-- | User supplied callback to create and initialize a resource.
-- (Sync-) Exceptions thrown from this function are caught,
-- and the broker continues.
--
-- * @k@ is the /key/ for the resource associated to an incoming
-- message
-- * @w@ is the type of the demultiplexed messages.
-- * @a@ specifies the resource type.
-- * @m@ is the monad of the returned action.
type ResourceCreator k w a m = k -> Maybe w -> RIO m a
-- | User supplied callback called _with exceptions masked_
-- when the 'MessageHandler' returns 'RemoveResource'
-- (Sync-) Exceptions thrown from this function are caught,
-- and do not prevent the removal of the resource, also the
-- broker continues.
--
-- * @k@ is the /key/ for the resource associated to an incoming
-- message
-- * @a@ specifies the resource type.
-- * @m@ is the monad of the returned action.
type ResourceCleaner k a m = k -> a -> RIO m ()
type BrokerState k a = Map k a
{-# NOINLINE brokerLoop #-}
brokerLoop ::
( HasLogFunc m,
Ord k,
Display k,
IsMessageBox msgBox
) =>
(forall x. RIO m x -> RIO m x) ->
msgBox w' ->
BrokerConfig k w' w a m ->
BrokerState k a ->
RIO m BrokerResult
brokerLoop unmask brokerBox config brokerState =
withException
( unmask (receive brokerBox)
>>= traverse (tryAny . onIncoming unmask config brokerState)
)
( \(ex :: SomeException) -> do
case fromException ex of
Just (_cancelled :: AsyncCancelled) ->
logDebug "broker loop: cancelled"
_ ->
logError
( "broker loop: exception while \
\receiving and dispatching messages: "
<> display ex
)
cleanupAllResources config brokerState
)
>>= maybe
( do
logError "broker loop: failed to receive next message"
cleanupAllResources config brokerState
return MkBrokerResult
)
( \res -> do
next <-
either
( \err -> do
logWarn
( "broker loop: Handling the last message\
\ caused an exception:"
<> display err
)
return brokerState
)
return
res
brokerLoop unmask brokerBox config next
)
{-# NOINLINE onIncoming #-}
onIncoming ::
(Ord k, HasLogFunc m, Display k) =>
(forall x. RIO m x -> RIO m x) ->
BrokerConfig k w' w a m ->
BrokerState k a ->
w' ->
RIO m (BrokerState k a)
onIncoming unmask config brokerState w' =
case demultiplexer config w' of
Initialize k mw ->
onInitialize unmask k config brokerState mw
Dispatch k w ->
onDispatch unmask k w config brokerState
onInitialize ::
(Ord k, HasLogFunc m, Display k) =>
(forall x. RIO m x -> RIO m x) ->
k ->
BrokerConfig k w' w a m ->
BrokerState k a ->
Maybe w ->
RIO m (BrokerState k a)
onInitialize unmask k config brokerState mw =
case Map.lookup k brokerState of
Just _ -> do
logError
( "cannot initialize a new worker, a worker with that ID exists: "
<> display k
)
return brokerState
Nothing ->
tryAny (unmask (resourceCreator config k mw))
>>= either
( \err -> do
logError
( "the resource creator for worker "
<> display k
<> " threw an exception: "
<> display err
)
return brokerState
)
( \res ->
let brokerState1 = Map.insert k res brokerState
in case mw of
Nothing ->
return brokerState1
Just w ->
onException
(onDispatch unmask k w config brokerState1)
( do
logError
( "exception while dispatching the "
<> "post-initialization message for worker: "
<> display k
)
resourceCleaner config k res
)
)
onDispatch ::
(Ord k, HasLogFunc m, Display k) =>
(forall x. RIO m x -> RIO m x) ->
k ->
w ->
BrokerConfig k w' w a m ->
BrokerState k a ->
RIO m (BrokerState k a)
onDispatch unmask k w config brokerState =
maybe notFound dispatch (Map.lookup k brokerState)
where
notFound = do
logWarn
( "cannot dispatch message, worker not found: "
<> display k
)
return brokerState
dispatch res =
tryAny (unmask (messageDispatcher config k w res))
>>= either
( \err -> do
logError
( "the message dispatcher callback for worker "
<> display k
<> " threw: "
<> display err
)
cleanupResource
k
config
brokerState
)
( \case
KeepResource ->
return brokerState
UpdateResource newRes ->
return (Map.insert k newRes brokerState)
RemoveResource mNewRes ->
cleanupResource
k
config
( maybe
brokerState
( \newRes ->
Map.insert k newRes brokerState
)
mNewRes
)
)
cleanupAllResources ::
BrokerConfig k w' w a m ->
BrokerState k a ->
RIO m ()
cleanupAllResources config brokerState =
traverse_
( uncurry
(tryResourceCleaner config)
)
(Map.assocs brokerState)
cleanupResource ::
(Ord k) =>
k ->
BrokerConfig k w' w a m ->
Map k a ->
RIO m (Map k a)
cleanupResource k config brokerState = do
traverse_ (tryResourceCleaner config k) (Map.lookup k brokerState)
return (Map.delete k brokerState)
tryResourceCleaner ::
BrokerConfig k w' w a m ->
k ->
a ->
RIO m ()
tryResourceCleaner config k res = do
void $ tryAny (resourceCleaner config k res)