1 files changed, 47 insertions, 31 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index 4920467a..57afb7fc 100644
--- a/Kademlia.hs
+++ b/Kademlia.hs
@@ -18,7 +18,6 @@ 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
@@ -36,6 +35,7 @@ import System.Timeout
 import Text.PrettyPrint               as PP hiding (($$), (<>))
 import Text.PrettyPrint.HughesPJClass hiding (($$), (<>))
 import System.IO
+import Tasks
 
 -- | The status of a given node with respect to a given routint table.
 data RoutingStatus
@@ -241,7 +241,7 @@ touchBucket space interval bkts psq tr
 -- 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.
-forkPollForRefresh :: POSIXTime -> TVar (Int.PSQ POSIXTime) -> (Int -> IO ()) -> IO ThreadId
+forkPollForRefresh :: POSIXTime -> TVar (Int.PSQ POSIXTime) -> (Int -> IO a) -> IO ThreadId
 forkPollForRefresh interval psq refresh = do
     fork $ do
       myThreadId >>= flip labelThread "pollForRefresh"
@@ -260,14 +260,15 @@ forkPollForRefresh interval psq refresh = do
                              -- Now fork the refresh operation.
                              -- TODO: We should probably propogate the kill signal to this thread.
                              fork $ do myThreadId >>= flip labelThread ("refresh."++show bktnum)
-                                       refresh bktnum
+                                       _ <- refresh bktnum
+                                       return ()
                              return ()
             seconds    -> threadDelay ( seconds * 1000000 )
         again
 
 refreshBucket :: forall nid ni 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 ()
+                 Search nid addr ni ni -> TVar (BucketList ni) -> nid -> Int -> IO Int
 refreshBucket sch var nid n = do
     tbl <- atomically (readTVar var)
     let count = bktCount tbl
@@ -275,19 +276,22 @@ 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 < count))
+    fin <- atomically $ newTVar False
     resultCounter <- atomically $ newTVar Set.empty
     let fullcount = R.defaultBucketSize
-    let checkBucketFull :: ni -> STM Bool
+        saveit True = writeTVar fin True >> return True
+        saveit _    = return False
+        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 (Set.insert found_node)
             resultCount <- readTVar resultCounter
-            case drop (n - 1) counts of
-                (cnt:_) | cnt < fullcount            -> return True
-                _ | Set.size resultCount < fullcount -> return True
-                _                                    -> return False
+            saveit $ case drop (n - 1) counts of
+                (cnt:_) | cnt < fullcount            -> True
+                _ | Set.size resultCount < fullcount -> True
+                _                                    -> False
 
     hPutStrLn stderr $ "Start refresh " ++ show (n,sample)
 
@@ -299,7 +303,11 @@ refreshBucket sch var nid n = do
             atomically $ searchIsFinished s >>= check
     atomically $ searchCancel s
     hPutStrLn stderr $ "Finish refresh " ++ show (n,sample)
-    return ()
+    rcount <- atomically $ do
+        c <- Set.size <$> readTVar resultCounter
+        b <- readTVar fin
+        return $ if b then 1 else c
+    return rcount
 
 bootstrap ::
     ( Show nid
@@ -312,33 +320,41 @@ bootstrap ::
     , Ord nid
     , Traversable t1
     , Traversable t
-    ) => Search nid addr ni ni -> TVar (BucketList ni) -> (a -> IO Bool) -> t a -> t1 a -> IO ()
+    ) => Search nid addr ni ni -> TVar (BucketList ni) -> (ni -> IO Bool) -> t ni -> t1 ni -> IO ()
 bootstrap sch var ping ns ns0 = do
+    gotPing <- atomically $ newTVar False
+
     -- First, ping the given nodes so that they are added to
     -- our routing table.
-    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
-        -- nodes.
-        when (not got_response) $ do
-            _ <- mapConcurrently g ping ns0
-            return ()
-
-    mvar <- newMVar ()
+    withTaskGroup "bootstrap.resume" 20 $ \g -> do
+        forM_ ns $ \n -> do
+            let lbl = show $ kademliaLocation (searchSpace sch) n
+            forkTask g lbl $ do
+                b <- ping n
+                when b $ atomically $ writeTVar gotPing True
+
+    -- We resort to the hardcoded fallback nodes only when we got no
+    -- responses.  This is to lesson the burden on well-known boostrap
+    -- nodes.
+    fallback <- atomically (readTVar gotPing) >>= return . when . not
+    fallback $ withTaskGroup "bootstrap.ping" 20 $ \g -> do
+        forM_ ns0 $ \n -> do
+            forkTask g (show $ kademliaLocation (searchSpace sch) n)
+                       (void $ ping n)
+    hPutStrLn stderr "Finished bootstrap pings."
+
     -- 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
+    flip fix 1 $ \again cnt -> do
+        when (cnt==0) $ do
             -- Force a delay in case the search returns too quickly
+            hPutStrLn stderr $ "Zero results, forcing 1 minute delay"
             threadDelay (60 * 1000000)
-            putMVar mvar ()
-        let shp    = reverse $ zip (R.shape tbl) [0 .. ]
+        tbl <- atomically $ readTVar var
+        let shp    = zip (R.shape tbl) [0 .. ]
             unfull = filter ( (< R.defaultBucketSize) . fst ) shp
         case dropWhile ((> R.defaultBucketCount - 1) . snd) unfull of
             [] -> do
@@ -346,15 +362,15 @@ bootstrap sch var ping ns ns0 = do
                     -- Not enough buckets, keep trying.
                     hPutStrLn stderr
                         $ "Not enough buckets, refresh " ++ show (R.defaultBucketCount - 1)
-                    refreshBucket sch var
+                    cnt <- refreshBucket sch var
                         (kademliaLocation (searchSpace sch) (thisNode tbl))
                         (R.defaultBucketCount - 1)
-                    again
+                    again cnt
             (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
+                cnt <- refreshBucket sch var (kademliaLocation (searchSpace sch) (thisNode tbl)) num'
+                again cnt