Modify kademlia search to distinguish a Canceled from timed-out query.

1 files changed, 7 insertions, 4 deletions

diff --git a/kad/src/Network/Kademlia/Search.hs b/kad/src/Network/Kademlia/Search.hs
index 03c18d0e..8d9c997b 100644
--- a/kad/src/Network/Kademlia/Search.hs
+++ b/kad/src/Network/Kademlia/Search.hs
@@ -29,7 +29,8 @@ import qualified Data.MinMaxPSQ   as MM
          ;import Data.MinMaxPSQ   (MinMaxPSQ, MinMaxPSQ')
 import qualified Data.Wrapper.PSQ as PSQ
          ;import Data.Wrapper.PSQ (pattern (:->), Binding, pattern Binding, Binding', PSQKey)
-import Network.Kademlia.Routing as R
+import Network.Kademlia.Routing   as R
+import Network.QueryResponse      (Result(..))
 #ifdef THREAD_DEBUG
 import Control.Concurrent.Lifted.Instrument
 #else
@@ -40,7 +41,7 @@ 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        :: nid -> ni -> IO (Result ([ni], [r], Maybe tok))
     , 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
@@ -138,12 +139,14 @@ sendQuery :: forall addr nid tok ni r.
             -> IO ()
 sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) = do
     myThreadId >>= flip labelThread ("searchQuery." ++ show searchTarget)
-    reply <- searchQuery searchTarget ni `catchIOError` const (return Nothing)
+    reply <- searchQuery searchTarget ni `catchIOError` const (return Canceled)
     -- (ns,rs)
     let tok = error "TODO: token"
     atomically $ do
         modifyTVar searchPendingCount pred
-        maybe (return ()) go reply
+        case reply of
+            Success x -> go x
+            _         -> return ()
  where
     go (ns,rs,tok) = do
         vs <- readTVar searchVisited