Fixes to IPv4 bootstrap.

1 files changed, 33 insertions, 12 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index 2c2ff2b9..fdc1fdd2 100644
--- a/Kademlia.hs
+++ b/Kademlia.hs
@@ -8,6 +8,7 @@ module Kademlia 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.DHT.Routing   as R
@@ -17,6 +18,7 @@ import Control.Concurrent.Lifted.Instrument
 import Control.Concurrent.Lifted
 import GHC.Conc                  (labelThread)
 #endif
+import Control.Concurrent.Async.Pool
 import Control.Concurrent.STM
 import Control.Monad
 import Data.Bits
@@ -33,7 +35,7 @@ import Network.DatagramServer.Types   (genBucketSample)
 import System.Timeout
 import Text.PrettyPrint               as PP hiding (($$), (<>))
 import Text.PrettyPrint.HughesPJClass hiding (($$), (<>))
-import Control.Concurrent.Async.Pool
+import System.IO
 
 -- | The status of a given node with respect to a given routint table.
 data RoutingStatus
@@ -264,7 +266,7 @@ forkPollForRefresh interval psq refresh = do
         again
 
 refreshBucket :: forall nid ni addr.
-                 ( FiniteBits nid, Serialize nid, Ord nid, Ord ni, Hashable nid, Hashable ni, Ord addr ) =>
+                 ( Show nid, FiniteBits nid, Serialize nid, Ord nid, Ord ni, Hashable nid, Hashable ni, Ord addr ) =>
                  Search nid addr ni ni -> TVar (BucketList ni) -> nid -> Int -> IO ()
 refreshBucket sch var nid n = do
     tbl <- atomically (readTVar var)
@@ -272,25 +274,35 @@ refreshBucket sch var nid n = do
                 then return nid              -- Yes? Search our own id.
                 else genBucketSample nid     -- No? Generate a random id.
                         (bucketRange n (n + 1 < bktCount tbl))
-    resultCounter <- atomically $ newTVar 0
+    resultCounter <- atomically $ newTVar Set.empty
     let fullcount = R.defaultBucketSize
     let 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 (+ 1)
+                 $ modifyTVar resultCounter (Set.insert found_node)
             resultCount <- readTVar resultCounter
             case drop (n - 1) counts of
-                (cnt:_) | cnt < fullcount       -> return True
-                _ | resultCount + 3 < fullcount -> return True -- +3 because maybe duplicates.
-                _                               -> return False
+                (cnt:_) | cnt < fullcount            -> return True
+                _ | Set.size resultCount < fullcount -> return True
+                _                                    -> return False
+
+    hPutStrLn stderr $ "Start refresh " ++ show (n,sample)
+
     -- Set 15 minute timeout in order to avoid overlapping refreshes.
-    _ <- timeout (15*60*1000000) $ search sch tbl sample checkBucketFull
+    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)
     return ()
 
 bootstrap ::
-    ( Serialize nid
+    ( Show nid
+    , Serialize nid
     , FiniteBits nid
     , Hashable ni
     , Hashable nid
@@ -319,13 +331,22 @@ bootstrap sch var ping ns ns0 = do
     fix $ \again -> do
         tbl <- atomically $ readTVar var
         let shp    = reverse $ zip (R.shape tbl) [0 .. ]
-            unfull = dropWhile ( (< R.defaultBucketSize) . fst ) shp
-        case take 1 unfull of
-            [] -> return () -- Bootstrap complete!
+            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)
+                    refreshBucket sch var
+                        (kademliaLocation (searchSpace sch) (thisNode tbl))
+                        (R.defaultBucketCount - 1)
+                    again
             (size,num):_ -> do
                 -- If we don't yet have enough buckets, we need to search our own id.
                 -- We indicate that by setting the bucket number to the target.
                 let num' | bktCount tbl < R.defaultBucketCount = R.defaultBucketCount - 1
                          | otherwise                           = num
+                hPutStrLn stderr $ "Bucket too small, refresh "++ show (num',(size,num),shp)
                 refreshBucket sch var (kademliaLocation (searchSpace sch) (thisNode tbl)) num'
                 again