@@ -5,9 +5,10 @@ module Network.Transport.TCP
6import Network.Transport
6import Network.Transport
8−import Control.Concurrent (forkIO, ThreadId, killThread)
8+import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
9import Control.Concurrent.Chan
9import Control.Concurrent.Chan
10import Control.Concurrent.MVar
10import Control.Concurrent.MVar
11+import Control.Exception (SomeException, throwTo, catch)
11import Control.Monad (forever, forM_)
12import Control.Monad (forever, forM_)
12import Data.ByteString.Lazy.Char8 (ByteString)
13import Data.ByteString.Lazy.Char8 (ByteString)
13import Data.IntMap (IntMap)
14import Data.IntMap (IntMap)
@@ -19,6 +20,7 @@ import Network.Socket
19 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
20 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
20 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
21 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
21import Safe
22import Safe
23+import System.IO (stderr, hPutStrLn)
23import qualified Data.Binary as B
25import qualified Data.Binary as B
24import qualified Data.ByteString.Lazy.Char8 as BS
26import qualified Data.ByteString.Lazy.Char8 as BS
@@ -55,7 +57,8 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
55 setSocketOption sock ReuseAddr 1
57 setSocketOption sock ReuseAddr 1
56 bindSocket sock (addrAddress serverAddr)
58 bindSocket sock (addrAddress serverAddr)
57 listen sock 5
59 listen sock 5
60+ threadId <- forkWithExceptions forkIO "Connection Listener" $
58− threadId <- forkIO $ procConnections chans sock
61+ procConnections chans sock
60 return Transport
63 return Transport
61 { newConnectionWith = {-# SCC "newConnectionWith" #-}\_ -> do
64 { newConnectionWith = {-# SCC "newConnectionWith" #-}\_ -> do
@@ -142,7 +145,8 @@ procConnections chans sock = forever $ do
142 putMVar chans (chanId', chanMap)
145 putMVar chans (chanId', chanMap)
143 error "procConnections: cannot find chanId"
146 error "procConnections: cannot find chanId"
144 Just (chan, socks) -> do
147 Just (chan, socks) -> do
148+ threadId <- forkWithExceptions forkIO "Message Listener" $
145− threadId <- forkIO $ procMessages chans chanId chan clientSock
149+ procMessages chans chanId chan clientSock
146 let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
150 let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
147 putMVar chans (chanId', chanMap')
151 putMVar chans (chanId', chanMap')
@@ -157,16 +161,17 @@ procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
157procMessages chans chanId chan sock = do
161procMessages chans chanId chan sock = do
158 sizeBSs <- recvExact sock 1
162 sizeBSs <- recvExact sock 1
159 case sizeBSs of
163 case sizeBSs of
160− [] -> closeSocket
164+ [] -> closeSocket
161− _ -> do
165+ [onebyte] -> do
162− let size = fromIntegral (B.decode . BS.concat $ sizeBSs :: Word8)
166+ let size = fromIntegral (B.decode onebyte :: Word8)
163 if size == 255
167 if size == 255
165 sizeBSs' <- recvExact sock 8
169 sizeBSs' <- recvExact sock 8
166 case sizeBSs' of
170 case sizeBSs' of
167 [] -> closeSocket
171 [] -> closeSocket
168− _ -> procMessage (B.decode . BS.concat $ sizeBSs' :: Int64)
172+ _ -> procMessage (B.decode . BS.concat$ sizeBSs' :: Int64)
169 else procMessage size
173 else procMessage size
174+ ls -> error "Shouldn't receive more than one bytestring when expecting a single byte!"
171 closeSocket :: IO ()
176 closeSocket :: IO ()
172 closeSocket = do
177 closeSocket = do
@@ -212,3 +217,15 @@ recvExact sock n = go [] sock n
212 then return []
217 then return []
213 else go (bs:bss) sock (n - BS.length bs)
218 else go (bs:bss) sock (n - BS.length bs)
222+forkWithExceptions :: (IO () -> IO ThreadId) -> String -> IO () -> IO ThreadId
223+forkWithExceptions forkit descr action = do
224+ parent <- myThreadId
226+ Control.Exception.catch action
228+ hPutStrLn stderr $ "Exception inside child thread "++descr++": "++show e
229+ throwTo parent (e::SomeException)