Finished reworked bootstrap algorithm.

1 files changed, 39 insertions, 16 deletions

diff --git a/src/Network/Kademlia/Bootstrap.hs b/src/Network/Kademlia/Bootstrap.hs
index 8ca0bb83..93cf08f3 100644
--- a/src/Network/Kademlia/Bootstrap.hs
+++ b/src/Network/Kademlia/Bootstrap.hs
@@ -138,7 +138,7 @@ forkPollForRefresh r@BucketRefresher{ refreshInterval
  where
     refresh :: Int -> IO Int
     refresh n = do
-        hPutStrLn stderr $ "Refresh time! "++ show n
+        -- hPutStrLn stderr $ "Refresh time! "++ show n
         refreshBucket r n
 
     go again ( bktnum :-> refresh_time ) = do
@@ -157,7 +157,7 @@ forkPollForRefresh r@BucketRefresher{ refreshInterval
                                        return ()
                              return ()
             picoseconds -> do
-                hPutStrLn stderr $ show (picoseconds `div` 10^12) ++ " seconds until refresh " ++ show bktnum
+                -- hPutStrLn stderr $ show (picoseconds `div` 10^12) ++ " seconds until refresh " ++ show bktnum
                 threadDelay ( picoseconds `div` 10^6 )
         again
 
@@ -197,13 +197,13 @@ onFinishedRefresh BucketRefresher { bootstrapCountdown
                                   , refreshQueue
                                   , refreshBuckets     } num now = do
     bootstrapping <- readTVar bootstrapMode
-    if not bootstrapping then return $ hPutStrLn stderr $ "Finished non-boostrapping refresh: "++show num
+    if not bootstrapping then return $ return () -- hPutStrLn stderr $ "Finished non-boostrapping refresh: "++show num
                          else do
     tbl <- readTVar refreshBuckets
     action <-
       if num /= R.bktCount tbl - 1
         then do modifyTVar' bootstrapCountdown (fmap pred)
-                return $ hPutStrLn stderr $ "BOOTSTRAP decrement"
+                return $ return () -- hPutStrLn stderr $ "BOOTSTRAP decrement"
         else do
             -- The last bucket finished.
             cnt <- readTVar bootstrapCountdown
@@ -220,17 +220,17 @@ onFinishedRefresh BucketRefresher { bootstrapCountdown
                         -- Schedule immediate refresh for unfull buckets (other than this one).
                         modifyTVar' refreshQueue $ Int.insert n (now - 1)
                     writeTVar bootstrapCountdown $! Just $! length unfull
-                    return $ hPutStrLn stderr $ "BOOTSTRAP scheduling: "++show unfull
+                    return $ return () -- hPutStrLn stderr $ "BOOTSTRAP scheduling: "++show unfull
                 Just n  -> do writeTVar bootstrapCountdown $! Just $! pred n
-                              return $ hPutStrLn stderr "BOOTSTRAP decrement (last bucket)"
+                              return $ return () -- hPutStrLn stderr "BOOTSTRAP decrement (last bucket)"
     cnt <- readTVar bootstrapCountdown
     if (cnt == Just 0)
     then do
         -- Boostrap finished!
         writeTVar bootstrapMode False
         writeTVar bootstrapCountdown Nothing
-        return $ do action ; hPutStrLn stderr $ "BOOTSTRAP complete (" ++ show num ++ ")."
-    else return $ do action ; hPutStrLn stderr $ "BOOTSTRAP progress " ++ show (num,R.bktCount tbl,cnt)
+        return $ do action ; hPutStrLn stderr $ "BOOTSTRAP complete (" ++ show (R.shape tbl) ++ ")."
+    else return $ do action ; hPutStrLn stderr $ "BOOTSTRAP progress " ++ show (num,R.shape tbl,cnt)
 
 refreshBucket :: (Show nid, Ord ni, Ord nid, Hashable nid, Hashable ni) =>
                  BucketRefresher nid ni -> Int -> IO Int
@@ -268,9 +268,25 @@ refreshBucket r@BucketRefresher{ refreshSearch  = sch
     return rcount
 
 refreshLastBucket :: SensibleNodeId nid ni => BucketRefresher nid ni -> IO ()
-refreshLastBucket r@BucketRefresher { refreshBuckets } = do
-    cnt <- atomically $ bktCount <$> readTVar refreshBuckets
-    void $ refreshBucket r (cnt - 1)
+refreshLastBucket r@BucketRefresher { refreshBuckets
+                                    , refreshQueue   } = do
+
+    now <- getPOSIXTime
+    atomically $ do
+    cnt <- bktCount <$> readTVar refreshBuckets
+    -- Schedule immediate refresh.
+    modifyTVar' refreshQueue $ Int.insert (cnt-1) (now - 1)
+
+restartBootstrap :: (Hashable ni, Hashable nid, Ord ni, Ord nid, Show nid) =>
+                 BucketRefresher nid ni -> STM (IO ())
+restartBootstrap r@BucketRefresher{ bootstrapMode, bootstrapCountdown } = do
+    unchanged <- readTVar bootstrapMode
+    writeTVar bootstrapMode True
+    writeTVar bootstrapCountdown Nothing
+    if not unchanged then return $ do
+                            hPutStrLn stderr "BOOTSTRAP entered bootstrap mode"
+                            refreshLastBucket r
+                     else return $ hPutStrLn stderr "BOOTSTRAP already bootstrapping"
 
 bootstrap :: (Ord ni, Ord nid, Hashable nid, Hashable ni, Foldable t, Foldable t1, Show nid) =>
              BucketRefresher nid ni
@@ -279,7 +295,8 @@ bootstrap :: (Ord ni, Ord nid, Hashable nid, Hashable ni, Foldable t, Foldable t
              -> IO ()
 bootstrap r@BucketRefresher { refreshSearch  = sch
                             , refreshBuckets = var
-                            , refreshPing    = ping } ns ns0 = do
+                            , refreshPing    = ping
+                            , bootstrapMode         } ns ns0 = do
     gotPing <- atomically $ newTVar False
 
     -- First, ping the given nodes so that they are added to
@@ -300,9 +317,10 @@ bootstrap r@BucketRefresher { refreshSearch  = sch
             forkTask g (show $ kademliaLocation (searchSpace sch) n)
                        (void $ ping n)
     hPutStrLn stderr "Finished bootstrap pings."
-    -- Now search our own Id by refreshing the last bucket.
-    refreshLastBucket r
-    -- That's it.
+    -- Now search our own Id by entering bootstrap mode from non-bootstrap mode.
+    join $ atomically $ do
+        writeTVar bootstrapMode False
+        restartBootstrap r
     --
     -- Hopefully 'forkPollForRefresh' was invoked and can take over
     -- maintenance.
@@ -402,7 +420,12 @@ touchBucket r@BucketRefresher{ refreshSearch
                 nid = kademliaLocation space transitioningNode
             tbl <- readTVar refreshBuckets
             let num = R.bucketNumber space nid tbl
-            let action = return ()
+            stamp <- readTVar refreshLastTouch
+            action <- case stamp /= 0 && (now - stamp > 60) of
+                True -> do
+                    -- It's been one minute since any bucket has been touched, re-enter bootstrap mode.
+                    restartBootstrap r
+                False -> return $ return ()
             interval <- effectiveRefreshInterval r num
             modifyTVar' refreshQueue $ Int.insert num (now + interval)
             writeTVar refreshLastTouch now