packages feed

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 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]     ]