@@ -27,7 +27,7 @@ import qualified Network.Socket as N
2727import qualified Network.Socket.ByteString.Lazy as NBS
2828
2929type ChanId = Int
30−type Chans = MVar (ChanId, IntMap (Chan ByteString, [(ThreadId, Socket)]))
30+type Chans = MVar (ChanId, IntMap (Chan [ByteString], [(ThreadId, Socket)]))
3131
3232-- | This deals with several different configuration properties:
3333-- * Buffer size, specified in Hints
@@ -58,16 +58,16 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
5858 threadId <- forkIO $ procConnections chans sock
5959
6060 return Transport
61− { newConnectionWith = \_ -> do
61+ { newConnectionWith = {-# SCC "newConnectionWith" #-}\_ -> do
6262 (chanId, chanMap) <- takeMVar chans
6363 chan <- newChan
6464 putMVar chans (chanId + 1, IntMap.insert chanId (chan, []) chanMap)
6565 return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
6666 , newMulticastWith = error "newMulticastWith: not defined"
67− , deserialize = \bs ->
67+ , deserialize = {-# SCC "deserialize" #-} \bs ->
6868 let (host, service, chanId) = B.decode bs in
6969 Just $ mkSourceAddr host service chanId
70− , closeTransport = do
70+ , closeTransport = {-# SCC "closeTransport" #-} do
7171 -- Kill the transport channel process
7272 killThread threadId
7373 sClose sock
@@ -81,8 +81,8 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
8181
8282mkSourceAddr :: HostName -> ServiceName -> ChanId -> SourceAddr
8383mkSourceAddr host service chanId = SourceAddr
84− { connectWith = \_ -> mkSourceEnd host service chanId
85− , serialize = B.encode (host, service, chanId)
84+ { connectWith = {-# SCC "connectWith" #-} \_ -> mkSourceEnd host service chanId
85+ , serialize = {-# SCC "serialize" #-} B.encode (host, service, chanId)
8686 }
8787
8888mkSourceEnd :: HostName -> ServiceName -> ChanId -> IO SourceEnd
@@ -96,7 +96,7 @@ mkSourceEnd host service chanId = withSocketsDo $ do
9696 N.connect sock (addrAddress serverAddr)
9797 NBS.sendAll sock $ B.encode (fromIntegral chanId :: Int64)
9898 return SourceEnd
99− { send = \bss -> do
99+ { send = {-# SCC "send" #-} \bss -> do
100100 let size = fromIntegral (sum . map BS.length $ bss) :: Int64
101101 if size < 255
102102 then
@@ -105,16 +105,14 @@ mkSourceEnd host service chanId = withSocketsDo $ do
105105 NBS.sendAll sock (B.encode (255 :: Word8))
106106 NBS.sendAll sock (B.encode size)
107107 mapM_ (NBS.sendAll sock) bss
108− , closeSourceEnd = sClose sock
108+ , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
109109 }
110110
111−mkTargetEnd :: Chans -> ChanId -> Chan ByteString -> TargetEnd
111+mkTargetEnd :: Chans -> ChanId -> Chan [ByteString] -> TargetEnd
112112mkTargetEnd chans chanId chan = TargetEnd
113113 { -- for now we will implement this as a Chan
114− receive = do
115− bs <- readChan chan
116− return [bs]
117− , closeTargetEnd = do
114+ receive = {-# SCC "receive" #-} readChan chan
115+ , closeTargetEnd = {-# SCC "closeTargetEnd" #-} do
118116 (chanId', chanMap) <- takeMVar chans
119117 case IntMap.lookup chanId chanMap of
120118 Nothing -> putMVar chans (chanId', chanMap)
@@ -132,10 +130,11 @@ procConnections :: Chans -> Socket -> IO ()
132130procConnections chans sock = forever $ do
133131 (clientSock, _clientAddr) <- accept sock
134132 -- decode the first message to find the correct chanId
135− mBs <- recvExact clientSock 8
136− case mBs of
137− Nothing -> error "procConnections: inbound chanId aborted"
138− Just bs -> do
133+ bss <- recvExact clientSock 8
134+ case bss of
135+ [] -> error "procConnections: inbound chanId aborted"
136+ bss -> do
137+ let bs = BS.concat bss
139138 let chanId = fromIntegral (B.decode bs :: Int64)
140139 (chanId', chanMap) <- takeMVar chans
141140 case IntMap.lookup chanId chanMap of
@@ -154,21 +153,19 @@ procConnections chans sock = forever $ do
154153-- extracted from the socket, and then written to the Chan only when
155154-- complete. If either of the first header size is null this indicates the
156155-- socket has closed.
157−procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
156+procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
158157procMessages chans chanId chan sock = do
159− mSizeBS <- recvExact sock 1
160− case mSizeBS of
161− Nothing -> closeSocket
162− Just sizeBS -> do
163− let size = fromIntegral (B.decode sizeBS :: Word8)
158+ sizeBSs <- recvExact sock 1
159+ case sizeBSs of
160+ [] -> closeSocket
161+ _ -> do
162+ let size = fromIntegral (B.decode . BS.concat $ sizeBSs :: Word8)
164163 if size == 255
165164 then do
166− mSizeBS' <- recvExact sock 8
167− case mSizeBS' of
168− Nothing -> closeSocket
169− Just sizeBS' -> do
170− let size' = B.decode sizeBS' :: Int64
171− procMessage size'
165+ sizeBSs' <- recvExact sock 8
166+ case sizeBSs' of
167+ [] -> closeSocket
168+ _ -> procMessage (B.decode . BS.concat $ sizeBSs' :: Int64)
172169 else procMessage size
173170 where
174171 closeSocket :: IO ()
@@ -185,31 +182,33 @@ procMessages chans chanId chan sock = do
185182 sClose sock
186183 procMessage :: Int64 -> IO ()
187184 procMessage size = do
188− mBs <- recvExact sock size
189− case mBs of
190− Nothing -> closeSocket
191− Just bs -> do
192− writeChan chan bs
185+ bss <- recvExact sock size
186+ case bss of
187+ [] -> closeSocket
188+ _ -> do
189+ writeChan chan bss
193190 procMessages chans chanId chan sock
194191
195−-- | The result of `recvExact sock n` is a `Maybe ByteString` of length `n`,
196−-- received from `sock`. No more bytes than necessary are read from the socket.
192+-- | The result of `recvExact sock n` is a `[ByteString]` whose concatenation
193+-- is of length `n`, received from `sock`. No more bytes than necessary are
194+-- read from the socket.
197195-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
198196-- superfluous input depending on the socket type. Also note that
199197-- if `recv` returns an empty `ByteString` then this means that the socket
200−-- was closed: in this case, we return the empty `ByteString`.
198+-- was closed: in this case, we return an empty list.
201199-- NB: It may be appropriate to change the return type to
202200-- IO (Either ByteString ByteString), where `return Left bs` indicates
203201-- a partial retrieval since the socket was closed, and `return Right bs`
204202-- indicates success. This hasn't been implemented, since a Transport `receive`
205203-- represents an atomic receipt of a message.
206−recvExact :: Socket -> Int64 -> IO (Maybe ByteString)
207−recvExact sock n = go BS.empty sock n
204+recvExact :: Socket -> Int64 -> IO [ByteString]
205+recvExact sock n = go [] sock n
208206 where
209− go bs sock 0 = return (Just bs)
210− go bs sock n = do
211− bs' <- NBS.recv sock n
212− if BS.null bs'
213− then return Nothing
214− else go (BS.append bs bs') sock (n - BS.length bs')
207+ go :: [ByteString] -> Socket -> Int64 -> IO [ByteString]
208+ go bss _ 0 = return (reverse bss)
209+ go bss sock n = do
210+ bs <- NBS.recv sock n
211+ if BS.null bs
212+ then return []
213+ else go (bs:bss) sock (n - BS.length bs)
215214