1 files changed, 24 insertions, 5 deletions

diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index 979dbb62..823982d4 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -1,6 +1,6 @@
 module Network.BitTorrent.DHT.ContactInfo
        ( PeerStore
-       , Network.BitTorrent.DHT.ContactInfo.lookup
+       , Network.BitTorrent.DHT.ContactInfo.freshPeers
        , Network.BitTorrent.DHT.ContactInfo.insertPeer
        , knownSwarms
        ) where
@@ -124,8 +124,8 @@ data SwarmData ip = SwarmData
 knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
 knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
 
-swarmSinglton :: Ord ip => PeerAddr ip -> SwarmData ip
-swarmSinglton a = SwarmData
+swarmSingleton :: Ord ip => PeerAddr ip -> SwarmData ip
+swarmSingleton a = SwarmData
     { peers = PSQ.singleton a 0
     , name = Nothing }
 
@@ -157,8 +157,27 @@ instance Serialize (PeerStore a) where
   put = undefined
 
 -- | Used in 'get_peers' DHT queries.
-lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
-lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
+-- lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
+-- lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
+
+batchSize = 64
+
+-- | Used in 'get_peers' DHT queries.
+freshPeers :: Ord a => InfoHash -> Timestamp -> PeerStore a -> ([PeerAddr a], PeerStore a)
+freshPeers ih tm (PeerStore m) = (ps, PeerStore m')
+ where
+    swarm = fromMaybe (SwarmData PSQ.empty Nothing) $ HM.lookup ih m
+    ps0 = take batchSize $ unfoldr (incomp minView) (peers swarm)
+    peers' = case reverse ps0 of
+                (_,psq):_ -> psq
+                _         -> peers swarm
+    ps = L.map (key . fst) ps0
+    m' = HM.insert ih swarm { peers = L.foldl' (\q p -> PSQ.insert p tm q) peers' ps } m
+
+incomp :: (x -> Maybe (r,x)) -> x -> Maybe ((r,x),x)
+incomp f x = do
+   (result,x') <- f x
+   pure $ ( (result,x'), x' )
 
 -- | Used in 'announce_peer' DHT queries.
 insertPeer :: Ord a => InfoHash -> Maybe ByteString -> PeerAddr a -> PeerStore a -> PeerStore a