1 files changed, 45 insertions, 31 deletions

diff --git a/kad/src/Network/Kademlia/Search.hs b/kad/src/Network/Kademlia/Search.hs
index 8d9c997b..e30c40da 100644
--- a/kad/src/Network/Kademlia/Search.hs
+++ b/kad/src/Network/Kademlia/Search.hs
@@ -19,6 +19,8 @@ import Control.Concurrent.Tasks
 import Control.Concurrent.STM
 import Control.Monad
 import Data.Function
+import qualified Data.Map.Strict as Map
+         ;import Data.Map.Strict (Map)
 import Data.Maybe
 import qualified Data.Set            as Set
          ;import Data.Set            (Set)
@@ -38,10 +40,11 @@ import Control.Concurrent.Lifted
 import GHC.Conc                  (labelThread)
 #endif
 
-data Search nid addr tok ni r = Search
+data Search nid addr tok ni r qk = Search
     { searchSpace        :: KademliaSpace nid ni
     , searchNodeAddress  :: ni -> addr
-    , searchQuery        :: nid -> ni -> IO (Result ([ni], [r], Maybe tok))
+    , searchQuery        :: nid -> ni -> (qk -> Result ([ni], [r], Maybe tok) -> IO ()) -> IO qk
+    , searchQueryCancel  :: 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
@@ -55,10 +58,11 @@ data Search nid addr tok ni r = Search
     , searchK            :: Int -- K = 16
     }
 
-data SearchState nid addr tok ni r = SearchState
+data SearchState nid addr tok ni r qk = SearchState
     { -- | The number of pending queries.  Incremented before any query is sent
       -- and decremented when we get a reply.
       searchPendingCount :: TVar Int
+    , searchPending      :: TVar (Map qk (STM ()))
       -- | Nodes scheduled to be queried (roughly at most K).
       --
       -- This will be set to Nothing when a search is canceled.
@@ -72,7 +76,7 @@ data SearchState nid addr tok ni r = SearchState
       -- set.  It also can be a bounded 'MinMaxPSQ', although searchAlpha
       -- should limit the number of outstanding queries.
     , searchVisited      :: TVar (Set addr)
-    , searchSpec         :: Search nid addr tok ni r
+    , searchSpec         :: Search nid addr tok ni r qk
     }
 
 
@@ -87,12 +91,13 @@ newSearch :: ( Ord addr
              -> (r -> STM Bool)        -- receives search results.
              -> nid                    -- target of search
              -}
-             Search nid addr tok ni r
+             Search nid addr tok ni r qk
              -> nid
              -> [ni]                   -- Initial nodes to query.
-             -> STM (SearchState nid addr tok ni r)
-newSearch s@(Search space nAddr qry _ _) target ns = do
+             -> STM (SearchState nid addr tok ni r qk)
+newSearch s@(Search space nAddr qry _ _ _) target ns = do
     c <- newTVar 0
+    p <- newTVar Map.empty
     q <- newTVar $ Just
                  $ MM.fromList
                  $ map (\n -> n :-> kademliaXor space target (kademliaLocation space n))
@@ -100,7 +105,7 @@ newSearch s@(Search space nAddr qry _ _) target ns = do
     i <- newTVar MM.empty
     v <- newTVar Set.empty
     return -- (Search space nAddr qry) , r , target
-           ( SearchState c q i v s )
+           ( SearchState c p q i v s )
 
 -- | Discard a value from a key-priority-value tuple. This is useful for
 -- swaping items from a "MinMaxPSQ'" to a "MinMaxPSQ".
@@ -110,11 +115,13 @@ stripValue (Binding ni _ nid) = (ni :-> nid)
 -- | Reset a 'SearchState' object to ready it for a repeated search.
 reset :: (Ord ni, Ord nid, Hashable ni, Hashable nid) =>
        (nid -> STM [ni])
-       -> Search nid addr1 tok1 ni r1
+       -> Search nid addr1 tok1 ni r1 qk
        -> nid
