2 files changed, 33 insertions, 17 deletions

diff --git a/dht/src/Network/Tox/Onion/Routes.hs b/dht/src/Network/Tox/Onion/Routes.hs
index c054b99e..d61c721e 100644
--- a/dht/src/Network/Tox/Onion/Routes.hs
+++ b/dht/src/Network/Tox/Onion/Routes.hs
@@ -14,6 +14,7 @@ import Network.QueryResponse
 import Network.QueryResponse.TCP
 import Network.Tox.NodeId
 import Network.Tox.Onion.Transport as Onion
+import Network.Tox.RelayPinger
 import qualified Data.Tox.Relay as TCP
 import qualified Network.Tox.TCP as TCP
 import qualified TCPProber as TCP
@@ -88,6 +89,7 @@ data OnionRouter = OnionRouter
     , tcpProberThread :: ThreadId
       -- | Kademlia table of TCP relays.
     , tcpBucketRefresher :: BucketRefresher NodeId TCP.NodeInfo
+    , tcpRelayPinger :: RelayPinger
       -- | Debug prints are written to this channel which is then flushed to
       -- 'routeLogger'.
     , routeLog        :: TChan String
@@ -207,7 +209,8 @@ newOnionRouter crypto perror tcp_enabled = do
                 writeTChan rlog $ unwords ["Selected TCP Gateway:",show ni,"via",show gw]
                 return gw
             }
-    or <- atomically $ do
+
+    or0 <- atomically $ do
         -- chan <- newTChan
         drg <- newTVar drg0
         -- forM_ [0..11] $ \n -> writeTChan chan $ BuildRoute (RouteId n)
@@ -243,12 +246,7 @@ newOnionRouter crypto perror tcp_enabled = do
                     }
                 , tcpMode = tcpmode
                 , tcpKademliaClient = tcp
-                    { TCP.tcpClient =
-                        let c = TCP.tcpClient tcp
-                        in c { clientNet = addHandler perror (handleMessage c)
-                                            $ onInbound (updateTCP o)
-                                            $ clientNet c }
-                    }
+                , tcpRelayPinger     = error "forkRouteBuilder: no RelayPinger specified"
                 , tcpBucketRefresher = refresher
                 , routeLog           = rlog
                 , routeThread        = error "forkRouteBuilder not invoked (missing onion route builder thread)."
@@ -259,10 +257,21 @@ newOnionRouter crypto perror tcp_enabled = do
                 , routeLogger        = perror
                 }
         return o
+    pinger <- forkRelayPinger (kademSpace $ refreshKademlia $ tcpBucketRefresher or0)
+                              (TCP.tcpClient tcp)
+    let or = or0 { tcpRelayPinger    = pinger
+                 , tcpKademliaClient = tcp
+                    { TCP.tcpClient =
+                        let c = TCP.tcpClient tcp
+                        in c { clientNet = addHandler perror (handleMessage c)
+                                            $ onInbound (updateTCP or)
+                                            $ clientNet c }
+                    }
+                 }
     return (or,tcptbl)
 
 updateTCP :: OnionRouter -> TCP.NodeInfo -> p -> IO ()
-updateTCP or addr x =
+updateTCP or addr x = do
     let refresher = tcpBucketRefresher or
         kademlia0 = refreshKademlia refresher
         kademlia = kademlia0 { kademIO = (kademIO kademlia0)
@@ -275,7 +284,8 @@ updateTCP or addr x =
                                     tblTransition (kademIO kademlia0) tr
                                 }
                              }
-    in insertNode kademlia addr
+    atomically $ bumpRelay (tcpRelayPinger or) addr
+    insertNode kademlia addr
 
 selectGateway :: TVar (R.BucketList TCP.NodeInfo) -> NodeInfo -> STM (Maybe TCP.NodeInfo)
 selectGateway tbl ni = do
diff --git a/dht/src/Network/Tox/RelayPinger.hs b/dht/src/Network/Tox/RelayPinger.hs
index 00c6f65a..5803c89f 100644
--- a/dht/src/Network/Tox/RelayPinger.hs
+++ b/dht/src/Network/Tox/RelayPinger.hs
@@ -58,12 +58,15 @@ addRelay :: RelayPinger -> TCP.NodeInfo -> STM ()
 addRelay (RelayPinger a kad cli que self) ni = do
     let nid = kademliaLocation kad ni
         nidk = decodeAnnounceKey a $ S.encode nid
-        pingit = ScheduledItem $ \a nidk tm -> do
-            scheduleRel a nidk pingit pinginterval -- Reschedule the next ping.
-            return $ do
-                _ <- tcpPing cli ni
-                -- TODO: Remove after ping timeout?
-                return ()
+        pingit = ScheduledItem $ \a nidk tm -> return $ do
+            m <- tcpPing cli ni
+            case m of
+                Just () -> atomically $ scheduleRel a nidk pingit pinginterval -- Reschedule the next ping.
+                Nothing -> do
+                    -- dput XMisc $ "relay-ping: Ping timeout for " ++ show ni
+                    atomically $ modifyTVar' que $
+                        \(n,q) -> ( (,) $! (if PSQ.member ni q then n - 1 else n) )
+                                  $ PSQ.delete ni q
     modifyTVar' que $ \(n,q) ->
                     ( (,) $! (if PSQ.member ni q then n else n+1) )
                     $ PSQ.insert
@@ -85,14 +88,17 @@ bumpRelay :: RelayPinger -> TCP.NodeInfo -> STM ()
 bumpRelay rp@(RelayPinger a kad cli que self) ni = do
     let nid = kademliaLocation kad ni
     (cnt,q) <- readTVar que
-    if cnt < 4
+    if cnt < 4 || PSQ.member ni q
         then addRelay rp ni
         else case PSQ.minView q of
             Nothing            -> addRelay rp ni
             Just (r :-> Down p, q')
                 | let pnew = RelayPriority (rankTCPPort ni) (kademliaXor kad self nid)
                 , pnew < p
-                -> delRelay rp r >> addRelay rp ni
+                -> do
+                    writeTVar que (3,q')
+                    unschedule a (decodeAnnounceKey a $ S.encode $ kademliaLocation kad r)
+                    addRelay rp ni
             _ -> return ()
 
 currentRelays :: RelayPinger -> STM (Int,[TCP.NodeInfo])