amqp 0.8.2 → 0.8.3

raw patch · 4 files changed

+81/−22 lines, 4 filesdep +vector

Dependencies added: vector

Files

• Network/AMQP.hs +1/−0
• Network/AMQP/ChannelAllocator.hs +50/−0
• Network/AMQP/Internal.hs +25/−19
• amqp.cabal +5/−3

@@ -54,6 +54,7 @@     openConnection,
     openConnection',
     openConnection'',
+    closeChannel,
     closeConnection,
     addConnectionClosedHandler,
 

@@ -0,0 +1,50 @@+module Network.AMQP.ChannelAllocator where
+
+import qualified Data.Vector.Mutable as V
+import Data.Word
+import Data.Bits
+
+newtype ChannelAllocator = ChannelAllocator (V.IOVector Word64)
+
+newChannelAllocator :: IO ChannelAllocator
+newChannelAllocator =
+    fmap ChannelAllocator $ V.replicate 1024 0
+
+allocateChannel :: ChannelAllocator -> IO Int
+allocateChannel (ChannelAllocator c) = do
+    maybeIx <- findFreeIndex c
+    case maybeIx of
+        Just chunk -> do
+            word <- V.read c chunk
+            let offset = findUnsetBit word
+            V.write c chunk (setBit word offset)
+            return (chunk*64 + offset)
+        Nothing -> error "all channels are already allocated"
+
+freeChannel :: ChannelAllocator -> Int -> IO Bool
+freeChannel (ChannelAllocator c) ix = do
+    let (chunk, offset) = divMod ix 64
+    word <- V.read c chunk
+    if testBit word offset
+        then do
+            V.write c chunk $ clearBit word offset
+            return True
+        else return False
+
+findUnsetBit :: Word64 -> Int
+findUnsetBit w = go 0
+  where
+    go 65 = error "findUnsetBit"
+    go ix | not (testBit w ix) = ix
+    go ix = go (ix+1)
+
+findFreeIndex :: V.IOVector Word64 -> IO (Maybe Int)
+findFreeIndex vec = go 0
+  where
+    -- TODO: make this faster
+    go 1024 = return Nothing
+    go ix = do
+        v <- V.read vec ix
+        if v /= 0xffffffffffffffff
+            then return $ Just ix
+            else go $! ix+1

@@ -28,6 +28,7 @@ import Network.AMQP.Types
 import Network.AMQP.Helpers
 import Network.AMQP.Generated
+import Network.AMQP.ChannelAllocator
 
 
 data DeliveryMode = Persistent -- ^ the message will survive server restarts (if the queue is durable)
@@ -127,13 +128,13 @@ 
 data Connection = Connection {
                     connHandle :: Conn.Connection,
-                    connChannels :: (MVar (IM.IntMap (Channel, ThreadId))), --open channels (channelID => (Channel, ChannelThread))
+                    connChanAllocator :: ChannelAllocator,
+                    connChannels :: MVar (IM.IntMap (Channel, ThreadId)), -- open channels (channelID => (Channel, ChannelThread))
                     connMaxFrameSize :: Int, --negotiated maximum frame size
                     connClosed :: MVar (Maybe String),
                     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 ()],
-                    lastChannelID :: MVar Int, -- for auto-incrementing the channelIDs
                     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
                 }
@@ -235,14 +236,15 @@ 
     --build Connection object
     cChannels <- newMVar IM.empty
-    lastChanID <- newMVar 0
     cClosed <- newMVar Nothing
+    cChanAllocator <- newChannelAllocator
+    _ <- allocateChannel cChanAllocator -- allocate channel 0
     writeLock <- newMVar ()
     ccl <- newEmptyMVar
     cClosedHandlers <- newMVar []
     cLastReceived <- getTimestamp >>= newMVar
     cLastSent <- getTimestamp >>= newMVar
-    let conn = Connection handle cChannels (fromIntegral maxFrameSize) cClosed ccl writeLock cClosedHandlers lastChanID cLastReceived cLastSent
+    let conn = Connection handle cChanAllocator cChannels (fromIntegral maxFrameSize) cClosed ccl writeLock cClosedHandlers cLastReceived cLastSent
 
     --spawn the connectionReceiver
     connThread <- forkIO $ CE.finally (connectionReceiver conn) $ do
