1 files changed, 61 insertions, 8 deletions

diff --git a/src/Network/Kademlia/Bootstrap.hs b/src/Network/Kademlia/Bootstrap.hs
index f160000d..8ca0bb83 100644
--- a/src/Network/Kademlia/Bootstrap.hs
+++ b/src/Network/Kademlia/Bootstrap.hs
@@ -189,6 +189,48 @@ checkBucketFull space var resultCounter fin n found_node = do
         _ | Set.size resultCount < fullcount -> True  -- we haven't got many results, keep going
         _                                    -> False -- okay, good enough, let's quit.
 
+-- | Called from 'refreshBucket' with the current time when a refresh of the
+-- supplied bucket number finishes.
+onFinishedRefresh :: BucketRefresher nid ni -> Int -> POSIXTime -> STM (IO ())
+onFinishedRefresh BucketRefresher { bootstrapCountdown
+                                  , bootstrapMode
+                                  , refreshQueue
+                                  , refreshBuckets     } num now = do
+    bootstrapping <- readTVar bootstrapMode
+    if not bootstrapping then 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"
+        else do
+            -- The last bucket finished.
+            cnt <- readTVar bootstrapCountdown
+            case cnt of
+                Nothing -> do
+                    let fullsize = R.defaultBucketSize
+                        notfull (n,len) | n==num         = False
+                                        | len>=fullsize  = False
+                                        | otherwise      = True
+                        unfull = case filter notfull $ zip [0..] (R.shape tbl) of
+                                    [] -> [(0,0)] -- Schedule at least 1 more refresh.
+                                    xs -> xs
+                    forM_ unfull $ \(n,_) -> do
+                        -- 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
+                Just n  -> do writeTVar bootstrapCountdown $! Just $! pred n
+                              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)
 
 refreshBucket :: (Show nid, Ord ni, Ord nid, Hashable nid, Hashable ni) =>
                  BucketRefresher nid ni -> Int -> IO Int
@@ -217,6 +259,8 @@ refreshBucket r@BucketRefresher{ refreshSearch  = sch
             atomically $ searchIsFinished s >>= check
     atomically $ searchCancel s
     hPutStrLn stderr $ "Finish refresh " ++ show (n,sample)
+    now <- getPOSIXTime
+    join $ atomically $ onFinishedRefresh r n now
     rcount <- atomically $ do
         c <- Set.size <$> readTVar resultCounter
         b <- readTVar fin
@@ -336,21 +380,30 @@ effectiveRefreshInterval BucketRefresher{ refreshInterval
 --
 -- We embed the result in the STM monad but currently, no STM state changes
 -- occur until the returned IO action is invoked.  TODO: simplify?
-touchBucket :: BucketRefresher nid ni
+touchBucket :: SensibleNodeId nid ni
+            => BucketRefresher nid ni
             -> RoutingTransition ni     -- ^ What happened to the bucket?
             -> STM (IO ())
-touchBucket BucketRefresher{ refreshSearch
-                           , refreshInterval
-                           , refreshBuckets
-                           , refreshQueue    }
+touchBucket r@BucketRefresher{ refreshSearch
+                             , refreshInterval
+                             , refreshBuckets
+                             , refreshQueue
+                             , refreshLastTouch
+                             , bootstrapMode
+                             , bootstrapCountdown }
             RoutingTransition{ transitionedTo
                              , transitioningNode }
     = case transitionedTo of
         Applicant -> return $ return () -- Ignore transition to applicant.
         _         -> return $ do        -- Reschedule for any other transition.
          now <- getPOSIXTime
-         atomically $ do
+         join $ atomically $ do
             let space = searchSpace refreshSearch
                 nid = kademliaLocation space transitioningNode
-            num <- R.bucketNumber space nid <$> readTVar refreshBuckets
-            modifyTVar' refreshQueue $ Int.insert num (now + refreshInterval)
+            tbl <- readTVar refreshBuckets
+            let num = R.bucketNumber space nid tbl
+            let action = return ()
+            interval <- effectiveRefreshInterval r num
+            modifyTVar' refreshQueue $ Int.insert num (now + interval)
+            writeTVar refreshLastTouch now
+            return action