-       -> SearchState nid addr tok ni r
-       -> STM (SearchState nid addr tok ni r)
+       -> SearchState nid addr tok ni r qk
+       -> STM (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
@@ -125,25 +132,29 @@ reset nearestNodes qsearch target st = do
     writeTVar (searchPendingCount st) 0
     return st
 
-sendQuery :: forall addr nid tok ni r.
+grokQuery :: forall addr nid tok ni r qk.
             ( Ord addr
             , PSQKey nid
             , PSQKey ni
             , Show nid
+            , Ord qk
             ) =>
-            Search nid addr tok ni r
+            Search nid addr tok ni r qk
             -> nid
             -> (r -> STM Bool) -- ^ return False to terminate the search.
-            -> SearchState nid addr tok ni r
+            -> SearchState nid addr tok ni r qk
             -> Binding ni nid
+            -> qk
+            -> Result ([ni],[r],Maybe tok)
             -> IO ()
-sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) = do
-    myThreadId >>= flip labelThread ("searchQuery." ++ show searchTarget)
-    reply <- searchQuery searchTarget ni `catchIOError` const (return Canceled)
+grokQuery Search{..} searchTarget withSearchResult sch@SearchState{..} (ni :-> d) qk reply = do
+    -- myThreadId >>= flip labelThread ("searchQuery." ++ show searchTarget)
+    -- reply <- searchQuery searchTarget ni `catchIOError` const (return Nothing)
     -- (ns,rs)
-    let tok = error "TODO: token"
+    -- let tok = error "TODO: token"
     atomically $ do
-        modifyTVar searchPendingCount pred
+        modifyTVar' searchPendingCount pred
+        modifyTVar' searchPending $ Map.delete qk
         case reply of
             Success x -> go x
             _         -> return ()
@@ -170,7 +181,7 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
         modifyTVar searchInformant $ MM.insertTake' (searchK - searchAlpha) ni tok d
         flip fix rs $ \loop -> \case
             r:rs' -> do
-                wanting <- searchResult r
+                wanting <- withSearchResult r
                 if wanting then loop rs'
                            else searchCancel sch
             [] -> return ()
@@ -178,7 +189,7 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
 
 searchIsFinished :: ( PSQKey nid
                     , PSQKey ni
-                    ) => SearchState nid addr tok ni r -> STM Bool
+                    ) => SearchState nid addr tok ni r qk -> STM Bool
 searchIsFinished SearchState{..} = do
     readTVar searchQueued >>= \case
         Just q -> do
@@ -191,9 +202,11 @@ searchIsFinished SearchState{..} = do
                                         <= PSQ.prio (fromJust $ MM.findMin q))))
         Nothing -> return True
 
-searchCancel :: SearchState nid addr tok ni r -> STM ()
+searchCancel :: SearchState nid addr tok ni r qk -> STM ()
 searchCancel SearchState{..} = do
     writeTVar searchQueued Nothing
+    m <- readTVar searchPending
+    foldr (>>) (return ()) m
 
 search ::
     ( Ord r
@@ -201,7 +214,8 @@ search ::
     , PSQKey nid
     , PSQKey ni
     , Show nid
-    ) => Search nid addr tok ni r -> R.BucketList ni -> nid -> (r -> STM Bool) -> IO (SearchState nid addr tok ni r, ThreadId)
+    , Ord qk
+    ) => Search nid addr tok ni r qk -> R.BucketList ni -> nid -> (r -> STM Bool) -> IO (SearchState nid addr tok ni r qk, ThreadId)
 search sch buckets target result = do
     let ns = R.kclosest (searchSpace sch) (searchK sch) target buckets
     st <- atomically $ newSearch sch target ns
@@ -211,11 +225,11 @@ search sch buckets target result = do
     atomically $ writeTVar v True
     return (st,t)
 
-searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni )
-           => Search nid addr tok ni r      -- ^ Query and distance methods.
-           -> nid                           -- ^ The target we are searching for.
-           -> (r -> STM Bool)               -- ^ Invoked on each result.  Return False to quit searching.
-           -> SearchState nid addr tok ni r -- ^ Search-related state.
+searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni, Ord qk )
+           => Search nid addr tok ni r qk      -- ^ Query and distance methods.
+           -> nid                              -- ^ The target we are searching for.
+           -> (r -> STM Bool)                  -- ^ Invoked on each result.  Return False to quit searching.
+           -> SearchState nid addr tok ni r qk -- ^ Search-related state.
            -> IO ()
 searchLoop sch@Search{..} target result s@SearchState{..} = do
     myThreadId >>= flip labelThread ("search."++show target)
@@ -240,8 +254,8 @@ 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)
+                            qk <- searchQuery target ni $
+                                    \qk reply -> grokQuery sch target result s (ni :-> d) qk reply
+                            atomically $ modifyTVar' searchPending $ Map.insert qk (searchQueryCancel qk)
                             again
                 _ -> searchIsFinished s >>= check >> return (return ())