1 files changed, 33 insertions, 21 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 7f20ad6d..99d8cdaf 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -73,6 +73,7 @@ import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
+import Control.Concurrent.STM
 
 {-----------------------------------------------------------------------
 --  Handlers
@@ -104,7 +105,10 @@ findNodeH = nodeHandler $ \ _ (FindNode nid) -> do
 -- | Default 'GetPeers' handler.
 getPeersH :: Address ip => NodeHandler ip
 getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
-  GotPeers <$> getPeerList ih <*> grantToken naddr
+  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.
 announceH :: Address ip => NodeHandler ip
@@ -236,20 +240,33 @@ refreshNodes nid = do
 insertNode :: Address ip => NodeInfo ip -> DHT ip ThreadId
 insertNode info = fork $ do
   var <- asks routingTable
-  t <- takeMVar var
-  t' <- do -- modifyMVar_ var $ \ t -> do
-    result <- routing (R.insert info t)
-    case result of
-      Nothing -> do
-        $(logDebugS) "insertNode" $ "Routing table is full: "
-                   <> T.pack (show (pPrint t))
-        return t
-      Just t' -> do
-        let logMsg = "Routing table updated: "
-                  <> pPrint t <> " -> " <> pPrint t'
-        $(logDebugS) "insertNode" (T.pack (render logMsg))
-        return t'
-  putMVar var t'
+  tm <- getTimestamp
+  let showTable = do
+          t <- liftIO $ atomically $ readTVar var
+          let logMsg = "Routing table: " <> pPrint t
+          $(logDebugS) "insertNode" (T.pack (render logMsg))
+  t <- liftIO $ atomically $ readTVar var
+  let arrival  = TryInsert info
+      arrival4 = TryInsert (fmap fromAddr info) :: Event (Maybe IPv4)
+  $(logDebugS) "insertNode" $ T.pack (show arrival4)
+  ps <- liftIO $ atomically $ do
+      t <- readTVar var
+      (ps,t') <- R.insert tm arrival t
+      writeTVar var t'
+      return ps
+  showTable
+  fork $ forM_ ps $ \(CheckPing ns)-> do
+    forM_ ns $ \n -> do
+        alive <- PingResult n <$> probeNode (nodeAddr n)
+        let PingResult _ b = alive
+        $(logDebugS) "insertNode" $ T.pack ("PingResult "++show (nodeId n,b))
+        tm <- getTimestamp
+        liftIO $ atomically $ do
+            t <- readTVar var
+            (_,t') <- R.insert tm alive t
+            writeTVar var t'
+        showTable
+  return ()
 
 -- | Throws exception if node is not responding.
 queryNode :: forall a b ip. Address ip => KRPC (Query a) (Response b)
@@ -268,9 +285,4 @@ q <@> addr = snd <$> queryNode addr q
 {-# INLINE (<@>) #-}
 
 restoreTable :: Address ip => Table ip -> DHT ip ()
-restoreTable tbl = do
-  tblvar <- asks routingTable
-  tbl0 <- liftIO $ takeMVar tblvar
-  mb <- routing $ merge tbl tbl0
-  maybe (return ()) (liftIO . putMVar tblvar) mb
-
+restoreTable tbl = mapM_ (insertNode . fst) $ L.concat $ R.toList tbl