Routing table snapshot/restore.

2 files changed, 16 insertions, 0 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index c5fcccb4..7f20ad6d 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -41,6 +41,7 @@ module Network.BitTorrent.DHT.Query
        , publish
 
        -- ** Routing table
+       , restoreTable
        , insertNode
        , refreshNodes
 
@@ -266,3 +267,10 @@ queryNode addr q = do
 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
+
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index fee52380..68edef56 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -46,6 +46,7 @@ module Network.BitTorrent.DHT.Routing
          -- * Construction
        , Network.BitTorrent.DHT.Routing.nullTable
        , Network.BitTorrent.DHT.Routing.insert
+       , Network.BitTorrent.DHT.Routing.merge
 
          -- * Conversion
        , Network.BitTorrent.DHT.Routing.TableEntry
@@ -457,3 +458,10 @@ toBucketList (One   b t) = b : toBucketList t
 
 toList :: Eq ip => Table ip -> [[TableEntry ip]]
 toList = L.map (L.map tableEntry . PSQ.toList) . toBucketList
+
+merge :: Eq ip => Table ip -> Table ip -> Routing ip (Table ip)
+merge a b = do
+    let ns = concatMap PSQ.toList $ toBucketList a
+    -- TODO: merge timestamps as well and let refresh take care of ping.
+    as <- filterM (needPing . nodeAddr . PSQ.key) ns
+    foldM (flip $ Network.BitTorrent.DHT.Routing.insert) b $ map PSQ.key as