@@ -5,9 +5,10 @@ module Network.Transport.TCP
55
66import Network.Transport
77
8−import Control.Concurrent (forkIO, ThreadId, killThread)
8+import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
99import Control.Concurrent.Chan
1010import Control.Concurrent.MVar
11+import Control.Exception (SomeException, throwTo, catch)
1112import Control.Monad (forever, forM_)
1213import Data.ByteString.Lazy.Char8 (ByteString)
1314import Data.IntMap (IntMap)
@@ -19,6 +20,7 @@ import Network.Socket
1920 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
2021 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
2122import Safe
23+import System.IO (stderr, hPutStrLn)
2224
2325import qualified Data.Binary as B
2426import qualified Data.ByteString.Lazy.Char8 as BS
@@ -55,7 +57,8 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
5557 setSocketOption sock ReuseAddr 1
5658 bindSocket sock (addrAddress serverAddr)
5759 listen sock 5
58− threadId <- forkIO $ procConnections chans sock
60+ threadId <- forkWithExceptions forkIO "Connection Listener" $
61+ procConnections chans sock
5962
6063 return Transport
6164 { newConnectionWith = {-# SCC "newConnectionWith" #-}\_ -> do
@@ -142,7 +145,8 @@ procConnections chans sock = forever $ do
142145 putMVar chans (chanId', chanMap)
143146 error "procConnections: cannot find chanId"
144147 Just (chan, socks) -> do
145− threadId <- forkIO $ procMessages chans chanId chan clientSock
148+ threadId <- forkWithExceptions forkIO "Message Listener" $
149+ procMessages chans chanId chan clientSock
146150 let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
147151 putMVar chans (chanId', chanMap')
148152
@@ -157,16 +161,17 @@ procMessages :: Chans -> ChanId -> Chan [ByteString] -> Socket -> IO ()
157161procMessages chans chanId chan sock = do
158162 sizeBSs <- recvExact sock 1
159163 case sizeBSs of
160− [] -> closeSocket
161− _ -> do
162− let size = fromIntegral (B.decode . BS.concat $ sizeBSs :: Word8)
164+ [] -> closeSocket
165+ [onebyte] -> do
166+ let size = fromIntegral (B.decode onebyte :: Word8)
163167 if size == 255
164168 then do
165169 sizeBSs' <- recvExact sock 8
166170 case sizeBSs' of
167171 [] -> closeSocket
168− _ -> procMessage (B.decode . BS.concat $ sizeBSs' :: Int64)
172+ _ -> procMessage (B.decode . BS.concat$ sizeBSs' :: Int64)
169173 else procMessage size
174+ ls -> error "Shouldn't receive more than one bytestring when expecting a single byte!"
170175 where
171176 closeSocket :: IO ()
172177 closeSocket = do
@@ -212,3 +217,15 @@ recvExact sock n = go [] sock n
212217 then return []
213218 else go (bs:bss) sock (n - BS.length bs)
214219
220+
221+
222+forkWithExceptions :: (IO () -> IO ThreadId) -> String -> IO () -> IO ThreadId
223+forkWithExceptions forkit descr action = do
224+ parent <- myThreadId
225+ forkit $
226+ Control.Exception.catch action
227+ (\ e -> do
228+ hPutStrLn stderr $ "Exception inside child thread "++descr++": "++show e
229+ throwTo parent (e::SomeException)
230+ )
231+