1 files changed, 127 insertions, 100 deletions

diff --git a/src/Network/Kademlia/Bootstrap.hs b/src/Network/Kademlia/Bootstrap.hs
index 283e054a..42bff665 100644
--- a/src/Network/Kademlia/Bootstrap.hs
+++ b/src/Network/Kademlia/Bootstrap.hs
@@ -1,11 +1,14 @@
 {-# LANGUAGE CPP                   #-}
+{-# LANGUAGE ConstraintKinds       #-}
 {-# LANGUAGE DeriveFunctor         #-}
 {-# LANGUAGE DeriveTraversable     #-}
 {-# LANGUAGE FlexibleContexts      #-}
 {-# LANGUAGE GADTs                 #-}
 {-# LANGUAGE KindSignatures        #-}
+{-# LANGUAGE NamedFieldPuns        #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE PatternSynonyms       #-}
+{-# LANGUAGE RankNTypes            #-}
 {-# LANGUAGE ScopedTypeVariables   #-}
 module Network.Kademlia.Bootstrap where
 
@@ -42,34 +45,90 @@ import Network.Kademlia.Search
 import Control.Concurrent.Tasks
 import Network.Kademlia
 
-
--- | > pollForRefresh interval queue 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:
 --
--- 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.
+-- >>> not_handled =
+-- >>>   , Stranger :--> Applicant -- A ping is pending, it's result is covered:
+-- >>>                             --     (Applicant :--> Stranger)
+-- >>>                             --     (Applicant :--> Accepted)
+-- >>>   , Accepted :--> Applicant -- Never happens
+-- >>>   ]
 --
--- 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
+
+type SensibleNodeId nid ni =
+                 ( Show nid
+                 , Ord nid
+                 , Ord ni
+                 , Hashable nid
+                 , Hashable ni )
+
+data BucketRefresher nid ni = forall tok addr. Ord addr =>  BucketRefresher
+    { -- | A staleness threshold (if a bucket goes this long without being
+      -- touched, a refresh will be triggered).
+      refreshInterval :: POSIXTime
+      -- | A TVar with the time-to-refresh schedule for each bucket.
+      --
+      -- 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.  (See 'touchBucket').
+    , refreshQueue    :: TVar (Int.PSQ POSIXTime)
+      -- | This is the kademlia node search specification.
+    , refreshSearch   :: Search nid addr tok ni ni
+      -- | The current kademlia routing table buckets.
+    , refreshBuckets  :: TVar (R.BucketList ni)
+      -- | Action to ping a node.  This is used only during initial bootstrap
+      -- to get some nodes in our table.  A 'True' result is interpreted as a a
+      -- pong, where 'False' is a non-response.
+    , refreshPing     :: ni -> IO Bool
+    }
+
+-- | Fork a refresh loop.  Kill the returned thread to terminate it.
+forkPollForRefresh :: SensibleNodeId nid ni => BucketRefresher nid ni -> IO ThreadId
+forkPollForRefresh BucketRefresher{ refreshInterval
+                                  , refreshQueue
+                                  , refreshBuckets
+                                  , refreshSearch   } = fork $ do
       myThreadId >>= flip labelThread "pollForRefresh"
       fix $ \again -> do
         join $ atomically $ do
-            nextup <- Int.findMin <$> readTVar psq
+            nextup <- Int.findMin <$> readTVar refreshQueue
             maybe retry (return . go again) nextup
  where
+    refresh :: Int -> IO Int
+    refresh = refreshBucket refreshSearch refreshBuckets
+
     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)
+                             atomically $ modifyTVar' refreshQueue
+                                        $ Int.insert bktnum (now + refreshInterval)
                              -- Now fork the refresh operation.
                              -- TODO: We should probably propogate the kill signal to this thread.
                              fork $ do myThreadId >>= flip labelThread ("refresh."++show bktnum)
@@ -79,15 +138,37 @@ forkPollForRefresh interval psq refresh = do
             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
+
+-- | This is a helper to 'refreshBucket' which does some book keeping to decide
+-- whether or not a bucket is sufficiently refreshed or not.  It will return
+-- false when we can terminate a node search.
+checkBucketFull :: Ord ni => KademliaSpace nid ni   -- ^ Obtain a node id from a node.
+                            -> TVar (BucketList ni) -- ^ The current routing table.
+                            -> TVar (Set.Set ni)    -- ^ In-range nodes found so far.
+                            -> TVar Bool            -- ^ The result will also be written here.
+                            -> Int                  -- ^ The bucket number of interest.
+                            -> ni                   -- ^ A newly found node.
+                            -> STM Bool
+checkBucketFull space var resultCounter fin n found_node = do
+    let fullcount = R.defaultBucketSize
+        saveit True = writeTVar fin True >> return True
+        saveit _    = return False
+    tbl <- readTVar var
+    let counts = R.shape tbl
+        nid = kademliaLocation space found_node
+    -- Update the result set with every found node that is in the
+    -- bucket of interest.
+    when (n == R.bucketNumber space nid tbl)
+         $ modifyTVar' resultCounter (Set.insert found_node)
+    resultCount <- readTVar resultCounter
+    saveit $ case drop (n - 1) counts of
+        (cnt:_) | cnt < fullcount            -> True  -- bucket not full, keep going
+        _ | Set.size resultCount < fullcount -> True  -- we haven't got many results, keep going
+        _                                    -> False -- okay, good enough, let's quit.
+
+
+refreshBucket :: (Hashable a, Hashable t, Ord t, Ord addr, Ord a, Show t) =>
+                 Search t addr tok a a -> TVar (BucketList a) -> Int -> IO Int
 refreshBucket sch var n = do
     tbl <- atomically (readTVar var)
     let count = bktCount tbl
@@ -100,27 +181,13 @@ refreshBucket sch var n = do
                         (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
+                                    else checkBucketFull (searchSpace sch) var resultCounter fin n
     _ <- timeout (15*60*1000000) $ do
             atomically $ searchIsFinished s >>= check
     atomically $ searchCancel s
@@ -131,20 +198,11 @@ refreshBucket sch var n = do
         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
+bootstrap :: (Ord ni, Ord nid, Hashable nid, Hashable ni, Foldable t, Foldable t1, Show nid) =>
+             BucketRefresher nid ni -> t1 ni -> t ni -> IO ()
+bootstrap BucketRefresher { refreshSearch  = sch
+                          , refreshBuckets = var
+                          , refreshPing    = ping } ns ns0 = do
     gotPing <- atomically $ newTVar False
 
     -- First, ping the given nodes so that they are added to
@@ -193,6 +251,9 @@ bootstrap sch var ping ns ns0 = do
                 cnt <- refreshBucket sch var num
                 again cnt
 
+-- | Reschedule a bucket's refresh-time.  It should be called whenever a bucket
+-- changes.  This will typically be invoked from 'tblTransition'.
+--
 -- 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
@@ -201,55 +262,21 @@ bootstrap sch var ping ns ns0 = do
 -- 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
+touchBucket :: BucketRefresher nid ni
+            -> RoutingTransition ni     -- ^ What happened to the bucket?
             -> STM (IO ())
-touchBucket space interval bkts psq tr
+touchBucket BucketRefresher{ refreshSearch
+                           , refreshInterval
+                           , refreshBuckets
+                           , refreshQueue
+                           }
+            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
--- >>>   ]
---
+            let space = searchSpace refreshSearch
+                nid = kademliaLocation space (transitioningNode tr)
+            num <- R.bucketNumber space nid <$> readTVar refreshBuckets
+            modifyTVar' refreshQueue $ Int.insert num (now + refreshInterval)