1 files changed, 53 insertions, 5 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index ecd427b8..631761ac 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -173,8 +173,8 @@ bucketQ = seqQ
 --   very unlikely that all nodes in bucket fail within an hour of
 --   each other.
 --
-data Bucket ni = Bucket { bktNodes :: !(PSQ ni Timestamp)
-                        , bktQ     :: !(BucketQueue ni)
+data Bucket ni = Bucket { bktNodes :: !(PSQ ni Timestamp) -- current routing nodes
+                        , bktQ     :: !(BucketQueue ni)   -- replacements pending time-outs
                         } deriving Generic
 
 deriving instance Show ni => Show (Bucket ni)
@@ -230,7 +230,12 @@ delta = 15 * 60
 insertBucket :: (Alternative f, Ord ni) => -- (Eq ip, Alternative f, Ord (NodeId))  =>
     Timestamp -> Event ni -> Bucket ni -> f ([CheckPing ni], Bucket ni)
 insertBucket curTime (TryInsert info) bucket
-  -- just update timestamp if a node is already in bucket
+  -- Just update timestamp if a node is already in bucket.
+  --
+  -- Note PingResult events should only occur for nodes we requested a ping for,
+  -- and those will always already be in the routing queue and will get their
+  -- timestamp updated here, since 'TryInsert' is called on every inbound packet,
+  -- including ping results.
   | already_have
     = pure ( [], map_ns $ PSQ.insertWith max info curTime )
   -- bucket is good, but not full => we can insert a new node
@@ -271,11 +276,20 @@ 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
+    upd | got_response     = id -- Timestamp was already updated by TryInsert.
         | 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
+                 -- XXX: curTime is the time somebody pinged out, not the time
+                 -- of which the popped node was originally queued.  That's not
+                 -- quite right.
+                 --
+                 -- It's the best we can do atm since the overflow BucketQueue
+                 -- does not include timestamps.
+                 --
+                 -- This means that initial time-outs for replacement nodes are up to
+                 -- twice as long as for the originals.
                  pure $ PSQ.insert info curTime nodes'
         | otherwise        = id
 
@@ -614,13 +628,47 @@ deriving instance ( Show ip
 #endif
 
 -- | Atomic 'Table' update
+--
+-- Deprecated, use these two functions instead:
+--
+--  [ 'updateForInbound' ]
+--
+--  [ 'updateForPingResult' ]
+--
 insert :: -- ( Eq ip , Applicative m , Ord (NodeId) , FiniteBits (NodeId)) =>
     (Applicative m, Ord ni) =>
     (nid -> Word -> Bool)
     -> (ni -> nid)
     -> Timestamp -> Event ni -> Table ni nid -> m ([CheckPing ni], Table ni nid)
-insert testIdBit nodeId tm event tbl = pure $ fromMaybe ([],tbl) $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni) (eventId nodeId event) (insertBucket tm event) tbl
+insert testIdBit nodeId tm event tbl
+    = pure $ fromMaybe ([],tbl)
+           $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni)
+                                    (eventId nodeId event)
+                                    (insertBucket tm event)
+                                    tbl
+
+
+-- | Call on every inbound packet (including requested ping results).  Returns
+-- a list of nodes to ping and an updated routing table.  The caller should
+-- also invoke 'updateForPingResult' on each node in the to-ping list to update
+-- the routing table with that information as it becomes available.
+updateForInbound :: Ord ni =>
+    (nid -> Word -> Bool)
+    -> (ni -> nid)
+    -> Timestamp -> ni -> Table ni nid -> ([ni], Table ni nid)
+updateForInbound testIdBit nodeId tm ni tbl =
+    let Identity (ps, tbl') = Network.DHT.Routing.insert testIdBit nodeId tm (TryInsert ni) tbl
+        decon (CheckPing ns) = ns
+    in (concatMap decon ps, tbl')
 
+updateForPingResult :: Ord ni =>
+    (nid -> Word -> Bool)
+    -> (ni -> nid)
+    -> Timestamp -> ni -> Bool -> Table ni nid -> Table ni nid
+updateForPingResult testIdBit nodeId tm ni got_reply tbl =
+    -- The _ should always be empty in this case.
+    let Identity (_, tbl') = Network.DHT.Routing.insert testIdBit nodeId tm (PingResult ni got_reply) tbl
+    in tbl'
 
 
 {-----------------------------------------------------------------------