2 files changed, 20 insertions, 17 deletions

diff --git a/kad/src/Network/Kademlia/Search.hs b/kad/src/Network/Kademlia/Search.hs
index e30c40da..19d0df69 100644
--- a/kad/src/Network/Kademlia/Search.hs
+++ b/kad/src/Network/Kademlia/Search.hs
@@ -44,7 +44,7 @@ data Search nid addr tok ni r qk = Search
     { searchSpace        :: KademliaSpace nid ni
     , searchNodeAddress  :: ni -> addr
     , searchQuery        :: nid -> ni -> (qk -> Result ([ni], [r], Maybe tok) -> IO ()) -> IO qk
-    , searchQueryCancel  :: qk -> STM ()
+    , searchQueryCancel  :: (IO () -> STM ()) -> qk -> STM ()
     , searchAlpha        :: Int -- α = 8
       -- | 'searchK' should be larger than 'searchAlpha'.  How much larger depends on
       -- how fast the queries are.  For Tox's much slower onion-routed queries, we
@@ -118,19 +118,21 @@ reset :: (Ord ni, Ord nid, Hashable ni, Hashable nid) =>
        -> Search nid addr1 tok1 ni r1 qk
        -> nid
        -> SearchState nid addr tok ni r qk
-       -> STM (SearchState nid addr tok ni r qk)
+       -> IO (SearchState nid addr tok ni r qk)
 reset nearestNodes qsearch target st = do
-    pc <- readTVar (searchPendingCount st)
-    check (pc == 0)
-    searchIsFinished st >>= check -- Wait for a search to finish before resetting.
-    bktNodes <- map (\ni -> ni :-> kademliaLocation (searchSpace qsearch) ni)
-                <$> nearestNodes target
-    priorInformants <- map stripValue . MM.toList <$> readTVar (searchInformant  st)
-    writeTVar (searchQueued       st) $ Just $ MM.fromList $ priorInformants ++ bktNodes
-    writeTVar (searchInformant    st) MM.empty
-    writeTVar (searchVisited      st) Set.empty
-    writeTVar (searchPendingCount st) 0
-    return st
+    atomically $ searchCancel st
+    atomically $ do
+        pc <- readTVar (searchPendingCount st)
+        check (pc == 0)
+        searchIsFinished st >>= check -- Wait for a search to finish before resetting.
+        bktNodes <- map (\ni -> ni :-> kademliaLocation (searchSpace qsearch) ni)
+                    <$> nearestNodes target
+        priorInformants <- map stripValue . MM.toList <$> readTVar (searchInformant  st)
+        writeTVar (searchQueued       st) $ Just $ MM.fromList $ priorInformants ++ bktNodes
+        writeTVar (searchInformant    st) MM.empty
+        writeTVar (searchVisited      st) Set.empty
+        writeTVar (searchPendingCount st) 0
+        return st
 
 grokQuery :: forall addr nid tok ni r qk.
             ( Ord addr
@@ -233,8 +235,9 @@ searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni,
            -> IO ()
 searchLoop sch@Search{..} target result s@SearchState{..} = do
     myThreadId >>= flip labelThread ("search."++show target)
-    withTaskGroup ("search.g."++show target) searchAlpha $ \g -> fix $ \again -> do
-        join $ atomically $ do
+    iochan <- atomically newTChan
+    fix $ \again -> do
+        join $ atomically $ orElse (fmap (>> again) $ readTChan iochan) $ do
             cnt <- readTVar $ searchPendingCount
             check (cnt <= searchAlpha) -- Only searchAlpha pending queries at a time.
             informants <- readTVar searchInformant
@@ -256,6 +259,6 @@ searchLoop sch@Search{..} target result s@SearchState{..} = do
                         return $ do
                             qk <- searchQuery target ni $
                                     \qk reply -> grokQuery sch target result s (ni :-> d) qk reply
-                            atomically $ modifyTVar' searchPending $ Map.insert qk (searchQueryCancel qk)
+                            atomically $ modifyTVar' searchPending $ Map.insert qk (searchQueryCancel (writeTChan iochan) qk)
                             again
                 _ -> searchIsFinished s >>= check >> return (return ())
diff --git a/kad/tests/searchCancel.hs b/kad/tests/searchCancel.hs
index 33860a2f..e8aa33c7 100644
--- a/kad/tests/searchCancel.hs
+++ b/kad/tests/searchCancel.hs
@@ -52,7 +52,7 @@ sch mbv var = Search
                                        let qk = maybe 0 (\(ns,_,_) -> head ns) r
                                        f qk $ maybe TimedOut Success r
                                        return qk
-    , searchQueryCancel = \_ -> return ()
+    , searchQueryCancel = \_ _ -> return ()
     , searchAlpha       = 4
     , searchK           = 8
     }