hedis 0.9.1 → 0.9.2
raw patch · 6 files changed
+592/−10 lines, 6 filesdep +asyncdep +stmdep +unordered-containersdep ~basePVP ok
version bump matches the API change (PVP)
Dependencies added: async, stm, unordered-containers
Dependency ranges changed: base
API changes (from Hackage documentation)
+ Database.Redis: addChannels :: MonadIO m => PubSubController -> [(RedisChannel, MessageCallback)] -> [(RedisPChannel, PMessageCallback)] -> m UnregisterCallbacksAction
+ Database.Redis: addChannelsAndWait :: MonadIO m => PubSubController -> [(RedisChannel, MessageCallback)] -> [(RedisPChannel, PMessageCallback)] -> m UnregisterCallbacksAction
+ Database.Redis: currentChannels :: MonadIO m => PubSubController -> m [RedisChannel]
+ Database.Redis: currentPChannels :: MonadIO m => PubSubController -> m [RedisPChannel]
+ Database.Redis: data PubSubController
+ Database.Redis: newPubSubController :: MonadIO m => [(RedisChannel, MessageCallback)] -> [(RedisPChannel, PMessageCallback)] -> m PubSubController
+ Database.Redis: pubSubForever :: Connection -> PubSubController -> IO () -> IO ()
+ Database.Redis: removeChannels :: MonadIO m => PubSubController -> [RedisChannel] -> [RedisPChannel] -> m ()
+ Database.Redis: removeChannelsAndWait :: MonadIO m => PubSubController -> [RedisChannel] -> [RedisPChannel] -> m ()
+ Database.Redis: type MessageCallback = ByteString -> IO ()
+ Database.Redis: type PMessageCallback = RedisChannel -> ByteString -> IO ()
+ Database.Redis: type RedisChannel = ByteString
+ Database.Redis: type RedisPChannel = ByteString
+ Database.Redis: type UnregisterCallbacksAction = IO ()
Files
- hedis.cabal +18/−2
- src/Database/Redis/Core.hs +1/−2
- src/Database/Redis/ProtocolPipelining.hs +7/−1
- src/Database/Redis/PubSub.hs +383/−4
- test/PubSubTest.hs +181/−0
- test/Test.hs +2/−1
hedis.cabal view
@@ -1,5 +1,5 @@ name: hedis-version: 0.9.1+version: 0.9.2 synopsis: Client library for the Redis datastore: supports full command set, pipelining.@@ -58,16 +58,23 @@ library hs-source-dirs: src ghc-options: -Wall -fwarn-tabs+ if flag(dev)+ ghc-options: -Werror+ if flag(dev)+ ghc-prof-options: -auto-all exposed-modules: Database.Redis build-depends: scanner >= 0.2,+ async >= 2.1, base >= 4.6 && < 5, bytestring >= 0.9, bytestring-lexing >= 0.5,+ unordered-containers, text, deepseq, mtl >= 2, network >= 2, resource-pool >= 0.2,+ stm, time, vector >= 0.9 @@ -89,15 +96,24 @@ hedis, time >= 1.2 ghc-options: -O2 -Wall -rtsopts+ if flag(dev)+ ghc-options: -Werror+ if flag(dev)+ ghc-prof-options: -auto-all test-suite hedis-test type: exitcode-stdio-1.0- main-is: test/Test.hs+ hs-source-dirs: test+ main-is: Test.hs+ other-modules: PubSubTest build-depends: base == 4.*, bytestring >= 0.10, hedis, HUnit,+ async,+ stm,+ text, mtl == 2.*, slave-thread, test-framework,
src/Database/Redis/Core.hs view
@@ -2,7 +2,7 @@ MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, CPP #-} module Database.Redis.Core (- Connection, connect,+ Connection(..), connect, ConnectInfo(..), defaultConnectInfo, Redis(), runRedis, unRedis, reRedis, RedisCtx(..), MonadRedis(..),@@ -39,7 +39,6 @@ deriving (Monad, MonadIO, Functor, Applicative) data RedisEnv = Env { envConn :: PP.Connection, envLastReply :: IORef Reply }- -- |This class captures the following behaviour: In a context @m@, a command -- will return it's result wrapped in a \"container\" of type @f@.
src/Database/Redis/ProtocolPipelining.hs view
@@ -15,7 +15,7 @@ -- module Database.Redis.ProtocolPipelining ( Connection,- connect, disconnect, request, send, recv,+ connect, disconnect, request, send, recv, flush, ConnectionLostException(..), HostName, PortID(..) ) where@@ -91,6 +91,12 @@ (r:rs) <- readIORef connReplies writeIORef connReplies rs return r++-- | Flush the socket. Normally, the socket is flushed in 'recv' (actually 'conGetReplies'), but+-- for the multithreaded pub/sub code, the sending thread needs to explicitly flush the subscription+-- change requests.+flush :: Connection -> IO ()+flush Conn{..} = hFlush connHandle -- |Send a request and receive the corresponding reply request :: Connection -> S.ByteString -> IO Reply
src/Database/Redis/PubSub.hs view
@@ -1,23 +1,42 @@ {-# LANGUAGE CPP, OverloadedStrings, RecordWildCards, EmptyDataDecls,- FlexibleInstances, FlexibleContexts #-}+ FlexibleInstances, FlexibleContexts, GeneralizedNewtypeDeriving #-} module Database.Redis.PubSub ( publish,++ -- ** Subscribing to channels+ -- $pubsubexpl++ -- *** Single-thread Pub/Sub pubSub, Message(..), PubSub(),- subscribe, unsubscribe, psubscribe, punsubscribe+ subscribe, unsubscribe, psubscribe, punsubscribe,+ -- *** Continuous Pub/Sub message controller+ pubSubForever,+ RedisChannel, RedisPChannel, MessageCallback, PMessageCallback,+ PubSubController, newPubSubController, currentChannels, currentPChannels,+ addChannels, addChannelsAndWait, removeChannels, removeChannelsAndWait,+ UnregisterCallbacksAction ) where #if __GLASGOW_HASKELL__ < 710 import Control.Applicative import Data.Monoid #endif+import Control.Concurrent.Async (withAsync, waitEitherCatch, waitEitherCatchSTM)+import Control.Concurrent.STM+import Control.Exception (throwIO) import Control.Monad import Control.Monad.State import Data.ByteString.Char8 (ByteString)+import Data.List (foldl')+import Data.Maybe (isJust)+import Data.Pool+import qualified Data.HashMap.Strict as HM import qualified Database.Redis.Core as Core-import Database.Redis.Protocol (Reply(..))+import qualified Database.Redis.ProtocolPipelining as PP+import Database.Redis.Protocol (Reply(..), renderRequest) import Database.Redis.Types -- |While in PubSub mode, we keep track of the number of current subscriptions@@ -83,6 +102,18 @@ lift $ Core.send (redisCmd cmd : changes cmd) modifyPending (updatePending cmd) +cmdCount :: Cmd a b -> Int+cmdCount DoNothing = 0+cmdCount (Cmd c) = length c++totalPendingChanges :: PubSub -> Int+totalPendingChanges (PubSub{..}) =+ cmdCount subs + cmdCount unsubs + cmdCount psubs + cmdCount punsubs++rawSendCmd :: (Command (Cmd a b)) => PP.Connection -> Cmd a b -> IO ()+rawSendCmd _ DoNothing = return ()+rawSendCmd conn cmd = PP.send conn $ renderRequest $ redisCmd cmd : changes cmd+ plusChangeCnt :: Cmd a b -> Int -> Int plusChangeCnt DoNothing = id plusChangeCnt (Cmd cs) = (+ length cs)@@ -212,6 +243,343 @@ PubSubState{..} <- get unless (subCnt == 0 && pending == 0) recv +-- | A Redis channel name+type RedisChannel = ByteString++-- | A Redis pattern channel name+type RedisPChannel = ByteString++-- | A handler for a message from a subscribed channel.+-- The callback is passed the message content.+--+-- Messages are processed synchronously in the receiving thread, so if the callback+-- takes a long time it will block other callbacks and other messages from being+-- received. If you need to move long-running work to a different thread, we suggest+-- you use 'TBQueue' with a reasonable bound, so that if messages are arriving faster+-- than you can process them, you do eventually block.+--+-- If the callback throws an exception, the exception will be thrown from 'pubSubForever'+-- which will cause the entire Redis connection for all subscriptions to be closed.+-- As long as you call 'pubSubForever' in a loop you will reconnect to your subscribed+-- channels, but you should probably add an exception handler to each callback to+-- prevent this.+type MessageCallback = ByteString -> IO ()++-- | A handler for a message from a psubscribed channel.+-- The callback is passed the channel the message was sent on plus the message content.+--+-- Similar to 'MessageCallback', callbacks are executed synchronously and any exceptions+-- are rethrown from 'pubSubForever'.+type PMessageCallback = RedisChannel -> ByteString -> IO ()++-- | An action that when executed will unregister the callbacks. It is returned from 'addChannels'+-- or 'addChannelsAndWait' and typically you would use it in 'bracket' to guarantee that you+-- unsubscribe from channels. For example, if you are using websockets to distribute messages to+-- clients, you could use something such as:+--+-- > websocketConn <- Network.WebSockets.acceptRequest pending+-- > let mycallback msg = Network.WebSockets.sendTextData websocketConn msg+-- > bracket (addChannelsAndWait ctrl [("hello", mycallback)] []) id $ const $ do+-- > {- loop here calling Network.WebSockets.receiveData -}+type UnregisterCallbacksAction = IO ()++newtype UnregisterHandle = UnregisterHandle Integer+ deriving (Eq, Show, Num)++-- | A controller that stores a set of channels, pattern channels, and callbacks.+-- It allows you to manage Pub/Sub subscriptions and pattern subscriptions and alter them at+-- any time throughout the life of your program.+-- You should typically create the controller at the start of your program and then store it+-- through the life of your program, using 'addChannels' and 'removeChannels' to update the+-- current subscriptions.+data PubSubController = PubSubController+ { callbacks :: TVar (HM.HashMap RedisChannel [(UnregisterHandle, MessageCallback)])+ , pcallbacks :: TVar (HM.HashMap RedisPChannel [(UnregisterHandle, PMessageCallback)])+ , sendChanges :: TBQueue PubSub+ , pendingCnt :: TVar Int+ , lastUsedCallbackId :: TVar UnregisterHandle+ }++-- | Create a new 'PubSubController'. Note that this does not subscribe to any channels, it just+-- creates the controller. The subscriptions will happen once 'pubSubForever' is called.+newPubSubController :: MonadIO m => [(RedisChannel, MessageCallback)] -- ^ the initial subscriptions+ -> [(RedisPChannel, PMessageCallback)] -- ^ the initial pattern subscriptions+ -> m PubSubController+newPubSubController x y = liftIO $ do+ cbs <- newTVarIO (HM.map (\z -> [(0,z)]) $ HM.fromList x)+ pcbs <- newTVarIO (HM.map (\z -> [(0,z)]) $ HM.fromList y)+ c <- newTBQueueIO 10+ pending <- newTVarIO 0+ lastId <- newTVarIO 0+ return $ PubSubController cbs pcbs c pending lastId++-- | Get the list of current channels in the 'PubSubController'. WARNING! This might not+-- exactly reflect the subscribed channels in the Redis server, because there is a delay+-- between adding or removing a channel in the 'PubSubController' and when Redis receives+-- and processes the subscription change request.+currentChannels :: MonadIO m => PubSubController -> m [RedisChannel]+currentChannels ctrl = HM.keys <$> (liftIO $ atomically $ readTVar $ callbacks ctrl)++-- | Get the list of current pattern channels in the 'PubSubController'. WARNING! This might not+-- exactly reflect the subscribed channels in the Redis server, because there is a delay+-- between adding or removing a channel in the 'PubSubController' and when Redis receives+-- and processes the subscription change request.+currentPChannels :: MonadIO m => PubSubController -> m [RedisPChannel]+currentPChannels ctrl = HM.keys <$> (liftIO $ atomically $ readTVar $ pcallbacks ctrl)++-- | Add channels into the 'PubSubController', and if there is an active 'pubSubForever', send the subscribe+-- and psubscribe commands to Redis. The 'addChannels' function is thread-safe. This function+-- does not wait for Redis to acknowledge that the channels have actually been subscribed; use+-- 'addChannelsAndWait' for that.+--+-- You can subscribe to the same channel or pattern channel multiple times; the 'PubSubController' keeps+-- a list of callbacks and executes each callback in response to a message.+--+-- The return value is an action 'UnregisterCallbacksAction' which will unregister the callbacks,+-- which should typically used with 'bracket'.+addChannels :: MonadIO m => PubSubController+ -> [(RedisChannel, MessageCallback)] -- ^ the channels to subscribe to+ -> [(RedisPChannel, PMessageCallback)] -- ^ the channels to pattern subscribe to+ -> m UnregisterCallbacksAction+addChannels _ [] [] = return $ return ()+addChannels ctrl newChans newPChans = liftIO $ do+ ident <- atomically $ do+ modifyTVar (lastUsedCallbackId ctrl) (+1)+ ident <- readTVar $ lastUsedCallbackId ctrl+ cm <- readTVar $ callbacks ctrl+ pm <- readTVar $ pcallbacks ctrl+ let newChans' = [ n | (n,_) <- newChans, not $ HM.member n cm]+ newPChans' = [ n | (n, _) <- newPChans, not $ HM.member n pm]+ ps = subscribe newChans' `mappend` psubscribe newPChans'+ writeTBQueue (sendChanges ctrl) ps+ writeTVar (callbacks ctrl) (HM.unionWith (++) cm (fmap (\z -> [(ident,z)]) $ HM.fromList newChans))+ writeTVar (pcallbacks ctrl) (HM.unionWith (++) pm (fmap (\z -> [(ident,z)]) $ HM.fromList newPChans))+ modifyTVar (pendingCnt ctrl) (+ totalPendingChanges ps)+ return ident+ return $ unsubChannels ctrl (map fst newChans) (map fst newPChans) ident++-- | Call 'addChannels' and then wait for Redis to acknowledge that the channels are actually subscribed.+--+-- Note that this function waits for all pending subscription change requests, so if you for example call+-- 'addChannelsAndWait' from multiple threads simultaneously, they all will wait for all pending+-- subscription changes to be acknowledged by Redis (this is due to the fact that we just track the total+-- number of pending change requests sent to Redis and just wait until that count reaches zero).+--+-- This also correctly waits if the network connection dies during the subscription change. Say that the+-- network connection dies right after we send a subscription change to Redis. 'pubSubForever' will throw+-- 'ConnectionLost' and 'addChannelsAndWait' will continue to wait. Once you recall 'pubSubForever'+-- with the same 'PubSubController', 'pubSubForever' will open a new connection, send subscription commands+-- for all channels in the 'PubSubController' (which include the ones we are waiting for),+-- and wait for the responses from Redis. Only once we receive the response from Redis that it has subscribed+-- to all channels in 'PubSubController' will 'addChannelsAndWait' unblock and return.+addChannelsAndWait :: MonadIO m => PubSubController+ -> [(RedisChannel, MessageCallback)] -- ^ the channels to subscribe to+ -> [(RedisPChannel, PMessageCallback)] -- ^ the channels to psubscribe to+ -> m UnregisterCallbacksAction+addChannelsAndWait _ [] [] = return $ return ()+addChannelsAndWait ctrl newChans newPChans = do+ unreg <- addChannels ctrl newChans newPChans+ liftIO $ atomically $ do+ r <- readTVar (pendingCnt ctrl)+ when (r > 0) retry+ return unreg++-- | Remove channels from the 'PubSubController', and if there is an active 'pubSubForever', send the+-- unsubscribe commands to Redis. Note that as soon as this function returns, no more callbacks will be+-- executed even if more messages arrive during the period when we request to unsubscribe from the channel+-- and Redis actually processes the unsubscribe request. This function is thread-safe.+--+-- If you remove all channels, the connection in 'pubSubForever' to redis will stay open and waiting for+-- any new channels from a call to 'addChannels'. If you really want to close the connection,+-- use 'Control.Concurrent.killThread' or 'Control.Concurrent.Async.cancel' to kill the thread running+-- 'pubSubForever'.+removeChannels :: MonadIO m => PubSubController+ -> [RedisChannel]+ -> [RedisPChannel]+ -> m ()+removeChannels _ [] [] = return ()+removeChannels ctrl remChans remPChans = liftIO $ atomically $ do+ cm <- readTVar $ callbacks ctrl+ pm <- readTVar $ pcallbacks ctrl+ let remChans' = filter (\n -> HM.member n cm) remChans+ remPChans' = filter (\n -> HM.member n pm) remPChans+ ps = (if null remChans' then mempty else unsubscribe remChans')+ `mappend` (if null remPChans' then mempty else punsubscribe remPChans')+ writeTBQueue (sendChanges ctrl) ps+ writeTVar (callbacks ctrl) (foldl' (flip HM.delete) cm remChans')+ writeTVar (pcallbacks ctrl) (foldl' (flip HM.delete) pm remPChans')+ modifyTVar (pendingCnt ctrl) (+ totalPendingChanges ps)++-- | Internal function to unsubscribe only from those channels matching the given handle.+unsubChannels :: PubSubController -> [RedisChannel] -> [RedisPChannel] -> UnregisterHandle -> IO ()+unsubChannels ctrl chans pchans h = liftIO $ atomically $ do+ cm <- readTVar $ callbacks ctrl+ pm <- readTVar $ pcallbacks ctrl++ -- only worry about channels that exist+ let remChans = filter (\n -> HM.member n cm) chans+ remPChans = filter (\n -> HM.member n pm) pchans++ -- helper functions to filter out handlers that match+ let filterHandle :: Maybe [(UnregisterHandle,a)] -> Maybe [(UnregisterHandle,a)]+ filterHandle Nothing = Nothing+ filterHandle (Just lst) = case filter (\x -> fst x /= h) lst of+ [] -> Nothing+ xs -> Just xs+ let removeHandles :: HM.HashMap ByteString [(UnregisterHandle,a)]+ -> ByteString+ -> HM.HashMap ByteString [(UnregisterHandle,a)]+ removeHandles m k = case filterHandle (HM.lookup k m) of -- recent versions of unordered-containers have alter+ Nothing -> HM.delete k m+ Just v -> HM.insert k v m++ -- maps after taking out channels matching the handle+ let cm' = foldl' removeHandles cm remChans+ pm' = foldl' removeHandles pm remPChans++ -- the channels to unsubscribe are those that no longer exist in cm' and pm'+ let remChans' = filter (\n -> not $ HM.member n cm') remChans+ remPChans' = filter (\n -> not $ HM.member n pm') remPChans+ ps = (if null remChans' then mempty else unsubscribe remChans')+ `mappend` (if null remPChans' then mempty else punsubscribe remPChans')++ -- do the unsubscribe+ writeTBQueue (sendChanges ctrl) ps+ writeTVar (callbacks ctrl) cm'+ writeTVar (pcallbacks ctrl) pm'+ modifyTVar (pendingCnt ctrl) (+ totalPendingChanges ps)+ return ()++-- | Call 'removeChannels' and then wait for all pending subscription change requests to be acknowledged+-- by Redis. This uses the same waiting logic as 'addChannelsAndWait'. Since 'removeChannels' immediately+-- notifies the 'PubSubController' to start discarding messages, you likely don't need this function and+-- can just use 'removeChannels'.+removeChannelsAndWait :: MonadIO m => PubSubController+ -> [RedisChannel]+ -> [RedisPChannel]+ -> m ()+removeChannelsAndWait _ [] [] = return ()+removeChannelsAndWait ctrl remChans remPChans = do+ removeChannels ctrl remChans remPChans+ liftIO $ atomically $ do+ r <- readTVar (pendingCnt ctrl)+ when (r > 0) retry++-- | Internal thread which listens for messages and executes callbacks.+-- This is the only thread which ever receives data from the underlying+-- connection.+listenThread :: PubSubController -> PP.Connection -> IO ()+listenThread ctrl rawConn = forever $ do+ msg <- PP.recv rawConn+ case decodeMsg msg of+ Msg (Message channel msgCt) -> do+ cm <- atomically $ readTVar (callbacks ctrl)+ case HM.lookup channel cm of+ Nothing -> return ()+ Just c -> mapM_ (\(_,x) -> x msgCt) c+ Msg (PMessage pattern channel msgCt) -> do+ pm <- atomically $ readTVar (pcallbacks ctrl)+ case HM.lookup pattern pm of+ Nothing -> return ()+ Just c -> mapM_ (\(_,x) -> x channel msgCt) c+ Subscribed -> atomically $+ modifyTVar (pendingCnt ctrl) (\x -> x - 1)+ Unsubscribed _ -> atomically $+ modifyTVar (pendingCnt ctrl) (\x -> x - 1)++-- | Internal thread which sends subscription change requests.+-- This is the only thread which ever sends data on the underlying+-- connection.+sendThread :: PubSubController -> PP.Connection -> IO ()+sendThread ctrl rawConn = forever $ do+ PubSub{..} <- atomically $ readTBQueue (sendChanges ctrl)+ rawSendCmd rawConn subs+ rawSendCmd rawConn unsubs+ rawSendCmd rawConn psubs+ rawSendCmd rawConn punsubs+ -- normally, the socket is flushed during 'recv', but+ -- 'recv' could currently be blocking on a message.+ PP.flush rawConn++-- | Open a connection to the Redis server, register to all channels in the 'PubSubController',+-- and process messages and subscription change requests forever. The only way this will ever+-- exit is if there is an exception from the network code or an unhandled exception+-- in a 'MessageCallback' or 'PMessageCallback'. For example, if the network connection to Redis+-- dies, 'pubSubForever' will throw a 'ConnectionLost'. When such an exception is+-- thrown, you can recall 'pubSubForever' with the same 'PubSubController' which will open a+-- new connection and resubscribe to all the channels which are tracked in the 'PubSubController'.+--+-- The general pattern is therefore during program startup create a 'PubSubController' and fork+-- a thread which calls 'pubSubForever' in a loop (using an exponential backoff algorithm+-- such as the <https://hackage.haskell.org/package/retry retry> package to not hammer the Redis+-- server if it does die). For example,+--+-- @+-- myhandler :: ByteString -> IO ()+-- myhandler msg = putStrLn $ unpack $ decodeUtf8 msg+--+-- onInitialComplete :: IO ()+-- onInitialComplete = putStrLn "Redis acknowledged that mychannel is now subscribed"+--+-- main :: IO ()+-- main = do+-- conn <- connect defaultConnectInfo+-- pubSubCtrl <- newPubSubController [("mychannel", myhandler)] []+-- forkIO $ forever $+-- pubSubForever conn pubSubCtrl onInitialComplete+-- \`catch\` (\\(e :: SomeException) -> do+-- putStrLn $ "Got error: " ++ show e+-- threadDelay $ 50*1000) -- TODO: use exponential backoff+--+-- {- elsewhere in your program, use pubSubCtrl to change subscriptions -}+-- @+--+-- At most one active 'pubSubForever' can be running against a single 'PubSubController' at any time. If+-- two active calls to 'pubSubForever' share a single 'PubSubController' there will be deadlocks. If+-- you do want to process messages using multiple connections to Redis, you can create more than one+-- 'PubSubController'. For example, create one PubSubController for each 'Control.Concurrent.getNumCapabilities'+-- and then create a Haskell thread bound to each capability each calling 'pubSubForever' in a loop.+-- This will create one network connection per controller/capability and allow you to+-- register separate channels and callbacks for each controller, spreading the load across the capabilities.+pubSubForever :: Core.Connection -- ^ The connection pool+ -> PubSubController -- ^ The controller which keeps track of all subscriptions and handlers+ -> IO () -- ^ This action is executed once Redis acknowledges that all the subscriptions in+ -- the controller are now subscribed. You can use this after an exception (such as+ -- 'ConnectionLost') to signal that all subscriptions are now reactivated.+ -> IO ()+pubSubForever (Core.Conn pool) ctrl onInitialLoad = withResource pool $ \rawConn -> do+ -- get initial subscriptions and write them into the queue.+ atomically $ do+ let loop = tryReadTBQueue (sendChanges ctrl) >>=+ \x -> if isJust x then loop else return ()+ loop+ cm <- readTVar $ callbacks ctrl+ pm <- readTVar $ pcallbacks ctrl+ let ps = subscribe (HM.keys cm) `mappend` psubscribe (HM.keys pm)+ writeTBQueue (sendChanges ctrl) ps+ writeTVar (pendingCnt ctrl) (totalPendingChanges ps)++ withAsync (listenThread ctrl rawConn) $ \listenT ->+ withAsync (sendThread ctrl rawConn) $ \sendT -> do++ -- wait for initial subscription count to go to zero or for threads to fail+ mret <- atomically $+ (Left <$> (waitEitherCatchSTM listenT sendT))+ `orElse`+ (Right <$> (readTVar (pendingCnt ctrl) >>=+ \x -> if x > 0 then retry else return ()))+ case mret of+ Right () -> onInitialLoad+ _ -> return () -- if there is an error, waitEitherCatch below will also see it++ -- wait for threads to end with error+ merr <- waitEitherCatch listenT sendT+ case merr of+ (Right (Left err)) -> throwIO err+ (Left (Left err)) -> throwIO err+ _ -> return () -- should never happen, since threads exit only with an error++ ------------------------------------------------------------------------------ -- Helpers --@@ -230,8 +598,19 @@ decodeMessage = Message <$> decode r1 <*> decode r2 decodePMessage = PMessage <$> decode r1 <*> decode r2 <*> decode (head rs) decodeCnt = fromInteger <$> decode r2- + decodeMsg r = errMsg r errMsg :: Reply -> a errMsg r = error $ "Hedis: expected pub/sub-message but got: " ++ show r+++-- $pubsubexpl+-- There are two Pub/Sub implementations. First, there is a single-threaded implementation 'pubSub'+-- which is simpler to use but has the restriction that subscription changes can only be made in+-- response to a message. Secondly, there is a more complicated Pub/Sub controller 'pubSubForever'+-- that uses concurrency to support changing subscriptions at any time but requires more setup.+-- You should only use one or the other. In addition, no types or utility functions (that are part+-- of the public API) are shared, so functions or types in one of the following sections cannot+-- be used for the other. In particular, be aware that they use different utility functions to subscribe+-- and unsubscribe to channels.
+ test/PubSubTest.hs view
@@ -0,0 +1,181 @@+{-# LANGUAGE CPP, OverloadedStrings #-}+module PubSubTest (testPubSubThreaded) where++#if __GLASGOW_HASKELL__ < 710+import Control.Applicative+import Data.Monoid (mappend)+#endif+import Control.Concurrent+import Control.Monad+import Control.Concurrent.Async+import Control.Exception+import Data.Typeable+--import Control.Monad.Trans+--import Data.Time+import qualified Data.List+import Data.Text+import Data.ByteString+import Control.Concurrent.STM+--import Data.Time.Clock.POSIX+import qualified Test.Framework as Test+import qualified Test.Framework.Providers.HUnit as Test (testCase)+import qualified Test.HUnit as HUnit++import Database.Redis++testPubSubThreaded :: [Connection -> Test.Test]+testPubSubThreaded = [removeAllTest, callbackErrorTest, removeFromUnregister]++-- | A handler label to be able to distinguish the handlers from one another+-- to help make sure we unregister the correct handler.+type HandlerLabel = Text++data TestMsg = MsgFromChannel HandlerLabel ByteString+ | MsgFromPChannel HandlerLabel RedisChannel ByteString+ deriving (Show, Eq)++type MsgVar = TVar [TestMsg]++-- | A handler that just writes the message into the TVar+handler :: HandlerLabel -> MsgVar -> MessageCallback+handler label ref msg = atomically $+ modifyTVar ref $ \x -> x ++ [MsgFromChannel label msg]++-- | A pattern handler that just writes the message into the TVar+phandler :: HandlerLabel -> MsgVar -> PMessageCallback+phandler label ref chan msg = atomically $+ modifyTVar ref $ \x -> x ++ [MsgFromPChannel label chan msg]++-- | Wait for a given message to be received+waitForMessage :: MsgVar -> HandlerLabel -> ByteString -> IO ()+waitForMessage ref label msg = atomically $ do+ let expected = MsgFromChannel label msg+ lst <- readTVar ref+ unless (expected `Prelude.elem` lst) retry+ writeTVar ref $ Prelude.filter (/= expected) lst++-- | Wait for a given pattern message to be received+waitForPMessage :: MsgVar -> HandlerLabel -> RedisChannel -> ByteString -> IO ()+waitForPMessage ref label chan msg = atomically $ do+ let expected = MsgFromPChannel label chan msg+ lst <- readTVar ref+ unless (expected `Prelude.elem` lst) retry+ writeTVar ref $ Prelude.filter (/= expected) lst++expectRedisChannels :: Connection -> [RedisChannel] -> IO ()+expectRedisChannels conn expected = do+ actual <- runRedis conn $ sendRequest ["PUBSUB", "CHANNELS"]+ case actual of+ Left err -> HUnit.assertFailure $ "Error geting channels: " ++ show err+ Right s -> HUnit.assertEqual "redis channels" (Data.List.sort s) (Data.List.sort expected)++-- | Test basic messages, plus using removeChannels+removeAllTest :: Connection -> Test.Test+removeAllTest conn = Test.testCase "Multithreaded Pub/Sub - basic" $ do+ msgVar <- newTVarIO []+ initialComplete <- newTVarIO False+ ctrl <- newPubSubController [("foo1", handler "InitialFoo1" msgVar), ("foo2", handler "InitialFoo2" msgVar)]+ [("bar1:*", phandler "InitialBar1" msgVar), ("bar2:*", phandler "InitialBar2" msgVar)]+ withAsync (pubSubForever conn ctrl (atomically $ writeTVar initialComplete True)) $ \_ -> do++ -- wait for initial+ atomically $ readTVar initialComplete >>= \b -> if b then return () else retry+ expectRedisChannels conn ["foo1", "foo2"]++ runRedis conn $ publish "foo1" "Hello"+ waitForMessage msgVar "InitialFoo1" "Hello"++ runRedis conn $ publish "bar2:zzz" "World"+ waitForPMessage msgVar "InitialBar2" "bar2:zzz" "World"++ -- subscribe to foo1 and bar1 again+ addChannelsAndWait ctrl [("foo1", handler "NewFoo1" msgVar)] [("bar1:*", phandler "NewBar1" msgVar)]+ expectRedisChannels conn ["foo1", "foo2"]++ runRedis conn $ publish "foo1" "abcdef"+ waitForMessage msgVar "InitialFoo1" "abcdef"+ waitForMessage msgVar "NewFoo1" "abcdef"++ -- unsubscribe from foo1 and bar1+ removeChannelsAndWait ctrl ["foo1", "unusued"] ["bar1:*", "unused:*"]+ expectRedisChannels conn ["foo2"]++ -- foo2 and bar2 are still subscribed+ runRedis conn $ publish "foo2" "12345"+ waitForMessage msgVar "InitialFoo2" "12345"++ runRedis conn $ publish "bar2:aaa" "0987"+ waitForPMessage msgVar "InitialBar2" "bar2:aaa" "0987"++data TestError = TestError ByteString+ deriving (Eq, Show, Typeable)+instance Exception TestError++-- | Test an error thrown from a message handler+callbackErrorTest :: Connection -> Test.Test+callbackErrorTest conn = Test.testCase "Multithreaded Pub/Sub - error in handler" $ do+ initialComplete <- newTVarIO False+ ctrl <- newPubSubController [("foo", throwIO . TestError)] []++ thread <- async (pubSubForever conn ctrl (atomically $ writeTVar initialComplete True))+ atomically $ readTVar initialComplete >>= \b -> if b then return () else retry++ runRedis conn $ publish "foo" "Hello"++ ret <- waitCatch thread+ case ret of+ Left (SomeException e) | cast e == Just (TestError "Hello") -> return ()+ _ -> HUnit.assertFailure $ "Did not properly throw error from message thread " ++ show ret++-- | Test removing channels by using the return value of 'addHandlersAndWait'.+removeFromUnregister :: Connection -> Test.Test+removeFromUnregister conn = Test.testCase "Multithreaded Pub/Sub - unregister handlers" $ do+ msgVar <- newTVarIO []+ initialComplete <- newTVarIO False+ ctrl <- newPubSubController [] []+ withAsync (pubSubForever conn ctrl (atomically $ writeTVar initialComplete True)) $ \_ -> do+ atomically $ readTVar initialComplete >>= \b -> if b then return () else retry++ -- register to some channels+ void $ addChannelsAndWait ctrl+ [("abc", handler "InitialAbc" msgVar), ("xyz", handler "InitialXyz" msgVar)]+ [("def:*", phandler "InitialDef" msgVar), ("uvw", phandler "InitialUvw" msgVar)]+ expectRedisChannels conn ["abc", "xyz"]++ runRedis conn $ publish "abc" "Hello"+ waitForMessage msgVar "InitialAbc" "Hello"++ -- register to some more channels+ unreg <- addChannelsAndWait ctrl+ [("abc", handler "SecondAbc" msgVar), ("123", handler "Second123" msgVar)]+ [("def:*", phandler "SecondDef" msgVar), ("890:*", phandler "Second890" msgVar)]+ expectRedisChannels conn ["abc", "xyz", "123"]++ -- check messages on all channels+ runRedis conn $ publish "abc" "World"+ waitForMessage msgVar "InitialAbc" "World"+ waitForMessage msgVar "SecondAbc" "World"++ runRedis conn $ publish "123" "World2"+ waitForMessage msgVar "Second123" "World2"++ runRedis conn $ publish "def:bbbb" "World3"+ waitForPMessage msgVar "InitialDef" "def:bbbb" "World3"+ waitForPMessage msgVar "SecondDef" "def:bbbb" "World3"++ runRedis conn $ publish "890:tttt" "World4"+ waitForPMessage msgVar "Second890" "890:tttt" "World4"++ -- unregister+ unreg++ -- we have no way of waiting until unregister actually happened, so just delay and hope+ threadDelay $ 1000*1000 -- 1 second+ expectRedisChannels conn ["abc", "xyz"]++ -- now only initial should be around. In particular, abc should still be subscribed+ runRedis conn $ publish "abc" "World5"+ waitForMessage msgVar "InitialAbc" "World5"++ runRedis conn $ publish "def:cccc" "World6"+ waitForPMessage msgVar "InitialDef" "def:cccc" "World6"
test/Test.hs view
@@ -16,7 +16,7 @@ import qualified Test.HUnit as HUnit import Database.Redis-+import PubSubTest ------------------------------------------------------------------------------ -- Main and helpers@@ -57,6 +57,7 @@ [ testsMisc, testsKeys, testsStrings, [testHashes], testsLists, testsSets, [testHyperLogLog] , testsZSets, [testPubSub], [testTransaction], [testScripting] , testsConnection, testsServer, [testScans], [testZrangelex]+ , testPubSubThreaded -- should always be run last as connection gets closed after it , [testQuit] ]