3 files changed, 41 insertions, 13 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
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 0bec867d..73b3d492 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -109,7 +109,6 @@ getPeersH :: Ord ip => Address ip => NodeHandler ip
 getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
   ps <- getPeerList ih
   tok <- grantToken naddr
-  $(logDebugS) "getPeersH" $ "INFO-HASH " <> T.pack (show (ih,fmap fromAddr naddr :: NodeAddr (Maybe IP)))
   return $ GotPeers ps tok
 
 -- | Default 'Announce' handler.
@@ -195,12 +194,6 @@ publish ih p = do
   _ <- sourceList [nodes] $= search ih (announceQ ih p) $$ C.take r
   return ()
 
-getTimestamp :: DHT ip Timestamp
-getTimestamp = do
-  utcTime <- liftIO $ getCurrentTime
-  $(logDebugS) "routing.make_timestamp" (T.pack (render (pPrint utcTime)))
-  return $ utcTimeToPOSIXSeconds utcTime
-
 
 probeNode :: Address ip => NodeAddr ip -> DHT ip (Bool, Maybe ReflectedIP)
 probeNode addr = do
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 339b18eb..4f861a1e 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -33,6 +33,7 @@ module Network.BitTorrent.DHT.Session
        , myNodeIdAccordingTo
        , routingInfo
        , routableAddress
+       , getTimestamp
 
          -- ** Initialization
        , LogFun
@@ -86,6 +87,10 @@ import Data.Set as S
 import Data.Time
 import Network (PortNumber)
 import System.Random (randomIO)
+import Data.Time.Clock.POSIX
+import Data.Text as Text
+import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
+
 
 import Data.Torrent as Torrent
 import Network.KRPC as KRPC hiding (Options, def)
@@ -435,7 +440,18 @@ lookupPeers :: Ord ip => InfoHash -> DHT ip [PeerAddr ip]
 lookupPeers ih = do
   refreshContacts
   var <- asks contactInfo
-  liftIO $ P.lookup ih <$> readTVarIO var
+  tm <- getTimestamp
+  liftIO $ atomically $ do
+    (ps,store') <- P.freshPeers ih tm <$> readTVar var
+    writeTVar var store'
+    return ps
+
+getTimestamp :: DHT ip Timestamp
+getTimestamp = do
+  utcTime <- liftIO $ getCurrentTime
+  $(logDebugS) "routing.make_timestamp" (Text.pack (render (pPrint utcTime)))
+  return $ utcTimeToPOSIXSeconds utcTime
+
 
 -- | Prepare result for 'get_peers' query.
 --