@@ -539,11 +541,14 @@ -- closes the channel internally; but doesn't tell the server
 closeChannel' :: Channel -> Text -> IO ()
 closeChannel' c reason = do
-    modifyMVar_ (connChannels $ connection c) $ \old -> return $ IM.delete (fromIntegral $ channelID c) old
-    -- mark channel as closed
     modifyMVar_ (chanClosed c) $ \x -> do
         if isNothing x
             then do
+                modifyMVar_ (connChannels $ connection c) $ \old -> do
+                    ret <- freeChannel (connChanAllocator $ connection c) $ fromIntegral $ channelID c
+                    when (not ret) $ putStrLn "closeChannel error: channel already freed"
+                    return $ IM.delete (fromIntegral $ channelID c) old
+
                 void $ killLock $ chanActive c
                 killOutstandingResponses $ outstandingResponses c
                 return $ Just $ maybe (T.unpack reason) id x
@@ -556,32 +561,33 @@             return undefined
 
 -- | opens a new channel on the connection
---
--- There's currently no closeChannel method, but you can always just close the connection (the maximum number of channels is 65535).
 openChannel :: Connection -> IO Channel
 openChannel c = do
     newInQueue <- newChan
     outRes <- newMVar Seq.empty
     lastConsTag <- newMVar 0
     ca <- newLock
-
     closed <- newMVar Nothing
     conss <- newMVar M.empty
     listeners <- newMVar []
 
-    --get a new unused channelID
-    newChannelID <- modifyMVar (lastChannelID c) $ \x -> return (x+1, x+1)
-
-    let newChannel = Channel c newInQueue outRes (fromIntegral newChannelID) lastConsTag ca closed conss listeners
-
-    thrID <- forkIO $ CE.finally (channelReceiver newChannel)
-        (closeChannel' newChannel "closed")
-
     --add new channel to connection's channel map
-    modifyMVar_ (connChannels c) (return . IM.insert newChannelID (newChannel, thrID))
+    newChannel <- modifyMVar (connChannels c) $ \mp -> do
+        newChannelID <- allocateChannel (connChanAllocator c)
+        let newChannel = Channel c newInQueue outRes (fromIntegral newChannelID) lastConsTag ca closed conss listeners
+        thrID <- forkIO $ CE.finally (channelReceiver newChannel)
+            (closeChannel' newChannel "closed")
+        when (IM.member newChannelID mp) $ CE.throwIO $ userError "openChannel fail: channel already open"
+        return (IM.insert newChannelID (newChannel, thrID) mp, newChannel)
 
-    (SimpleMethod (Channel_open_ok _)) <- request newChannel (SimpleMethod (Channel_open (ShortString "")))
+    SimpleMethod (Channel_open_ok _) <- request newChannel (SimpleMethod (Channel_open (ShortString "")))
     return newChannel
+
+-- | closes a channel. It is typically not necessary to manually call this as closing a connection will implicitly close all channels.
+closeChannel :: Channel -> IO ()
+closeChannel c = do
+    SimpleMethod Channel_close_ok <- request c $ SimpleMethod $ Channel_close 0 (ShortString "") 0 0
+    closeChannel' c "user called closeChannel"
 
 -- | writes multiple frames to the channel atomically
 writeFrames :: Channel -> [FramePayload] -> IO ()

@@ -1,5 +1,5 @@ Name:                amqp
-Version:             0.8.2
+Version:             0.8.3
 Synopsis:            Client library for AMQP servers (currently only RabbitMQ)
 Description:         Client library for AMQP servers (currently only RabbitMQ)
                      .
@@ -29,9 +29,10 @@     split>=0.2,
     clock >= 0.4.0.1,
     monad-control >= 0.3,
-    connection == 0.2.*
+    connection == 0.2.*,
+    vector
   Exposed-modules:    Network.AMQP, Network.AMQP.Types, Network.AMQP.Lifted
-  Other-modules:      Network.AMQP.Generated, Network.AMQP.Helpers, Network.AMQP.Protocol, Network.AMQP.Internal
+  Other-modules:      Network.AMQP.ChannelAllocator, Network.AMQP.Generated, Network.AMQP.Helpers, Network.AMQP.Protocol, Network.AMQP.Internal
   GHC-Options:        -Wall
 
 Executable amqp-builder
@@ -76,3 +77,4 @@     , hspec              >= 1.3
     , hspec-expectations >= 0.3.3
     , connection == 0.2.*
+    , vector