@@ -1,3 +1,5 @@
1+{-# LANGUAGE CPP #-}
2+
13module Network.Transport.TCP
24 ( mkTransport
35 , TCPConfig (..)
@@ -12,19 +14,18 @@ import Control.Concurrent.Chan
1214import Control.Concurrent.MVar
1315import Control.Monad (forever, forM_)
1416
15−import Data.ByteString.Char8 (ByteString)
16−import qualified Data.ByteString.Char8 as BS
17−import Data.ByteString.Internal
17+import qualified Data.ByteString.Internal as BSI
1818
1919import Data.IntMap (IntMap)
2020import qualified Data.IntMap as IntMap
2121
22−import Data.Int
23−import Data.Serialize
22+import Data.Int (Int32)
23+import qualified Data.Serialize as Ser
2424import Data.Word
2525
2626import Foreign.ForeignPtr (withForeignPtr)
2727import Foreign.Ptr (Ptr, plusPtr)
28+import Foreign.Storable (sizeOf)
2829
2930import Network.Socket
3031 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
@@ -32,14 +33,30 @@ import Network.Socket
3233 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
3334 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
3435import qualified Network.Socket as N
35−import qualified Network.Socket.ByteString as NBS
3636
3737import Text.Printf
3838
3939import Safe
4040
41−import qualified Data.Binary as B
42−
41+#ifndef LAZY
42+import Data.ByteString.Char8 (ByteString)
43+import qualified Data.ByteString.Char8 as BS
44+import qualified Network.Socket.ByteString as NBS
45+sendNBS = NBS.sendMany
46+encode = Ser.encode
47+decode = Ser.decode
48+#else
49+import Data.ByteString.Lazy.Char8 (ByteString)
50+import qualified Data.ByteString.Lazy.Char8 as BS
51+import qualified Network.Socket.ByteString.Lazy as NBS
52+sendNBS sock bs = NBS.sendAll sock $! head bs
53+encode = Ser.encodeLazy
54+decode = Ser.decodeLazy
55+#endif
56+{-# INLINE encode #-}
57+{-# INLINE decode #-}
58+encode :: Ser.Serialize a => a -> ByteString
59+decode :: Ser.Serialize a => ByteString -> Either String a
4360
4461type ChanId = Int
4562type Chans = MVar (ChanId, IntMap (Chan [ByteString], [(ThreadId, Socket)]))
@@ -112,13 +129,13 @@ mkSourceEnd host service chanId = withSocketsDo $ do
112129 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
113130 setSocketOption sock ReuseAddr 1
114131 N.connect sock (addrAddress serverAddr)
115− NBS.sendAll sock $ encode (fromIntegral chanId :: Int64)
132+ NBS.sendAll sock $ encode (fromIntegral chanId :: Int32)
116133 return SourceEnd
117134 { send = {-# SCC "send" #-} \bss -> do
118− let size = fromIntegral (sum . map BS.length $ bss) :: Int64
135+ let size = fromIntegral (sum . map BS.length $ bss) :: Int32
119136 sizeStr | size < 255 = encode (fromIntegral size :: Word8)
120137 | otherwise = BS.cons (toEnum 255) (encode size)
121− NBS.sendMany sock (sizeStr:bss)
138+ sendNBS sock (sizeStr:bss)
122139 , closeSourceEnd = {-# SCC "closeSourceEnd" #-} sClose sock
123140 }
124141
@@ -145,11 +162,11 @@ procConnections chans sock = forever $ do
145162 let err m = error $ printf "procConnections: %s" m
146163 (clientSock, _clientAddr) <- accept sock
147164 -- decode the first message to find the correct chanId
148− bss <- recvExact clientSock 8
165+ bss <- recvExact clientSock (fromIntegral $ sizeOf (0 :: Int32))
149166 case bss of
150167 Left _ -> err "inbound chanId aborted"
151168 Right bs -> do
152− let chanId = either err fromIntegral (decode bs :: Either String Int64)
169+ let chanId = either err fromIntegral (decode bs :: Either String Int32)
153170 (chanId', chanMap) <- takeMVar chans
154171 case IntMap.lookup chanId chanMap of
155172 Nothing -> do
@@ -175,7 +192,7 @@ procMessages chans chanId chan sock = do
175192 Right sizeBS -> do
176193 case decode sizeBS :: Either String Word8 of
177194 Right 255 -> do
178− esizeBS' <- recvExact sock 8
195+ esizeBS' <- recvExact sock (fromIntegral $ sizeOf (0 :: Int32))
179196 case esizeBS' of
180197 Left _ -> closeSocket
181198 Right sizeBS' -> procMessage $ decode sizeBS'
@@ -194,7 +211,7 @@ procMessages chans chanId chan sock = do
194211 let chanMap' = IntMap.insert chanId (chan, socks') chanMap
195212 putMVar chans (chanId', chanMap')
196213 sClose sock
197− procMessage :: Either String Int64 -> IO ()
214+ procMessage :: Either String Int32 -> IO ()
198215 procMessage (Left msg) = err msg
199216 procMessage (Right size) = do
200217 ebs <- recvExact sock size
@@ -216,22 +233,33 @@ procMessages chans chanId chan sock = do
216233-- a partial retrieval since the socket was closed, and `return Right bs`
217234-- indicates success. This hasn't been implemented, since a Transport `receive`
218235-- represents an atomic receipt of a message.
219−recvExact :: Socket -> Int64 -> IO (Either ByteString ByteString)
236+#ifndef LAZY
237+recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
220238recvExact sock l = do
221− res <- createAndTrim (fromIntegral l) (go 0)
239+ res <- BSI.createAndTrim (fromIntegral l) (go 0)
222240 case BS.length res of
223241 n | n == (fromIntegral l) -> return $ Right res
224− n -> do
225− printf "recvExact: expected %dB, got %dB\n" l n
226− return $ Left res
242+ n -> return $ Left res
227243 where
228244 go :: Int -> Ptr Word8 -> IO Int
229245 go n ptr | n == (fromIntegral l) = return n
230246 go n ptr = do
231− (p, off, len) <- toForeignPtr <$> NBS.recv sock (fromIntegral l)
247+ (p, off, len) <- BSI.toForeignPtr <$> NBS.recv sock (min (fromIntegral l-n) 4096)
232248 if len == 0
233249 then return n
234250 else withForeignPtr p $ \p -> do
235− memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)
251+ BSI.memcpy (ptr `plusPtr` n) (p `plusPtr` off) (fromIntegral len)
236252 go (n+len) ptr
237−
253+#else
254+recvExact :: Socket -> Int32 -> IO (Either ByteString ByteString)
255+recvExact sock n =
256+ go [] sock n
257+ where
258+ go :: [ByteString] -> Socket -> Int32 -> IO (Either ByteString ByteString)
259+ go bss _ 0 = return (Right $ BS.concat (reverse bss))
260+ go bss sock n = do
261+ bs <- NBS.recv sock (min (fromIntegral n) 4096)
262+ if BS.null bs
263+ then return (Left $ BS.concat (reverse bss))
264+ else go (bs:bss) sock (n - (fromIntegral $ BS.length bs))
265+#endif