1 files changed, 14 insertions, 8 deletions

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])