stomp-patterns-0.5.0: Network/Mom/Stompl/Patterns/Registry.hs
{-# Language BangPatterns #-}
module Registry
where
import Types
import Network.Mom.Stompl.Client.Queue
import System.Timeout
import Data.Time
import Data.Char (isDigit, toUpper)
import Data.List (nub)
import Data.Maybe (fromMaybe)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Sequence (Seq, (|>), (<|), ViewL(..))
import qualified Data.Sequence as S
import Data.Foldable (toList)
import Codec.MIME.Type (nullType)
import Control.Exception (throwIO, catches)
import Control.Concurrent
import Control.Monad (forever)
import Control.Applicative ((<$>))
-----------------------------------------------------------------------
-- | JobType: Service, Task or Topic
-----------------------------------------------------------------------
data JobType = Service | Task | Topic
deriving (Eq, Show)
-----------------------------------------------------------------------
-- | Safe read method for JobType
-----------------------------------------------------------------------
readJobType :: String -> Maybe JobType
readJobType s =
case map toUpper s of
"SERVICE" -> Just Service
"TASK" -> Just Task
"TOPIC" -> Just Topic
_ -> Nothing
------------------------------------------------------------------------
-- | A helper that shall ease the use of the registers.
-- A registry to which a call wants to connect is described as
--
-- * The 'QName' through which the registry receives requests;
--
-- * The 'Timeout' in microseconds, /i.e./ the time the caller
-- will wait before the request fails;
--
-- * A triple of heartbeat specifications:
-- the /best/ value, /i.e./
-- the rate at which the caller
-- prefers to send heartbeats,
-- the /minimum/ rate at which the caller
-- can accept to send heartbeats,
-- the /maximum/ rate at which the caller
-- can accept to send heartbeats.
-- Note that all these values are in milliseconds!
------------------------------------------------------------------------
type RegistryDesc = (QName, Int, (Int, Int, Int))
------------------------------------------------------------------------
-- | Connect to a registry:
-- The caller registers itself at the registry.
-- The owner of the registry will then
-- use the caller depending on its purpose.
--
-- * 'Con': Connection to a Stomp broker;
--
-- * 'JobName': The name of the job provided by the caller;
--
-- * 'JobType': The type of the job provided by the caller;
--
-- * 'QName': The registry's registration queue;
--
-- * 'QName': The queue to register;
-- this is the queue the register will actually
-- use (for forwarding requests or whatever
-- it does in this specific case).
-- The registry, internally,
-- uses 'JobName' together with this queue
-- as a /key/ to identify the provider.
--
-- * Int: Timeout in microseconds;
--
-- * Int: Preferred heartbeat in milliseconds
-- (0 for no heartbeats).
--
-- The function returns a tuple of 'StatusCode'
-- and the heartbeat proposed by the registry
-- (which may differ from the preferred heartbeat of the caller).
-- Whenever the 'StatusCode' is not 'OK',
-- the heartbeat is 0.
-- If the 'JobName' is null, the 'StatusCode' will be 'BadRequest'.
-- If the timeout expires, register throws 'TimeoutX'.
------------------------------------------------------------------------
register :: Con -> JobName -> JobType ->
QName -> QName ->
Int -> Int -> IO (StatusCode, Int)
register c j t o i to me | null j = return (BadRequest,0)
| otherwise =
let i' = o ++ "/" ++ j ++ "/" ++ i
hs = [("__type__", "register"),
("__job-type__", show t),
("__job__", j),
("__queue__", i),
("__hb__", show me),
("__channel__", i')]
in withWriter c "RegistryW" o [] [] nobody $ \w ->
withReader c "RegistryR" i' [] [] ignorebody $ \r -> do
writeQ w nullType hs ()
mbF <- timeout to $ readQ r
case mbF of
Nothing -> throwIO $ TimeoutX
"No response from registry"
Just m -> do eiS <- getSC m
case eiS of
Left s -> throwIO $ BadStatusCodeX s
Right OK -> do h <- getHB m
return (OK, h)
Right sc -> return (sc, 0)
------------------------------------------------------------------------
-- | Disconnect from a registry:
-- The caller disconnects from a registry
-- to which it has registered before.
-- For the case that the registry is not receiving heartbeats
-- from the caller,
-- it is essential to unregister, when
-- the service is no longer provided.
-- Otherwise, the registry has no way to know
-- that it should not send requests to this provider anymore.
--
-- * 'Con': Connection to a Stomp broker;
--
-- * 'JobName': The 'JobName' to unregister;
--
-- * 'QName': The registry's registration queue ;
--
-- * 'QName': The queue to unregister;
--
-- * Int: The timeout in microseconds.
--
-- The function returns a 'StatusCode'.
-- If 'JobName' is null, the 'StatusCode' will be 'BadRequest'.
-- If the timeout expires, the function will throw 'TimeoutX'.
------------------------------------------------------------------------
unRegister :: Con -> JobName ->
QName -> QName ->
Int -> IO StatusCode
unRegister c j o i tmo | null j = return BadRequest
| otherwise =
let i' = o ++ "/" ++ j ++ i
hs = [("__type__", "unreg"),
("__job__", j),
("__queue__", i),
("__channel__", i')]
in withWriter c "RegistryW" o [] [] nobody $ \w ->
withReader c "RegistryR" i' [] [] ignorebody $ \r -> do
writeQ w nullType hs ()
mbF <- timeout tmo $ readQ r
case mbF of
Nothing -> throwIO $ TimeoutX "No response from register"
Just m -> do eiS <- getSC m
case eiS of
Left s -> throwIO $ BadStatusCodeX s
Right sc -> return sc
------------------------------------------------------------------------
-- | Send heartbeats:
--
-- * 'MVar' 'HB': An MVar of type 'HB', this MVar will be used
-- to keep track of when the heartbeat has actually to
-- be sent.
--
-- * 'Writer' (): The writer through which to send the heartbeat;
-- The queue name of the writer is the registration queue
-- of the registry; note that its type is ():
-- heartbeats are empty messages.
--
-- * 'JobName': The 'JobName' for which to send heartbeats;
--
-- * 'QName': The queue for which to send heartbeats.
------------------------------------------------------------------------
heartbeat :: MVar HB -> Writer () -> JobName -> QName -> IO ()
heartbeat m w j q | null q = return ()
| otherwise =
let hs = [("__type__", "hb"),
("__job__", j),
("__queue__", q)]
in do now <- getCurrentTime
modifyMVar_ m (go now hs)
where go now hs hb@(HB me nxt)
| me > 0 && nxt < now = do writeQ w nullType hs ()
return hb{hbMeNext = timeAdd now me}
| otherwise = return hb
------------------------------------------------------------------------
-- | A provider is an opaque data type;
-- most of its attributes are used only internally by the registry.
-- Interesting for user applications, however, is the queue
-- that identifies the provider.
------------------------------------------------------------------------
data Provider = Provider {
-- | Queue through which the job is provided
prvQ :: QName,
prvHb :: Int,
prvNxt :: UTCTime
}
deriving Show
------------------------------------------------------------------------
-- Two providers are identical if they have the same queue name
------------------------------------------------------------------------
instance Eq Provider where
x == y = prvQ x == prvQ y
------------------------------------------------------------------------
-- Add provider to seq or, if already in,
-- update according to the values of the new node.
------------------------------------------------------------------------
updOrAddProv :: Bool -> (Provider -> Provider) -> Provider ->
Seq Provider -> Seq Provider
updOrAddProv add upd p s =
case S.viewl s of
S.EmptyL -> if add then S.singleton p else S.empty
x :< ss -> if prvQ x == prvQ p
then upd x <| ss
else x <| updOrAddProv add upd p ss
------------------------------------------------------------------------
-- Remove one provider from the seq
------------------------------------------------------------------------
remProv :: QName -> Seq Provider -> Seq Provider
remProv q s =
case S.viewl s of
S.EmptyL -> S.empty
x :< ss -> if prvQ x == q then ss
else x <| remProv q ss
------------------------------------------------------------------------
-- Get head of seq and add to end of sequence;
-- remove all "dead" nodes on the way
------------------------------------------------------------------------
getHeads :: UTCTime -> Seq Provider -> ([Provider], Seq Provider)
getHeads now s =
case S.viewl s of
S.EmptyL -> ([], S.empty)
x :< ss -> if prvHb x > 0 &&
prvNxt x < now then getHeads now ss
else ([x], ss |> x)
------------------------------------------------------------------------
-- Job: 'JobType' plus 'Sequence' of 'Provider's
------------------------------------------------------------------------
data JobNode = JobNode {
jobType :: JobType,
jobProvs :: Seq Provider
}
------------------------------------------------------------------------
-- The inner heart of the registry:
-- a 'Map' of 'JobName', 'JobNode'
------------------------------------------------------------------------
data Reg = Reg {
regName :: String,
regWork :: Map JobName JobNode
}
------------------------------------------------------------------------
-- | Registry: An opaque data type
------------------------------------------------------------------------
data Registry = Registry {
regM :: MVar Reg
}
------------------------------------------------------------------------
-- Use registry (with return value)
------------------------------------------------------------------------
useRegistry :: Registry -> (Reg -> IO (Reg, r)) -> IO r
useRegistry r = modifyMVar (regM r)
------------------------------------------------------------------------
-- Use registry (without return value)
------------------------------------------------------------------------
useRegistry_ :: Registry -> (Reg -> IO Reg) -> IO ()
useRegistry_ r = modifyMVar_ (regM r)
------------------------------------------------------------------------
-- Add provider to job
------------------------------------------------------------------------
insertR :: Registry -> JobName -> JobType -> QName -> Int -> IO ()
insertR r jn w qn i =
useRegistry_ r $ \reg -> do now <- getCurrentTime
return reg{regWork = ins now $ regWork reg}
where ins now m =
let j = fromMaybe (JobNode w S.empty) $ M.lookup jn m
p = Provider qn i $ nextHB now True i
ps = updOrAddProv True (upd p) p $ jobProvs j
in M.insert jn j{jobProvs = ps} m
upd n _ = n
------------------------------------------------------------------------
-- Update heartbeat of provider
------------------------------------------------------------------------
updR :: Registry -> JobName -> QName -> IO ()
updR r jn qn =
useRegistry_ r $ \reg -> do now <- getCurrentTime
return reg{regWork = ins now $ regWork reg}
where ins now m =
case M.lookup jn m of
Nothing -> m
Just j -> let p = Provider qn 0 now
ps = updOrAddProv False (upd now) p
(jobProvs j)
in M.insert jn j{jobProvs = ps} m
upd now o = o{prvNxt = nextHB now True $ tolerance * prvHb o}
------------------------------------------------------------------------
-- Remove 'Provider' from the job
------------------------------------------------------------------------
removeR :: Registry -> JobName -> QName -> IO ()
removeR r jn qn =
useRegistry_ r $ \reg -> return reg{regWork = ins $ regWork reg}
where ins m =
case M.lookup jn m of
Nothing -> m
Just j ->
let ps = remProv qn $ jobProvs j
in if S.null ps then M.delete jn m
else M.insert jn j{jobProvs = ps} m
------------------------------------------------------------------------
-- | Map action to 'Provider's of job 'JobName';
-- mapping means different things for:
--
-- * Serice, Task: action is applied to the first
-- active provider of a list of providers
-- and this provider
-- is then sent to the back of the list,
-- hence, implementing a balancer.
--
-- * Topic: action is applied to all providers,
-- hence, implementing a publisher.
--
-- Parameters:
--
-- * 'Registry': The registry to use;
--
-- * 'JobName': The job to which to apply the action;
--
-- * ('Provider' -> IO ()): The action to apply.
--
-- The function returns False iff the requested job is not available
-- and True otherwise. (Note that a job without providers is removed;
-- when the function returns True, the job, thus,
-- was applied at least once.
------------------------------------------------------------------------
mapR :: Registry -> JobName -> (Provider -> IO ()) -> IO Bool
mapR r jn f =
useRegistry r $ \reg -> getCurrentTime >>= \now ->
ins now (regWork reg) >>= \(js,t) ->
return (reg{regWork = js},t)
where ins now m =
case M.lookup jn m of
Nothing -> return (m, False)
Just j ->
let (xs, ps) = if jobType j `elem` [Service, Task]
then getHeads now $ jobProvs j
else (toList $ jobProvs j,
jobProvs j)
in mapM_ f xs >>
return (M.insert jn j{jobProvs = ps} m, True)
------------------------------------------------------------------------
-- | Map function of type
--
-- > 'Provider' -> 'Provider'
--
-- to all 'Provider's of job 'JobName'
-- (independent of 'JobType')
------------------------------------------------------------------------
mapAllR :: Registry -> JobName -> (Provider -> Provider) -> IO ()
mapAllR r jn f =
useRegistry_ r $ \reg -> ins (regWork reg) >>= \m ->
return reg{regWork = m}
where ins m =
case M.lookup jn m of
Nothing -> return m
Just j -> return (M.insert jn j{jobProvs = go $ jobProvs j} m)
go s = case S.viewl s of
S.EmptyL -> S.empty
x :< ss -> f x <| go ss
------------------------------------------------------------------------
-- | Retrieves /n/ 'Provider's of a certain job;
-- getProvider works, for all 'JobType's
-- according to the work balancer logic, /i.e./:
-- it returns the first n providers of the list for this job
-- and moves them to the end of the list.
-- 'getProvider' is used, for instance, in the Desk pattern.
--
-- * 'Registry': The registry in use;
--
-- * 'JobName': The job for which the caller needs a provider;
--
-- * Int: The number /n/ of providers to retrieve;
-- if less than /n/ providers are available for this job,
-- all available providers will be returned,
-- but no error event is created.
------------------------------------------------------------------------
getProvider :: Registry -> JobName -> Int -> IO [Provider]
getProvider r jn n =
useRegistry r $ \reg -> do now <- getCurrentTime
let (x,m) = ins now $ regWork reg
return (reg{regWork = m}, x)
where ins now m = case M.lookup jn m of
Nothing -> ([], m)
Just j ->
let (x,ps) = go now (jobProvs j) n
in (x, M.insert jn j{jobProvs = ps} m)
go now ps i | i <= 0 = ([],ps)
| otherwise = let (!x ,ps1) = getHeads now ps
(!x',ps2) = go now ps1 (i-1)
in (nub (x++x'), ps2)
------------------------------------------------------------------------
-- | This function shows all jobs with all their providers
-- in a registry; the function is intended for debugging only.
------------------------------------------------------------------------
showRegistry :: Registry -> IO ()
showRegistry r =
useRegistry_ r $ \reg -> let l = map fst $ M.toList (regWork reg)
p = map (getProvs reg) l
lp = zip l p
in print lp >> return reg
where getProvs reg jn = case M.lookup jn $ regWork reg of
Nothing -> []
Just ps -> toList $ jobProvs ps
------------------------------------------------------------------------
-- | A registry is used through a function
-- that, internally, creates a registry
-- and defines its lifetime in terms of the scope of an action
-- passed in to the function:
--
-- * 'Con': Connection to a Stomp broker;
--
-- * String: Name of the registry used for error handling;
--
-- * 'QName': Name of the registration queue.
-- It is this queue to which 'register'
-- sends a registration request;
--
-- * (Int, Int): Minimal and maximal accepted heartbeat interval;
--
-- * 'OnError': Error handler;
--
-- * ('Registry' -> IO r): The action that defines
-- the registry's lifetime;
-- the result of this action, /r/,
-- is also the result of /withRegistry/.
------------------------------------------------------------------------
withRegistry :: Con -> String -> QName -> (Int, Int)
-> OnError -> (Registry -> IO r) -> IO r
withRegistry c n rq (mn, mx) onErr action =
-- always start the reader in the main thread -------------
-- for if started in the background thread -------------
-- the action may send a message -------------
-- without the reader having subscribed to its queue ------
withReader c (n ++ "Reader") rq [] [] ignorebody $ \r -> do
let nm = n ++ "Registry"
reg <- Registry <$> newMVar (Reg nm M.empty)
withThread (startReg reg r nm) (action reg)
where startReg reg r nm =
withWriter c (n ++ "Writer") "unknown" [] [] nobody $ \w ->
forever $ catches
(do m <- readQ r
t <- getMType m
case t of
"register" -> handleRegister reg m w (mn,mx)
"unreg" -> handleUnRegister reg m w
"hb" -> handleHeartbeat reg m
x -> throwIO $ HeaderX "__type__" $
"Unknown type: " ++ x)
(ignoreHandler nm onErr)
------------------------------------------------------------------------
-- Handle registration request
------------------------------------------------------------------------
handleRegister :: Registry -> Message m -> Writer () -> (Int, Int) -> IO ()
handleRegister r m w (mn,mx) = do
(j,q) <- getJobQueue m
ch <- getChannel m
t <- getJobType m
hb <- getHB m
let h | hb < mn || hb > mx = if (mn - hb) < (hb - mx) then mn else mx
| otherwise = hb
insertR r j t q h
let hs = [("__sc__", show OK),
("__hb__", show h)]
writeAdHoc w ch nullType hs ()
------------------------------------------------------------------------
-- Handle unRegister request
------------------------------------------------------------------------
handleUnRegister :: Registry -> Message m -> Writer () -> IO ()
handleUnRegister r m w = do
(j,q) <- getJobQueue m
ch <- getChannel m
removeR r j q
let hs=[("__sc__", show OK)]
writeAdHoc w ch nullType hs ()
------------------------------------------------------------------------
-- Handle heartbeat
------------------------------------------------------------------------
handleHeartbeat :: Registry -> Message m -> IO ()
handleHeartbeat r m = do
(j,q) <- getJobQueue m
updR r j q
-- print $ msgHdrs m -- test
------------------------------------------------------------------------
-- | Get JobQueue
-- (and throw an exception if at least
-- one of the headers does not exist)
------------------------------------------------------------------------
getJobQueue :: Message m -> IO (String, String)
getJobQueue m = getJobName m >>= \j -> getQueue m >>= \q -> return (j,q)
------------------------------------------------------------------------
-- | Get Message Type from headers
-- (and throw an exception if the header does not exist)
------------------------------------------------------------------------
getMType :: Message m -> IO String
getMType = getHeader "__type__" "No message type in headers"
------------------------------------------------------------------------
-- | Get Job name from headers
-- (and throw an exception if the header does not exist)
------------------------------------------------------------------------
getJobName :: Message m -> IO String
getJobName = getHeader "__job__" "No job name in headers"
------------------------------------------------------------------------
-- | Get Reply queue (channel) from headers
-- (and throw an exception if the header does not exist)
------------------------------------------------------------------------
getChannel :: Message m -> IO String
getChannel = getHeader "__channel__" "No response q in headers"
------------------------------------------------------------------------
-- | Get Queue name from headers
-- (and throw an exception if the header does not exist)
------------------------------------------------------------------------
getQueue :: Message m -> IO String
getQueue = getHeader "__queue__" "No queue q in headers"
------------------------------------------------------------------------
-- | Get Job type from headers
-- (and throw an exception if the header does not exist
-- or contains an invalid value)
------------------------------------------------------------------------
getJobType :: Message m -> IO JobType
getJobType m =
getHeader "__job-type__" "No job type in headers" m >>= \x ->
case readJobType x of
Nothing -> throwIO $ HeaderX "__job-type__" $
"unknown type: " ++ x
Just t -> return t
------------------------------------------------------------------------
-- | Get Heartbeat specification from headers
-- (and throw an exception if the header does not exist
-- or if its value is not numeric)
------------------------------------------------------------------------
getHB :: Message m -> IO Int
getHB m =
case lookup "__hb__" $ msgHdrs m of
Nothing -> return 0
Just v -> if all isDigit v
then return $ read v
else throwIO $ HeaderX "__hb__" $
"heartbeat not numeric: " ++ show v
------------------------------------------------------------------------
-- | Get Status code from headers
-- (and throw an exception if the header does not exist)
------------------------------------------------------------------------
getSC :: Message m -> IO (Either String StatusCode)
getSC m = readStatusCode <$> getHeader "__sc__"
"No status code in message" m
------------------------------------------------------------------------
-- | Get Generic function to retrieve a header value
-- (and throw an exception if the header does not exist):
--
-- * String: Key of the wanted header
--
-- * String: Error message in case there is no such header
--
-- * 'Message' m: The message whose headers we want to search
------------------------------------------------------------------------
getHeader :: String -> String -> Message m -> IO String
getHeader h e m = case lookup h $ msgHdrs m of
Nothing -> throwIO $ HeaderX h e
Just v -> return v