1 files changed, 11 insertions, 4 deletions

diff --git a/dht/src/Network/QueryResponse/TCP.hs b/dht/src/Network/QueryResponse/TCP.hs
index 45ff73a6..defdf6a2 100644
--- a/dht/src/Network/QueryResponse/TCP.hs
+++ b/dht/src/Network/QueryResponse/TCP.hs
@@ -55,6 +55,7 @@ data TCPCache st = TCPCache
     , tcpMax :: Int
     }
 
+-- This is a suitable /st/ parameter to 'TCPCache'
 data SessionProtocol x y = SessionProtocol
     { streamGoodbye :: IO ()        -- ^ "Goodbye" protocol upon termination.
     , streamDecode  :: IO (Maybe x) -- ^ Parse inbound messages.
@@ -74,6 +75,9 @@ showStat :: IsString p => TCPSession st -> p
 showStat r = case r of PendingTCPSession -> "pending."
                        TCPSession {}     -> "established."
 
+tcp_timeout :: Int
+tcp_timeout = 10000000
+
 acquireConnection :: MVar (Maybe (Either a (x, addr)))
                            -> TCPCache (SessionProtocol x y)
                            -> StreamHandshake addr x y
@@ -96,7 +100,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
     case entry of
         Nothing -> fmap join $ forM (guard bDoCon) $ \() -> do
             proto <- getProtocolNumber "tcp"
-            mh <- catchIOError (do h <- timeout 10000000 $ do
+            mh <- catchIOError (do h <- timeout tcp_timeout $ do
                                        sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
                                        connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
                                        h <- socketToHandle sock ReadWriteMode
@@ -104,6 +108,9 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                                        return h
                                    return h)
                                $ \e -> return Nothing
+            when (isNothing mh)
+              $ atomically $ modifyTVar' (lru tcpcache)
+                           $ MM.delete (TCPAddress $ streamAddr stream addr)
             ret <- fmap join $ forM mh $ \h -> do
              mst <- catchIOError (Just <$> streamHello stream addr h)
                                  (\e -> return Nothing)
@@ -122,7 +129,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                     dput XTCP $ "TCP streamDecode " ++ show (streamAddr stream addr) ++ " --> " ++ maybe "Nothing" (const "got") x
                     case x of
                         Just u -> do
-                            m <- timeout (1000000) $ putMVar mvar $ Just $ Right (u, addr)
+                            m <- timeout tcp_timeout $ putMVar mvar $ Just $ Right (u, addr)
                             when (isNothing m) $ do
                                 dput XTCP $ "TCP "++show (streamAddr stream addr) ++ " dropped packet."
                                 tryTakeMVar mvar
@@ -136,7 +143,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                             now <- getPOSIXTime
                             forM_ (zip [1..] $ MM.toList c) $ \(i,MM.Binding (TCPAddress addr) r (Down tm)) -> do
                                 dput XTCP $ unwords [show i ++ ".", "Still connected:",  show addr, show (now - tm), showStat r]
-                            mreport <- timeout 10000000 $ threadReport False -- XXX: Paranoid timeout
+                            mreport <- timeout tcp_timeout $ threadReport False -- XXX: Paranoid timeout
                             case mreport of
                                 Just treport -> dput XTCP treport
                                 Nothing      -> dput XTCP "TCP ERROR: threadReport timed out."
@@ -168,7 +175,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
             return ret
         Just (tm, PendingTCPSession)
             | not bDoCon -> return Nothing
-            | otherwise  -> fmap join $ timeout 10000000 $ atomically $ do
+            | otherwise  -> fmap join $ timeout tcp_timeout $ atomically $ do
                 c <- readTVar (lru tcpcache)
                 let v = MM.lookup' (TCPAddress $ streamAddr stream addr) c
                 case v of