4 files changed, 23 insertions, 17 deletions

diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index 180ae82d..83865c98 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -1074,7 +1074,7 @@ mainlineSearch :: (NodeId -> NodeInfo -> IO (Maybe ([NodeInfo], [r], Maybe tok))
 mainlineSearch qry = Search
     { searchSpace       = mainlineSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = qry
+    , searchQuery       = Left qry
     , searchAlpha       = 8
     , searchK           = 16
     }
diff --git a/src/Network/Kademlia/Search.hs b/src/Network/Kademlia/Search.hs
index e87a8618..1be1afc1 100644
--- a/src/Network/Kademlia/Search.hs
+++ b/src/Network/Kademlia/Search.hs
@@ -31,7 +31,8 @@ import GHC.Conc                  (labelThread)
 data Search nid addr tok ni r = Search
     { searchSpace        :: KademliaSpace nid ni
     , searchNodeAddress  :: ni -> addr
-    , searchQuery        :: nid -> ni -> IO (Maybe ([ni], [r], Maybe tok))
+    , searchQuery        :: Either (nid -> ni -> IO (Maybe ([ni], [r], Maybe tok)))
+                                   (nid -> ni -> (Maybe ([ni],[r],Maybe tok) -> IO ()) -> IO ())
     , 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
@@ -112,7 +113,7 @@ reset nearestNodes qsearch target st = do
     writeTVar (searchPendingCount st) 0
     return st
 
-sendQuery :: forall addr nid tok ni r.
+sendAsyncQuery :: forall addr nid tok ni r.
             ( Ord addr
             , PSQKey nid
             , PSQKey ni
@@ -123,15 +124,22 @@ sendQuery :: forall addr nid tok ni r.
             -> (r -> STM Bool) -- ^ return False to terminate the search.
             -> SearchState nid addr tok ni r
             -> Binding ni nid
+            -> TaskGroup
             -> IO ()
-sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) = do
-    myThreadId >>= flip labelThread ("searchQuery." ++ show searchTarget)
-    reply <- searchQuery searchTarget ni `catchIOError` const (return Nothing)
-    -- (ns,rs)
-    let tok = error "TODO: token"
-    atomically $ do
-        modifyTVar searchPendingCount pred
-        maybe (return ()) go reply
+sendAsyncQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) g =
+    case searchQuery of
+        Left blockingQuery ->
+          forkTask g "searchQuery" $ do
+            myThreadId >>= flip labelThread ("searchQuery." ++ show searchTarget)
+            reply <- blockingQuery searchTarget ni `catchIOError` const (return Nothing)
+            atomically $ do
+                modifyTVar searchPendingCount pred
+                maybe (return ()) go reply
+        Right nonblockingQuery -> do
+            nonblockingQuery searchTarget ni $ \reply ->
+                atomically $ do
+                    modifyTVar searchPendingCount pred
+                    maybe (return ()) go reply
  where
     go (ns,rs,tok) = do
         vs <- readTVar searchVisited
@@ -221,9 +229,7 @@ searchLoop sch@Search{..} target result s@SearchState{..} = do
                         modifyTVar searchVisited $ Set.insert (searchNodeAddress ni)
                         modifyTVar searchPendingCount succ
                         return $ do
-                            forkTask g
-                                     "searchQuery"
-                                     $ sendQuery sch target result s (ni :-> d)
+                            sendAsyncQuery sch target result s (ni :-> d) g
                             again
                 _ -> -- Otherwise, we are finished.
                      do check (cnt == 0)
diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index afdf2cc3..abd607c3 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -194,7 +194,7 @@ newRouting addr crypto update4 update6 = do
             nullSearch = Search
                 { searchSpace       = toxSpace
                 , searchNodeAddress = nodeIP &&& nodePort
-                , searchQuery       = \_ _ -> return Nothing
+                , searchQuery       = Left $ \_ _ -> return Nothing
                 , searchAlpha       = 1
                 , searchK           = 2
                 }
@@ -525,7 +525,7 @@ nodeSearch :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> Search NodeI
 nodeSearch client cbvar = Search
     { searchSpace       = toxSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = getNodes client cbvar
+    , searchQuery       = Left $ getNodes client cbvar
     , searchAlpha       = 8
     , searchK           = 16
 
diff --git a/src/Network/Tox/Onion/Handlers.hs b/src/Network/Tox/Onion/Handlers.hs
index 52cc298d..edbbbb49 100644
--- a/src/Network/Tox/Onion/Handlers.hs
+++ b/src/Network/Tox/Onion/Handlers.hs
@@ -218,7 +218,7 @@ toxidSearch :: (TransactionId -> OnionDestination r -> STM (OnionDestination r,
 toxidSearch getTimeout crypto client = Search
     { searchSpace       = toxSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = getRendezvous getTimeout crypto client
+    , searchQuery       = Left $ getRendezvous getTimeout crypto client
     , searchAlpha = 3
     , searchK     = 6
     }