Factored out Network.Kademlia.Bootstrap.

4 files changed, 257 insertions, 210 deletions

diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index eb06564a..c7ef560c 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -54,6 +54,7 @@ import qualified Data.Wrapper.PSQInt          as Int
 import Debug.Trace
 import Network.BitTorrent.MainlineDHT.Symbols
 import Network.Kademlia
+import Network.Kademlia.Bootstrap
 import Network.Address                        (Address, fromAddr, fromSockAddr,
                                                setPort, sockAddrPort, testIdBit,
                                                toSockAddr, genBucketSample', WantIP(..),
diff --git a/src/Network/Kademlia.hs b/src/Network/Kademlia.hs
index f1ffda19..8956df2c 100644
--- a/src/Network/Kademlia.hs
+++ b/src/Network/Kademlia.hs
@@ -178,213 +178,3 @@ insertNode (Kademlia reporter space io) node = do
 
   return ()
 
-
--- TODO: Bootstrap/Refresh
---
--- From BEP 05:
---
--- Each bucket should maintain a "last changed" property to indicate how
--- "fresh" the contents are.
---
--- Note: We will use a "time to next refresh" property instead and store it in
--- a priority search queue.
---
--- When...
---
--- >>> data BucketTouchEvent = RoutingStatus :--> RoutingStatus
--- >>> bucketEvents =
--- >>>  [ Applicant :--> Stranger -- a node in a bucket is pinged and it responds,
--- >>>
--- >>>  , Stranger  :--> Accepted -- or a node is added to a bucket,
--- >>>
--- >>>  , Accepted  :--> Stranger -- or a node in a bucket is replaced
--- >>>  , Applicant :--> Accepted -- with another node,
--- >>>  ]
---
--- the bucket's last changed property should be updated. Buckets
--- that have not been changed in 15 minutes should be "refreshed." This is done
--- by picking a random ID in the range of the bucket and performing a
--- find_nodes search on it.
---
--- The only other possible BucketTouchEvents are as follows:
---
--- >>> not_handled =
--- >>>   , Stranger :--> Applicant -- A ping is pending, it's result is covered:
--- >>>                             --     (Applicant :--> Stranger)
--- >>>                             --     (Applicant :--> Accepted)
--- >>>   , Accepted :--> Applicant -- Never happens
--- >>>   ]
---
-
--- XXX: This will be redundantly triggered twice upon every node replacement
--- because we do not currently distinguish between standalone
--- insertion/deletion events and an insertion/deletion pair constituting
--- replacement.
---
--- It might also be better to pass the timestamp of the transition here and
--- keep the refresh queue in better sync with the routing table by updating it
--- within the STM monad.
-touchBucket :: KademliaSpace nid ni
-            -> POSIXTime
-            -> TVar (BucketList ni)
-            -> TVar (Int.PSQ POSIXTime)
-            -> RoutingTransition ni
-            -> STM (IO ())
-touchBucket space interval bkts psq tr
-    | (transitionedTo tr == Applicant)
-                = return $ return ()
-    | otherwise = return $ do
-         now <- getPOSIXTime
-         atomically $ do
-            let nid = kademliaLocation space (transitioningNode tr)
-            num <- R.bucketNumber space nid <$> readTVar bkts
-            modifyTVar' psq $ Int.insert num (now + interval)
-
--- | > pollForRefresh interval queue refresh
---
--- Fork a refresh loop.  Kill the returned thread to terminate it.  The
--- arguments are: a staleness threshold (if a bucket goes this long without
--- being touched, a refresh will be triggered), a TVar with the time-to-refresh
--- schedule for each bucket, and a refresh action to be forked when a bucket
--- excedes the staleness threshold.
---
--- TO "touch" a bucket and prevent it from being refreshed, reschedule it's
--- refresh time to some time into the future by modifying the 'Int.PSQ' in the
--- TVar.
-forkPollForRefresh :: POSIXTime -> TVar (Int.PSQ POSIXTime) -> (Int -> IO a) -> IO ThreadId
-forkPollForRefresh interval psq refresh = do
-    fork $ do
-      myThreadId >>= flip labelThread "pollForRefresh"
-      fix $ \again -> do
-        join $ atomically $ do
-            nextup <- Int.findMin <$> readTVar psq
-            maybe retry (return . go again) nextup
- where
-    go again ( bktnum :-> refresh_time ) = do
-        now <- getPOSIXTime
-        case fromEnum (refresh_time - now) of
-            x | x <= 0 -> do -- Refresh time!
-                             -- Move it to the back of the refresh queue.
-                             atomically $ modifyTVar' psq
-                                        $ Int.insert bktnum (now + interval)
-                             -- Now fork the refresh operation.
-                             -- TODO: We should probably propogate the kill signal to this thread.
-                             fork $ do myThreadId >>= flip labelThread ("refresh."++show bktnum)
-                                       _ <- refresh bktnum
-                                       return ()
-                             return ()
-            seconds    -> threadDelay ( seconds * 1000000 )
-        again
-
-refreshBucket :: forall nid tok ni addr.
-                 ( Show nid
-                 , Serialize nid
-                 , Ord nid
-                 , Ord ni
-                 , Ord addr
-                 , Hashable nid
-                 , Hashable ni ) =>
-                 Search nid addr tok ni ni -> TVar (BucketList ni) -> Int -> IO Int
-refreshBucket sch var n = do
-    tbl <- atomically (readTVar var)
-    let count = bktCount tbl
-        nid = kademliaLocation (searchSpace sch) (thisNode tbl)
-    sample <- if n+1 >= count                         -- Is this the last bucket?
-                then return nid                       -- Yes? Search our own id.
-                else kademliaSample (searchSpace sch) -- No? Generate a random id.
-                        getEntropy
-                        nid
-                        (bucketRange n (n + 1 < count))
-    fin <- atomically $ newTVar False
-    resultCounter <- atomically $ newTVar Set.empty
-    let fullcount = R.defaultBucketSize
-        saveit True = writeTVar fin True >> return True
-        saveit _    = return False
-        checkBucketFull :: ni -> STM Bool
-        checkBucketFull found_node = do
-            tbl <- readTVar var
-            let counts = R.shape tbl
-            when (n == R.bucketNumber (searchSpace sch) (kademliaLocation (searchSpace sch) found_node) tbl)
-                 $ modifyTVar' resultCounter (Set.insert found_node)
-            resultCount <- readTVar resultCounter
-            saveit $ case drop (n - 1) counts of
-                (cnt:_) | cnt < fullcount            -> True
-                _ | Set.size resultCount < fullcount -> True
-                _                                    -> False
-
-    hPutStrLn stderr $ "Start refresh " ++ show (n,sample)
-
-    -- Set 15 minute timeout in order to avoid overlapping refreshes.
-    s <- search sch tbl sample $ if n+1 == R.defaultBucketCount
-                                    then const $ return True -- Never short-circuit the last bucket.
-                                    else checkBucketFull
-    _ <- timeout (15*60*1000000) $ do
-            atomically $ searchIsFinished s >>= check
-    atomically $ searchCancel s
-    hPutStrLn stderr $ "Finish refresh " ++ show (n,sample)
-    rcount <- atomically $ do
-        c <- Set.size <$> readTVar resultCounter
-        b <- readTVar fin
-        return $ if b then 1 else c
-    return rcount
-
-bootstrap ::
-    ( Show nid
-    , Serialize nid
-    -- , FiniteBits nid
-    , Hashable ni
-    , Hashable nid
-    , Ord ni
-    , Ord addr
-    , Ord nid
-    , Traversable t1
-    , Traversable t
-    ) => Search nid addr tok ni ni -> TVar (BucketList ni) -> (ni -> IO Bool) -> t ni -> t1 ni -> IO ()
-bootstrap sch var ping ns ns0 = do
-    gotPing <- atomically $ newTVar False
-
-    -- First, ping the given nodes so that they are added to
-    -- our routing table.
-    withTaskGroup "bootstrap.resume" 20 $ \g -> do
-        forM_ ns $ \n -> do
-            let lbl = show $ kademliaLocation (searchSpace sch) n
-            forkTask g lbl $ do
-                b <- ping n
-                when b $ atomically $ writeTVar gotPing True
-
-    -- We resort to the hardcoded fallback nodes only when we got no
-    -- responses.  This is to lesson the burden on well-known boostrap
-    -- nodes.
-    fallback <- atomically (readTVar gotPing) >>= return . when . not
-    fallback $ withTaskGroup "bootstrap.ping" 20 $ \g -> do
-        forM_ ns0 $ \n -> do
-            forkTask g (show $ kademliaLocation (searchSpace sch) n)
-                       (void $ ping n)
-    hPutStrLn stderr "Finished bootstrap pings."
-
-    -- Now run searches until all the buckets are full.  On a small network,
-    -- this may never quit.
-    --
-    -- TODO: For small networks, we should give up on filling a nearby bucket
-    -- at some point and move on to one farther away.
-    flip fix 1 $ \again cnt -> do
-        when (cnt==0) $ do
-            -- Force a delay in case the search returns too quickly
-            hPutStrLn stderr $ "Zero results, forcing 1 minute delay"
-            threadDelay (60 * 1000000)
-        tbl <- atomically $ readTVar var
-        let shp    = zip (R.shape tbl) [0 .. ]
-            unfull = filter ( (< R.defaultBucketSize) . fst ) shp
-        case dropWhile ((> R.defaultBucketCount - 1) . snd) unfull of
-            [] -> do
-                when (length shp < R.defaultBucketCount) $ do
-                    -- Not enough buckets, keep trying.
-                    hPutStrLn stderr
-                        $ "Not enough buckets, refresh " ++ show (R.defaultBucketCount - 1)
-                    cnt <- refreshBucket sch var
-                        (R.defaultBucketCount - 1)
-                    again cnt
-            (size,num):_ -> do
-                hPutStrLn stderr $ "Bucket too small, refresh "++ show (size,num,shp)
-                cnt <- refreshBucket sch var num
-                again cnt
diff --git a/src/Network/Kademlia/Bootstrap.hs b/src/Network/Kademlia/Bootstrap.hs
new file mode 100644
index 00000000..283e054a
--- /dev/null
+++ b/src/Network/Kademlia/Bootstrap.hs
@@ -0,0 +1,255 @@
+{-# LANGUAGE CPP                   #-}
+{-# LANGUAGE DeriveFunctor         #-}
+{-# LANGUAGE DeriveTraversable     #-}
+{-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE GADTs                 #-}
+{-# LANGUAGE KindSignatures        #-}
+{-# LANGUAGE PartialTypeSignatures #-}
+{-# LANGUAGE PatternSynonyms       #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+module Network.Kademlia.Bootstrap where
+
+import Data.Function
+import Data.Maybe
+import qualified Data.Set as Set
+import Data.Time.Clock       (getCurrentTime)
+import Data.Time.Clock.POSIX (getPOSIXTime, utcTimeToPOSIXSeconds)
+import Network.Kademlia.Routing   as R
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import Control.Concurrent.Lifted
+import GHC.Conc                  (labelThread)
+#endif
+import Control.Concurrent.STM
+import Control.Monad
+import Data.Bits
+import Data.Hashable
+import Data.IP
+import Data.Monoid
+import Data.Serialize                 (Serialize)
+import Data.Time.Clock.POSIX          (POSIXTime)
+import System.Entropy
+import System.Timeout
+import Text.PrettyPrint               as PP hiding (($$), (<>))
+import Text.PrettyPrint.HughesPJClass hiding (($$), (<>))
+import System.IO
+
+import qualified Data.Wrapper.PSQInt  as Int
+         ;import Data.Wrapper.PSQInt  (pattern (:->))
+import Network.Address                (bucketRange,genBucketSample)
+import Network.Kademlia.Search
+import Control.Concurrent.Tasks
+import Network.Kademlia
+
+
+-- | > pollForRefresh interval queue refresh
+--
+-- Fork a refresh loop.  Kill the returned thread to terminate it.  The
+-- arguments are: a staleness threshold (if a bucket goes this long without
+-- being touched, a refresh will be triggered), a TVar with the time-to-refresh
+-- schedule for each bucket, and a refresh action to be forked when a bucket
+-- excedes the staleness threshold.
+--
+-- TO "touch" a bucket and prevent it from being refreshed, reschedule it's
+-- refresh time to some time into the future by modifying the 'Int.PSQ' in the
+-- TVar.
+forkPollForRefresh :: POSIXTime -> TVar (Int.PSQ POSIXTime) -> (Int -> IO a) -> IO ThreadId
+forkPollForRefresh interval psq refresh = do
+    fork $ do
+      myThreadId >>= flip labelThread "pollForRefresh"
+      fix $ \again -> do
+        join $ atomically $ do
+            nextup <- Int.findMin <$> readTVar psq
+            maybe retry (return . go again) nextup
+ where
+    go again ( bktnum :-> refresh_time ) = do
+        now <- getPOSIXTime
+        case fromEnum (refresh_time - now) of
+            x | x <= 0 -> do -- Refresh time!
+                             -- Move it to the back of the refresh queue.
+                             atomically $ modifyTVar' psq
+                                        $ Int.insert bktnum (now + interval)
+                             -- Now fork the refresh operation.
+                             -- TODO: We should probably propogate the kill signal to this thread.
+                             fork $ do myThreadId >>= flip labelThread ("refresh."++show bktnum)
+                                       _ <- refresh bktnum
+                                       return ()
+                             return ()
+            seconds    -> threadDelay ( seconds * 1000000 )
+        again
+
+refreshBucket :: forall nid tok ni addr.
+                 ( Show nid
+                 , Serialize nid
+                 , Ord nid
+                 , Ord ni
+                 , Ord addr
+                 , Hashable nid
+                 , Hashable ni ) =>
+                 Search nid addr tok ni ni -> TVar (BucketList ni) -> Int -> IO Int
+refreshBucket sch var n = do
+    tbl <- atomically (readTVar var)
+    let count = bktCount tbl
+        nid = kademliaLocation (searchSpace sch) (thisNode tbl)
+    sample <- if n+1 >= count                         -- Is this the last bucket?
+                then return nid                       -- Yes? Search our own id.
+                else kademliaSample (searchSpace sch) -- No? Generate a random id.
+                        getEntropy
+                        nid
+                        (bucketRange n (n + 1 < count))
+    fin <- atomically $ newTVar False
+    resultCounter <- atomically $ newTVar Set.empty
+    let fullcount = R.defaultBucketSize
+        saveit True = writeTVar fin True >> return True
+        saveit _    = return False
+        checkBucketFull :: ni -> STM Bool
+        checkBucketFull found_node = do
+            tbl <- readTVar var
+            let counts = R.shape tbl
+            when (n == R.bucketNumber (searchSpace sch) (kademliaLocation (searchSpace sch) found_node) tbl)
+                 $ modifyTVar' resultCounter (Set.insert found_node)
+            resultCount <- readTVar resultCounter
+            saveit $ case drop (n - 1) counts of
+                (cnt:_) | cnt < fullcount            -> True
+                _ | Set.size resultCount < fullcount -> True
+                _                                    -> False
+
+    hPutStrLn stderr $ "Start refresh " ++ show (n,sample)
+
+    -- Set 15 minute timeout in order to avoid overlapping refreshes.
+    s <- search sch tbl sample $ if n+1 == R.defaultBucketCount
+                                    then const $ return True -- Never short-circuit the last bucket.
+                                    else checkBucketFull
+    _ <- timeout (15*60*1000000) $ do
+            atomically $ searchIsFinished s >>= check
+    atomically $ searchCancel s
+    hPutStrLn stderr $ "Finish refresh " ++ show (n,sample)
+    rcount <- atomically $ do
+        c <- Set.size <$> readTVar resultCounter
+        b <- readTVar fin
+        return $ if b then 1 else c
+    return rcount
+
+
+bootstrap ::
+    ( Show nid
+    , Serialize nid
+    -- , FiniteBits nid
+    , Hashable ni
+    , Hashable nid
+    , Ord ni
+    , Ord addr
+    , Ord nid
+    , Traversable t1
+    , Traversable t
+    ) => Search nid addr tok ni ni -> TVar (BucketList ni) -> (ni -> IO Bool) -> t ni -> t1 ni -> IO ()
+bootstrap sch var ping ns ns0 = do
+    gotPing <- atomically $ newTVar False
+
+    -- First, ping the given nodes so that they are added to
+    -- our routing table.
+    withTaskGroup "bootstrap.resume" 20 $ \g -> do
+        forM_ ns $ \n -> do
+            let lbl = show $ kademliaLocation (searchSpace sch) n
+            forkTask g lbl $ do
+                b <- ping n
+                when b $ atomically $ writeTVar gotPing True
+
+    -- We resort to the hardcoded fallback nodes only when we got no
+    -- responses.  This is to lesson the burden on well-known boostrap
+    -- nodes.
+    fallback <- atomically (readTVar gotPing) >>= return . when . not
+    fallback $ withTaskGroup "bootstrap.ping" 20 $ \g -> do
+        forM_ ns0 $ \n -> do
+            forkTask g (show $ kademliaLocation (searchSpace sch) n)
+                       (void $ ping n)
+    hPutStrLn stderr "Finished bootstrap pings."
+
+    -- Now run searches until all the buckets are full.  On a small network,
+    -- this may never quit.
+    --
+    -- TODO: For small networks, we should give up on filling a nearby bucket
+    -- at some point and move on to one farther away.
+    flip fix 1 $ \again cnt -> do
+        when (cnt==0) $ do
+            -- Force a delay in case the search returns too quickly
+            hPutStrLn stderr $ "Zero results, forcing 1 minute delay"
+            threadDelay (60 * 1000000)
+        tbl <- atomically $ readTVar var
+        let shp    = zip (R.shape tbl) [0 .. ]
+            unfull = filter ( (< R.defaultBucketSize) . fst ) shp
+        case dropWhile ((> R.defaultBucketCount - 1) . snd) unfull of
+            [] -> do
+                when (length shp < R.defaultBucketCount) $ do
+                    -- Not enough buckets, keep trying.
+                    hPutStrLn stderr
+                        $ "Not enough buckets, refresh " ++ show (R.defaultBucketCount - 1)
+                    cnt <- refreshBucket sch var
+                        (R.defaultBucketCount - 1)
+                    again cnt
+            (size,num):_ -> do
+                hPutStrLn stderr $ "Bucket too small, refresh "++ show (size,num,shp)
+                cnt <- refreshBucket sch var num
+                again cnt
+
+-- XXX: This will be redundantly triggered twice upon every node replacement
+-- because we do not currently distinguish between standalone
+-- insertion/deletion events and an insertion/deletion pair constituting
+-- replacement.
+--
+-- It might also be better to pass the timestamp of the transition here and
+-- keep the refresh queue in better sync with the routing table by updating it
+-- within the STM monad.
+touchBucket :: KademliaSpace nid ni
+            -> POSIXTime
+            -> TVar (BucketList ni)
+            -> TVar (Int.PSQ POSIXTime)
+            -> RoutingTransition ni
+            -> STM (IO ())
+touchBucket space interval bkts psq tr
+    | (transitionedTo tr == Applicant)
+                = return $ return ()
+    | otherwise = return $ do
+         now <- getPOSIXTime
+         atomically $ do
+            let nid = kademliaLocation space (transitioningNode tr)
+            num <- R.bucketNumber space nid <$> readTVar bkts
+            modifyTVar' psq $ Int.insert num (now + interval)
+
+-- TODO: Bootstrap/Refresh
+--
+-- From BEP 05:
+--
+-- Each bucket should maintain a "last changed" property to indicate how
+-- "fresh" the contents are.
+--
+-- Note: We will use a "time to next refresh" property instead and store it in
+-- a priority search queue.
+--
+-- When...
+--
+-- >>> data BucketTouchEvent = RoutingStatus :--> RoutingStatus
+-- >>> bucketEvents =
+-- >>>  [ Applicant :--> Stranger -- a node in a bucket is pinged and it responds,
+-- >>>
+-- >>>  , Stranger  :--> Accepted -- or a node is added to a bucket,
+-- >>>
+-- >>>  , Accepted  :--> Stranger -- or a node in a bucket is replaced
+-- >>>  , Applicant :--> Accepted -- with another node,
+-- >>>  ]
+--
+-- the bucket's last changed property should be updated. Buckets
+-- that have not been changed in 15 minutes should be "refreshed." This is done
+-- by picking a random ID in the range of the bucket and performing a
+-- find_nodes search on it.
+--
+-- The only other possible BucketTouchEvents are as follows:
+--
+-- >>> not_handled =
+-- >>>   , Stranger :--> Applicant -- A ping is pending, it's result is covered:
+-- >>>                             --     (Applicant :--> Stranger)
+-- >>>                             --     (Applicant :--> Accepted)
+-- >>>   , Accepted :--> Applicant -- Never happens
+-- >>>   ]
+--
diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index fc28e2d2..494e319b 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -11,6 +11,7 @@ import Crypto.Tox
 import Network.Kademlia.Search
 import qualified Data.Wrapper.PSQInt as Int
 import Network.Kademlia
+import Network.Kademlia.Bootstrap    (touchBucket)
 import Network.Address               (WantIP (..), ipFamily, testIdBit,fromSockAddr, sockAddrPort)
 import qualified Network.Kademlia.Routing as R
 import Control.TriadCommittee