1 files changed, 21 insertions, 11 deletions

diff --git a/src/Network/QueryResponse/TCP.hs b/src/Network/QueryResponse/TCP.hs
index 328a19e1..efeab305 100644
--- a/src/Network/QueryResponse/TCP.hs
+++ b/src/Network/QueryResponse/TCP.hs
@@ -19,9 +19,12 @@ import Data.Time.Clock.POSIX
 import Data.Word
 import Network.BSD
 import Network.Socket
+import System.Timeout
 import System.IO
 import System.IO.Error
 
+import DebugTag
+import DPut
 import Connection.Tcp (socketFamily)
 import qualified Data.MinMaxPSQ as MM
 import Network.QueryResponse
@@ -61,11 +64,12 @@ acquireConnection :: MVar (Maybe (Either a (x, addr)))
                            -> TCPCache (SessionProtocol x y)
                            -> StreamHandshake addr x y
                            -> addr
+                           -> Bool
                            -> IO (Maybe (y -> IO ()))
-acquireConnection mvar tcpcache stream addr = do
+acquireConnection mvar tcpcache stream addr bDoCon = do
     cache <- atomically $ readTVar (lru tcpcache)
     case MM.lookup' (TCPAddress $ streamAddr stream addr) cache of
-        Nothing -> do
+        Nothing -> fmap join $ forM (guard bDoCon) $ \() -> do
             proto <- getProtocolNumber "tcp"
             mh <- catchIOError (do sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
                                    connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
@@ -77,10 +81,12 @@ acquireConnection mvar tcpcache stream addr = do
             t <- getPOSIXTime
             rthread <- forkIO $ fix $ \loop -> do
                 x <- streamDecode st
-                putMVar mvar $ fmap (\u -> Right (u, addr)) x
                 case x of
-                    Just _ -> loop
+                    Just u -> do
+                        timeout (1000000) $ putMVar mvar $ Just $ Right (u, addr)
+                        loop
                     Nothing -> do
+                        dput XTCP $ "TCP disconnected: " ++ show (streamAddr stream addr)
                         atomically $ modifyTVar' (lru tcpcache)
                                    $ MM.delete (TCPAddress $ streamAddr stream addr)
                         hClose h
@@ -91,14 +97,18 @@ acquireConnection mvar tcpcache stream addr = do
                     , tcpState  = st
                     , tcpThread = rthread
                     }
-            let (retires,cache') = MM.takeView (tcpMax tcpcache)
-                                    $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down t) cache
+            retires <- atomically $ do
+                c <- readTVar (lru tcpcache)
+                let (rs,c') = MM.takeView (tcpMax tcpcache)
+                            $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down t) c
+                writeTVar (lru tcpcache) c'
+                return rs
             forM_ retires $ \(MM.Binding (TCPAddress k) r _) -> void $ forkIO $ do
                 myThreadId >>= flip labelThread ("tcp-close:"++show k)
+                dput XTCP $ "TCP dropped: " ++ show k
                 killThread (tcpThread r)
                 streamGoodbye st
                 hClose (tcpHandle r)
-                atomically $ modifyTVar' (lru tcpcache) $ MM.delete (TCPAddress k)
 
             return $ Just $ streamEncode st
         Just (tm,v) -> do
@@ -119,15 +129,15 @@ closeAll tcpcache stream = do
 
 tcpTransport :: Int -- ^ maximum number of TCP links to maintain.
                 -> StreamHandshake addr x y
-                -> IO (TransportA err addr x y)
+                -> IO (TransportA err addr x (Bool,y))
 tcpTransport maxcon stream = do
     msgvar <- newEmptyMVar
     tcpcache <- atomically $ (`TCPCache` maxcon) <$> newTVar (MM.empty)
     return Transport
-        { awaitMessage   = (takeMVar msgvar >>=)
-        , sendMessage    = \addr y -> do
+        { awaitMessage   = \f -> takeMVar msgvar >>= \x -> f x `catchIOError` (\e -> dput XTCP ("TCP transport stopped. " ++ show e) >> f Nothing)
+        , sendMessage    = \addr (bDoCon,y) -> do
                                 t <- forkIO $ do
-                                    msock <- acquireConnection msgvar tcpcache stream addr
+                                    msock <- acquireConnection msgvar tcpcache stream addr bDoCon
                                     mapM_ ($ y) msock
                                   `catchIOError` \e -> return ()
                                 labelThread t "tcp-send"