1 files changed, 45 insertions, 23 deletions

diff --git a/src/Network/Kademlia/Bootstrap.hs b/src/Network/Kademlia/Bootstrap.hs
index 3540b24e..595659d2 100644
--- a/src/Network/Kademlia/Bootstrap.hs
+++ b/src/Network/Kademlia/Bootstrap.hs
@@ -5,6 +5,7 @@
 {-# LANGUAGE FlexibleContexts      #-}
 {-# LANGUAGE GADTs                 #-}
 {-# LANGUAGE KindSignatures        #-}
+{-# LANGUAGE LambdaCase            #-}
 {-# LANGUAGE NamedFieldPuns        #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE PatternSynonyms       #-}
@@ -56,24 +57,32 @@ type SensibleNodeId nid 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
+      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)
+      -- its refresh time to some time into the future by modifying its
+      -- priority in this priority search queue.
+    , refreshQueue     :: TVar (Int.PSQ POSIXTime)
       -- | This is the kademlia node search specification.
-    , refreshSearch   :: Search nid addr tok ni ni
+    , refreshSearch    :: Search nid addr tok ni ni
       -- | The current kademlia routing table buckets.
-    , refreshBuckets  :: TVar (R.BucketList ni)
+    , 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
+    , refreshPing      :: ni -> IO Bool
+    , -- | Timestamp of last bucket event.
+      refreshLastTouch :: TVar POSIXTime
+    , -- | This variable indicates whether or not we are in bootstrapping mode.
+      bootstrapMode   :: TVar Bool
+    , -- | When this countdown reaches 0, we exit bootstrap mode.  It is decremented on
+      -- every finished refresh.
+      bootstrapCountdown :: TVar (Maybe Int)
     }
 
-newBucketRefresher :: (Ord addr, Hashable addr, SensibleNodeId nid ni)
+newBucketRefresher :: ( Ord addr, Hashable addr
+                      , SensibleNodeId nid ni )
                     => ni
                     -> Search nid addr tok ni ni
                     -> (ni -> IO Bool)
@@ -83,12 +92,18 @@ newBucketRefresher template_ni sch ping = do
         nodeId = kademliaLocation spc
     bkts <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) template_ni R.defaultBucketCount
     sched <- newTVar Int.empty
+    lasttouch <- newTVar 0 -- Would use getPOSIXTime here, or minBound, but alas...
+    bootstrapVar <- newTVar True -- Start in bootstrapping mode.
+    bootstrapCnt <- newTVar Nothing
     return  BucketRefresher
-                    { refreshInterval = 15 * 60
-                    , refreshQueue    = sched
-                    , refreshSearch   = sch
-                    , refreshBuckets  = bkts
-                    , refreshPing     = ping
+                    { refreshInterval    = 15 * 60
+                    , refreshQueue       = sched
+                    , refreshSearch      = sch
+                    , refreshBuckets     = bkts
+                    , refreshPing        = ping
+                    , refreshLastTouch   = lasttouch
+                    , bootstrapMode      = bootstrapVar
+                    , bootstrapCountdown = bootstrapCnt
                     }
 
 -- | This was added to avoid the compile error "Record update for
@@ -99,11 +114,14 @@ updateRefresherIO :: Ord addr
                      -> (ni -> IO Bool)
                      -> BucketRefresher nid ni ->  BucketRefresher nid ni
 updateRefresherIO sch ping BucketRefresher{..} = BucketRefresher
-    { refreshSearch   = sch
-    , refreshPing     = ping
-    , refreshInterval = refreshInterval
-    , refreshBuckets  = refreshBuckets
-    , refreshQueue    = refreshQueue
+    { refreshSearch      = sch
+    , refreshPing        = ping
+    , refreshInterval    = refreshInterval
+    , refreshBuckets     = refreshBuckets
+    , refreshQueue       = refreshQueue
+    , refreshLastTouch   = refreshLastTouch
+    , bootstrapMode      = bootstrapMode
+    , bootstrapCountdown = bootstrapCountdown
     }
 
 -- | Fork a refresh loop.  Kill the returned thread to terminate it.
@@ -197,14 +215,19 @@ refreshBucket sch var n = do
         return $ if b then 1 else c
     return rcount
 
+refreshLastBucket :: SensibleNodeId nid ni => BucketRefresher nid ni -> IO ()
+refreshLastBucket r@BucketRefresher { refreshBuckets, refreshSearch } = do
+    cnt <- atomically $ bktCount <$> readTVar refreshBuckets
+    void $ refreshBucket refreshSearch refreshBuckets (cnt - 1)
+
 bootstrap :: (Ord ni, Ord nid, Hashable nid, Hashable ni, Foldable t, Foldable t1, Show nid) =>
              BucketRefresher nid ni
              -> t1 ni -- ^ Nodes to bootstrap from.
              -> t ni  -- ^ Fallback nodes; used only if the others are unresponsive.
              -> IO ()
-bootstrap BucketRefresher { refreshSearch  = sch
-                          , refreshBuckets = var
-                          , refreshPing    = ping } ns ns0 = do
+bootstrap r@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
@@ -226,8 +249,7 @@ bootstrap BucketRefresher { refreshSearch  = sch
                        (void $ ping n)
     hPutStrLn stderr "Finished bootstrap pings."
     -- Now search our own Id by refreshing the last bucket.
-    last <- atomically $ bktCount <$> readTVar var
-    void $ refreshBucket sch var last
+    refreshLastBucket r
     -- That's it.
     --
     -- Hopefully 'forkPollForRefresh' was invoked and can take over