@@ -5,10 +5,10 @@ module Network.Transport.TCP
55
66import Network.Transport
77
8−import Control.Concurrent (forkIO)
8+import Control.Concurrent (forkIO, ThreadId, killThread)
99import Control.Concurrent.Chan
1010import Control.Concurrent.MVar
11−import Control.Monad (forever)
11+import Control.Monad (forever, forM_)
1212import Data.ByteString.Lazy.Char8 (ByteString)
1313import Data.IntMap (IntMap)
1414import Data.Int
@@ -25,8 +25,10 @@ import qualified Data.IntMap as IntMap
2525import qualified Network.Socket as N
2626import qualified Network.Socket.ByteString.Lazy as NBS
2727
28+import Debug.Trace
29+
2830type ChanId = Int
29−type Chans = MVar (ChanId, IntMap (Chan ByteString, MVar [Socket]))
31+type Chans = MVar (ChanId, IntMap (Chan ByteString, [(ThreadId, Socket)]))
3032
3133-- | This deals with several different configuration properties:
3234-- * Buffer size, specified in Hints
@@ -42,7 +44,7 @@ data TCPConfig = TCPConfig Hints HostName ServiceName
4244-- Messages are all queued using an unbounded Chan.
4345mkTransport :: TCPConfig -> IO Transport
4446mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
45− channels <- newMVar (0, IntMap.empty)
47+ chans <- newMVar (0, IntMap.empty)
4648 serverAddrs <- getAddrInfo
4749 (Just (N.defaultHints { addrFlags = [AI_PASSIVE] }))
4850 Nothing
@@ -51,21 +53,31 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
5153 [] -> error "mkTransport: getAddrInfo returned []"
5254 as -> head as
5355 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
56+ setSocketOption sock ReuseAddr 1
5457 bindSocket sock (addrAddress serverAddr)
5558 listen sock 5
56− _ <- forkIO $ procConnections channels sock
59+ threadId <- forkIO $ procConnections chans sock
5760
5861 return Transport
5962 { newConnectionWith = \_ -> do
60− (chanId, chanMap) <- takeMVar channels
63+ (chanId, chanMap) <- takeMVar chans
6164 chan <- newChan
62− socks <- newMVar [sock]
63− putMVar channels (chanId + 1, IntMap.insert chanId (chan, socks) chanMap)
64− return (mkSourceAddr host service chanId, mkTargetEnd chan socks)
65+ putMVar chans (chanId + 1, IntMap.insert chanId (chan, []) chanMap)
66+ return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
6567 , newMulticastWith = error "newMulticastWith: not defined"
6668 , deserialize = \bs ->
6769 let (host, service, chanId) = B.decode bs in
6870 Just $ mkSourceAddr host service chanId
71+ , closeTransport = do
72+ -- Kill the transport channel process
73+ killThread threadId
74+ sClose sock
75+ -- Kill all target end processes
76+ (chanId, chanMap) <- takeMVar chans
77+ forM_ (IntMap.elems chanMap) (\(chan, socks) ->
78+ forM_ socks (\(threadId', sock') -> do
79+ killThread threadId'
80+ sClose sock'))
6981 }
7082
7183 where
@@ -93,13 +105,22 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
93105 , closeSourceEnd = sClose sock
94106 }
95107
96− mkTargetEnd :: Chan ByteString -> MVar [Socket] -> TargetEnd
97− mkTargetEnd chan socks = TargetEnd
108+ mkTargetEnd :: Chans -> ChanId -> Chan ByteString -> TargetEnd
109+ mkTargetEnd chans chanId chan = TargetEnd
98110 { -- for now we will implement this as a Chan
99111 receive = do
100112 bs <- readChan chan
101113 return [bs]
102− , closeTargetEnd = takeMVar socks >>= mapM_ sClose
114+ , closeTargetEnd = do
115+ (chanId', chanMap) <- takeMVar chans
116+ case IntMap.lookup chanId chanMap of
117+ Nothing -> putMVar chans (chanId', chanMap)
118+ Just (_, socks) -> do
119+ forM_ socks $ \(threadId, sock) -> do
120+ killThread threadId
121+ sClose sock
122+ let chanMap' = IntMap.delete chanId chanMap
123+ putMVar chans (chanId', chanMap')
103124 }
104125
105126 procConnections :: Chans -> Socket -> IO ()
@@ -108,36 +129,50 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
108129 -- decode the first message to find the correct chanId
109130 bs <- recvExact clientSock 8
110131 let chanId = fromIntegral (B.decode bs :: Int64)
111− (_, chanMap) <- readMVar chans
132+ (chanId', chanMap) <- takeMVar chans
112133 case IntMap.lookup chanId chanMap of
113− Nothing -> error "procConnections: cannot find chanId"
134+ Nothing -> do
135+ putMVar chans (chanId', chanMap)
136+ error "procConnections: cannot find chanId"
114137 Just (chan, socks) -> do
115− modifyMVar_ socks (return . (clientSock :))
116− forkIO $ procMessages chan clientSock
138+ threadId <- forkIO $ procMessages chans chanId chan clientSock
139+ let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
140+ putMVar chans (chanId', chanMap')
117141
118142 -- This function first extracts a header of type Int64, which determines
119143 -- the size of the ByteString that follows. The ByteString is then
120144 -- extracted from the socket, and then written to the Chan only when
121− -- complete.
122− procMessages :: Chan ByteString -> Socket -> IO ()
123− procMessages chan sock = forever $ do
145+ -- complete. A length of -1 indicates that the socket has finished
146+ -- communicating.
147+ procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
148+ procMessages chans chanId chan sock = do
124149 sizeBS <- recvExact sock 8
125150 let size = B.decode sizeBS
126− bs <- recvExact sock size
127− writeChan chan bs
151+ if size == -1
152+ then do
153+ (chanId', chanMap) <- takeMVar chans
154+ case IntMap.lookup chanId chanMap of
155+ Nothing -> do
156+ putMVar chans (chanId', chanMap)
157+ error "procMessages: chanId not found."
158+ Just (chan, socks) -> do
159+ let socks' = filter ((/= sock) . snd) socks
160+ let chanMap' = IntMap.insert chanId (chan, socks') chanMap
161+ putMVar chans (chanId', chanMap')
162+ sClose sock
163+ else do
164+ bs <- recvExact sock size
165+ writeChan chan bs
166+ procMessages chans chanId chan sock
128167
129−-- The result of `recvExact sock n` is a `ByteString` whose length is exactly
130−-- `n`. No more bytes than necessary are read from the socket.
168+-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
169+-- from `sock`. No more bytes than necessary are read from the socket.
131170-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
132171-- superfluous input depending on the socket type.
133172recvExact :: Socket -> Int64 -> IO ByteString
134−recvExact sock n = do
135− bs <- NBS.recv sock n
136− let remainder = n - BS.length bs
137− if remainder > 0
138− then do
139− bs' <- recvExact sock remainder
140− return (BS.append bs bs')
141− else
142− return bs
143−
173+recvExact sock n = go BS.empty sock n
174+ where
175+ go bs sock 0 = return bs
176+ go bs sock n = do
177+ bs' <- NBS.recv sock n
178+ go (BS.append bs bs') sock (n - BS.length bs')