Bootstrap fixes.

1 files changed, 13 insertions, 5 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index fdc1fdd2..4920467a 100644
--- a/Kademlia.hs
+++ b/Kademlia.hs
@@ -270,10 +270,11 @@ refreshBucket :: forall nid ni addr.
                  Search nid addr ni ni -> TVar (BucketList ni) -> nid -> Int -> IO ()
 refreshBucket sch var nid n = do
     tbl <- atomically (readTVar var)
-    sample <- if n+1 == R.defaultBucketCount -- Is this the last bucket?
-                then return nid              -- Yes? Search our own id.
-                else genBucketSample nid     -- No? Generate a random id.
-                        (bucketRange n (n + 1 < bktCount tbl))
+    let count = bktCount tbl
+    sample <- if n+1 >= count            -- Is this the last bucket?
+                then return nid          -- Yes? Search our own id.
+                else genBucketSample nid -- No? Generate a random id.
+                        (bucketRange n (n + 1 < count))
     resultCounter <- atomically $ newTVar Set.empty
     let fullcount = R.defaultBucketSize
     let checkBucketFull :: ni -> STM Bool
@@ -315,7 +316,7 @@ bootstrap ::
 bootstrap sch var ping ns ns0 = do
     -- First, ping the given nodes so that they are added to
     -- our routing table.
-    withTaskGroup 20 $ \g -> do
+    withTaskGroup 50 $ \g -> do
         got_response <- or <$> mapConcurrently g ping ns
         -- We resort to the hardcoded fallback nodes only when we got no
         -- responses.  This is to lesson the burden on well-known boostrap
@@ -323,13 +324,20 @@ bootstrap sch var ping ns ns0 = do
         when (not got_response) $ do
             _ <- mapConcurrently g ping ns0
             return ()
+
+    mvar <- newMVar ()
     -- 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.
     fix $ \again -> do
+        takeMVar mvar
         tbl <- atomically $ readTVar var
+        fork $ do
+            -- Force a delay in case the search returns too quickly
+            threadDelay (60 * 1000000)
+            putMVar mvar ()
         let shp    = reverse $ zip (R.shape tbl) [0 .. ]
             unfull = filter ( (< R.defaultBucketSize) . fst ) shp
         case dropWhile ((> R.defaultBucketCount - 1) . snd) unfull of