diff --git a/Network/AMQP.hs b/Network/AMQP.hs --- a/Network/AMQP.hs +++ b/Network/AMQP.hs @@ -65,6 +65,7 @@ closeChannel, closeConnection, addConnectionClosedHandler, + addConnectionBlockedHandler, getServerProperties, -- * Channel diff --git a/Network/AMQP/Generated.hs b/Network/AMQP/Generated.hs --- a/Network/AMQP/Generated.hs +++ b/Network/AMQP/Generated.hs @@ -110,6 +110,8 @@ put (Connection_open_ok a) = putWord16be 10 >> putWord16be 41 >> put a put (Connection_close a b c d) = putWord16be 10 >> putWord16be 50 >> put a >> put b >> put c >> put d put Connection_close_ok = putWord16be 10 >> putWord16be 51 + put (Connection_blocked a) = putWord16be 10 >> putWord16be 60 >> put a + put Connection_unblocked = putWord16be 10 >> putWord16be 61 put (Channel_open a) = putWord16be 20 >> putWord16be 10 >> put a put (Channel_open_ok a) = putWord16be 20 >> putWord16be 11 >> put a put (Channel_flow a) = putWord16be 20 >> putWord16be 20 >> put a @@ -174,6 +176,8 @@ (10,41) -> get >>= \a -> return (Connection_open_ok a) (10,50) -> get >>= \a -> get >>= \b -> get >>= \c -> get >>= \d -> return (Connection_close a b c d) (10,51) -> return Connection_close_ok + (10,60) -> get >>= \a -> return (Connection_blocked a) + (10,61) -> return Connection_unblocked (20,10) -> get >>= \a -> return (Channel_open a) (20,11) -> get >>= \a -> return (Channel_open_ok a) (20,20) -> get >>= \a -> return (Channel_flow a) @@ -273,6 +277,12 @@ Connection_close_ok | + Connection_blocked + ShortString -- reason + | + Connection_unblocked + + | Channel_open ShortString -- reserved-1 | @@ -502,4 +512,4 @@ Confirm_select_ok - deriving Show + deriving Show diff --git a/Network/AMQP/Internal.hs b/Network/AMQP/Internal.hs --- a/Network/AMQP/Internal.hs +++ b/Network/AMQP/Internal.hs @@ -149,6 +149,7 @@ connClosedLock :: MVar (), -- used by closeConnection to block until connection-close handshake is complete connWriteLock :: MVar (), -- to ensure atomic writes to the socket connClosedHandlers :: MVar [IO ()], + connBlockedHandlers :: MVar [(Text -> IO (), IO ())], connLastReceived :: MVar Int64, -- the timestamp from a monotonic clock when the last frame was received connLastSent :: MVar Int64, -- the timestamp from a monotonic clock when the last frame was written connServerProperties :: FieldTable -- the server properties sent in Connection_start @@ -206,6 +207,8 @@ writeFrame (connHandle conn) $ Frame 0 $ MethodPayload Connection_close_ok myThreadId >>= killConnection conn Abnormal (CE.toException . ConnectionClosedException Abnormal . T.unpack $ errorMsg) forwardToChannel 0 HeartbeatPayload = return () + forwardToChannel 0 (MethodPayload (Connection_blocked reason)) = handleBlocked reason + forwardToChannel 0 (MethodPayload Connection_unblocked) = handleUnblocked forwardToChannel 0 payload = hPutStrLn stderr $ "Got unexpected msg on channel zero: " ++ show payload forwardToChannel chanID payload = do --got asynchronous msg => forward to registered channel @@ -214,6 +217,16 @@ Just c -> writeChan (inQueue $ fst c) payload Nothing -> hPutStrLn stderr $ "ERROR: channel not open " ++ show chanID + handleBlocked (ShortString reason) = do + withMVar (connBlockedHandlers conn) $ \listeners -> + forM_ listeners $ \(l, _) -> CE.catch (l reason) $ \(ex :: CE.SomeException) -> + hPutStrLn stderr $ "connection blocked listener threw exception: "++ show ex + + handleUnblocked = do + withMVar (connBlockedHandlers conn) $ \listeners -> + forM_ listeners $ \(_, l) -> CE.catch l $ \(ex :: CE.SomeException) -> + hPutStrLn stderr $ "connection unblocked listener threw exception: "++ show ex + -- | Opens a connection to a broker specified by the given 'ConnectionOpts' parameter. openConnection'' :: ConnectionOpts -> IO Connection openConnection'' connOpts = withSocketsDo $ do @@ -263,9 +276,10 @@ writeLock <- newMVar () ccl <- newEmptyMVar cClosedHandlers <- newMVar [] + cBlockedHandlers <- newMVar [] cLastReceived <- getTimestamp >>= newMVar cLastSent <- getTimestamp >>= newMVar - let conn = Connection handle cChanAllocator cChannels (fromIntegral maxFrameSize) cClosed ccl writeLock cClosedHandlers cLastReceived cLastSent serverProps + let conn = Connection handle cChanAllocator cChannels (fromIntegral maxFrameSize) cClosed ccl writeLock cClosedHandlers cBlockedHandlers cLastReceived cLastSent serverProps --spawn the connectionReceiver connThread <- forkFinally' (connectionReceiver conn) $ \res -> do @@ -333,7 +347,8 @@ , ("capabilities", FVFieldTable clientCapabilities) ] ++ maybe [] (\x -> [("connection_name", FVString x)]) (coName connOpts) - clientCapabilities = FieldTable $ M.fromList $ [ ("consumer_cancel_notify", FVBool True) ] + clientCapabilities = FieldTable $ M.fromList $ [ ("consumer_cancel_notify", FVBool True), + ("connection.blocked", FVBool True) ] handleSecureUntilTune handle sasl = do tuneOrSecure <- readFrame handle @@ -433,6 +448,14 @@ -- otherwise add it to the list _ -> modifyMVar_ (connClosedHandlers conn) $ \old -> return $ handler:old + +-- | @addConnectionBlockedHandler conn blockedHandler unblockedHandler@ adds handlers that will be called +-- when a connection gets blocked/unlocked due to server resource constraints. +-- +-- More information: +addConnectionBlockedHandler :: Connection -> (Text -> IO ()) -> IO () -> IO () +addConnectionBlockedHandler conn blockedHandler unblockedHandler = + modifyMVar_ (connBlockedHandlers conn) $ \old -> return $ (blockedHandler, unblockedHandler):old readFrame :: Conn.Connection -> IO Frame readFrame handle = do diff --git a/amqp.cabal b/amqp.cabal --- a/amqp.cabal +++ b/amqp.cabal @@ -1,5 +1,5 @@ Name: amqp -Version: 0.18.0 +Version: 0.18.1 Synopsis: Client library for AMQP servers (currently only RabbitMQ) Description: Client library for AMQP servers (currently only RabbitMQ) License: BSD3 diff --git a/changelog.md b/changelog.md --- a/changelog.md +++ b/changelog.md @@ -1,3 +1,7 @@ +### Version 0.18.1 + +* new function `addConnectionBlockedHandler` to be notified when the connection is blocked (due to the server being resource-constrained) + ### Version 0.18.0 * `ConnectionClosedException` and `ChannelClosedException` now specify whether the close was normal (user-initiated) or abnormal (caused by an AMQP exception)