Avoid overfilling buckets in DHT routing table.

1 files changed, 24 insertions, 4 deletions

diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index c230929e..38207be5 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -268,27 +268,47 @@ insertBucket curTime (TryInsert info) bucket
   -- then one is replaced by the new node in the next insertBucket
   -- iteration.
   | not (L.null stales)
-    = pure ( [CheckPing stales],  map_q $ pushBack bucketQ info )
+    = pure ( [CheckPing stales]
+           , bucket { -- Update timestamps so that we don't redundantly ping.
+                      bktNodes = updateStamps curTime stales $ bktNodes bucket
+                      -- Update queue with the pending NodeInfo in case of ping fail.
+                    , bktQ     = runIdentity $ pushBack bucketQ info $ bktQ bucket } )
   -- When the bucket is full of good nodes, the new node is simply discarded.
   -- We must return 'A.empty' here to ensure that bucket splitting happens
   -- inside 'modifyBucket'.
   | otherwise = A.empty
  where
-    stales = map key $ PSQ.atMost (curTime - delta) $ bktNodes bucket
+    -- We (take 1) to keep a 1-to-1 correspondence between pending pings and
+    -- waiting nodes in the bktQ.  This way, we don't have to worry about what
+    -- to do with failed pings for which there is no ready replacements.
+    stales = -- One stale:
+             do (n :-> t) <- maybeToList $ PSQ.findMin (bktNodes bucket)
+                guard (t < curTime - delta)
+                return n
+             -- All stale:
+             -- map key $ PSQ.atMost (curTime - delta) $ bktNodes bucket
 
     already_have = maybe False (const True) $ PSQ.lookup info (bktNodes bucket)
 
     map_ns f = bucket { bktNodes = f (bktNodes bucket) }
-    map_q f = bucket { bktQ = runIdentity $ f (bktQ bucket) }
+    -- map_q f = bucket { bktQ = runIdentity $ f (bktQ bucket) }
 
 insertBucket curTime (PingResult bad_node got_response) bucket
     = pure ([], Bucket (upd $ bktNodes bucket) popped)
  where
     (top, popped) = runIdentity $ popFront bucketQ (bktQ bucket)
     upd | got_response     = id
-        | Just info <- top = PSQ.insert info curTime . PSQ.delete bad_node
+        | Just info <- top = \nodes ->
+            fromMaybe nodes $ do
+                 _ <- PSQ.lookup bad_node nodes -- Insert only if there's a removal.
+                 let nodes' = PSQ.delete bad_node nodes
+                 pure $ PSQ.insert info curTime nodes'
         | otherwise        = id
 
+updateStamps :: Eq ip => Timestamp -> [NodeInfo ip] -> PSQ (NodeInfo ip) Timestamp -> PSQ (NodeInfo ip) Timestamp
+updateStamps curTime stales nodes = foldl' (\q n -> PSQ.insert n curTime q) nodes stales
+
+
 type BitIx = Word
 
 partitionQ :: Monad f => QueueMethods f elem b -> (elem -> Bool) -> b -> f (b, b)