Fix search bug.

1 files changed, 19 insertions, 11 deletions

diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
index 1fe73c30..79cc9489 100644
--- a/src/Network/BitTorrent/DHT/Search.hs
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -67,26 +67,34 @@ sendQuery IterativeSearch{..} (ni :-> d) = do
         modifyTVar searchInformant $ MM.insertTake searchK ni d
         modifyTVar searchResults   $ \s -> foldr Set.insert s rs
 
+
+searchIsFinished :: Ord ip => IterativeSearch ip r -> STM Bool
+searchIsFinished IterativeSearch{..} = do
+    q <- readTVar searchQueued
+    cnt <- readTVar searchPendingCount
+    informants <- readTVar searchInformant
+    return $ cnt == 0
+             && ( MM.null q
+                  || ( MM.size informants >= searchK
+                       && ( PSQ.prio (fromJust $ MM.findMax informants)
+                            <= PSQ.prio (fromJust $ MM.findMin q))))
+
 search ::
     (Ord r, Ord ip) =>
     IterativeSearch ip r -> IO ()
 search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do
     fix $ \again -> do
         join $ atomically $ do
-            found <- MM.minView <$> readTVar searchQueued
             cnt <- readTVar $ searchPendingCount
+            informants <- readTVar searchInformant
+            found <- MM.minView <$> readTVar searchQueued
             case found of
-                Nothing            -> retry
-                Just (ni :-> d, q) -> do
-                    informants <- readTVar searchInformant
-                    if MM.size informants < searchK
-                       && (cnt > 0 || not (MM.null q))
-                       || PSQ.prio (fromJust $ MM.findMax informants) > d
-                    then do
-                        writeTVar searchQueued q
+                Just (ni :-> d, q)
+                  | (MM.size informants < searchK) && (cnt > 0 || not (MM.null q))
+                    || (PSQ.prio (fromJust $ MM.findMax informants) > d)
+                  -> do writeTVar searchQueued q
                         modifyTVar searchVisited $ Set.insert (nodeAddr ni)
                         modifyTVar searchPendingCount succ
                         return $ withAsync g (sendQuery s (ni :-> d)) (const again)
-                    else do
-                        check (cnt == 0)
+                _ -> do check (cnt == 0)
                         return $ return ()