1 files changed, 15 insertions, 8 deletions

diff --git a/kad/tests/searchCancel.hs b/kad/tests/searchCancel.hs
index e8aa33c7..393f10c1 100644
--- a/kad/tests/searchCancel.hs
+++ b/kad/tests/searchCancel.hs
@@ -2,6 +2,7 @@ import Control.Concurrent
 import Control.Concurrent.STM
 import Control.Monad
 import Data.Bits
+import GHC.Event
 
 import Network.Kademlia
 import Network.Kademlia.Bootstrap
@@ -14,15 +15,13 @@ import Network.QueryResponse as QR
 import qualified Data.MinMaxPSQ as MM
 import qualified Data.Set as Set
 
-makeSchResults :: TVar (Maybe (SearchState Int Int () Int Int Int)) -> TVar Int -> IO (Maybe ([Int],[Int],Maybe ()))
+makeSchResults :: TVar (Maybe (SearchState Int Int () Int Int Int)) -> TVar Int -> IO ([Int],[Int],Maybe ())
 makeSchResults mbv var = do
-    putStrLn "makeSchResults"
-    threadDelay 200000
     (r,io) <- atomically $ do
         n <- readTVar var
         let ns = take 4 [n .. ]
         writeTVar var $! n + 4
-        let r = Just (ns, ns, Just ()) -- Maybe ([ni], [r], Maybe tok)
+        let r = (ns, ns, Just ()) -- Maybe ([ni], [r], Maybe tok)
             stmio = if n > 490
                         then do
                             ms <- readTVar mbv
@@ -48,10 +47,18 @@ sch :: TVar (Maybe (SearchState Int Int () Int Int Int)) -> TVar Int -> Search I
 sch mbv var = Search
     { searchSpace       = kad
     , searchNodeAddress = id
-    , searchQuery       = \_ _ f -> do r <- makeSchResults mbv var
-                                       let qk = maybe 0 (\(ns,_,_) -> head ns) r
-                                       f qk $ maybe TimedOut Success r
-                                       return qk
+    -- searchQuery :: nid -> ni -> (qk -> Result ([ni], [r], Maybe tok) -> IO ()) -> IO qk
+    , searchQuery       = \_ _ f -> do
+        val <- atomically $ do
+            n <- readTVar var
+            writeTVar var $! n + 1
+            return n
+        tm <- getSystemTimerManager
+        k <- registerTimeout tm 200000 $ do
+            putStrLn $ "makeSchResults " ++ show val
+            r <- makeSchResults mbv var
+            f val (Success r)
+        return val
     , searchQueryCancel = \_ _ -> return ()
     , searchAlpha       = 4
     , searchK           = 8