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

1 files changed, 13 insertions, 13 deletions

diff --git a/dht/src/Network/BitTorrent/MainlineDHT.hs b/dht/src/Network/BitTorrent/MainlineDHT.hs
index fc69fedd..8532b492 100644
--- a/dht/src/Network/BitTorrent/MainlineDHT.hs
+++ b/dht/src/Network/BitTorrent/MainlineDHT.hs
@@ -70,7 +70,7 @@ import Network.Kademlia.Search          (Search (..))
 import Network.BitTorrent.DHT.Token           as Token
 import qualified Network.Kademlia.Routing          as R
          ;import Network.Kademlia.Routing          (getTimestamp)
-import Network.QueryResponse
+import Network.QueryResponse                  as QR
 import Network.Socket
 import System.IO.Error
 import System.IO.Unsafe                       (unsafeInterleaveIO)
@@ -569,7 +569,7 @@ newClient swarms addr = do
         -- We defer initializing the refreshSearch and refreshPing until we
         -- have a client to send queries with.
         let nullPing = const $ return False
-            nullSearch = mainlineSearch $ \_ _ -> return Nothing
+            nullSearch = mainlineSearch $ \_ _ -> return Canceled
         tbl4 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info R.defaultBucketCount
         refresher4 <- newBucketRefresher tbl4 nullSearch nullPing
         tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info6 R.defaultBucketCount
@@ -1045,14 +1045,14 @@ mainlineSend :: ( BEncode xqry
                   -> MainlineClient
                   -> qry
                   -> NodeInfo
-                  -> IO (Maybe rsp)
+                  -> IO (QR.Result rsp)
 mainlineSend meth unwrap msg client nid addr = do
     reply <- sendQuery client serializer (msg nid) addr
     return $ case reply of
-        Success (Right x) -> Just x
-        Success (Left e)  -> Nothing -- TODO: Do something with parse errors.
-        Canceled          -> Nothing
-        TimedOut          -> Nothing
+        Success (Right x) -> Success x
+        Success (Left e)  -> Canceled -- TODO: Do something with parse errors.
+        Canceled          -> Canceled
+        TimedOut          -> TimedOut
  where
     serializer = MethodSerializer
         { methodTimeout  = \ni -> return (ni, 5000000)
@@ -1066,23 +1066,23 @@ mainlineSend meth unwrap msg client nid addr = do
 
 ping :: MainlineClient -> NodeInfo -> IO Bool
 ping client addr =
-    fromMaybe False
+    fromMaybe False . resultToMaybe
     <$> mainlineSend (Method "ping") (\Pong -> True) (const Ping) client () addr
 
 -- searchQuery        :: ni -> IO (Maybe [ni], [r], tok))
-getNodes :: MainlineClient -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
+getNodes :: MainlineClient -> NodeId -> NodeInfo -> IO (QR.Result ([NodeInfo],[NodeInfo],Maybe ()))
 getNodes = mainlineSend (Method "find_node") unwrapNodes $ flip FindNode (Just Want_Both)
 
 unwrapNodes :: NodeFound -> ([NodeInfo], [NodeInfo], Maybe ())
 unwrapNodes (NodeFound ns4 ns6) = (ns4++ns6, ns4++ns6, Just ())
 
-getPeers :: MainlineClient -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[PeerAddr],Maybe Token))
+getPeers :: MainlineClient -> NodeId -> NodeInfo -> IO (QR.Result ([NodeInfo],[PeerAddr],Maybe Token))
 getPeers = mainlineSend (Method "get_peers") unwrapPeers $ flip GetPeers (Just Want_Both) . coerce
 
 unwrapPeers :: GotPeers -> ([NodeInfo], [PeerAddr], Maybe Token)
 unwrapPeers (GotPeers ps (NodeFound ns4 ns6) tok) = (ns4++ns6, ps, Just tok)
 
-mainlineSearch :: (NodeId -> NodeInfo -> IO (Maybe ([NodeInfo], [r], Maybe tok)))
+mainlineSearch :: (NodeId -> NodeInfo -> IO (QR.Result ([NodeInfo], [r], Maybe tok)))
                   -> Search NodeId (IP, PortNumber) tok NodeInfo r
 mainlineSearch qry = Search
     { searchSpace       = mainlineSpace
@@ -1140,5 +1140,5 @@ resolve want hostAndPort = do
 
 
 announce :: MainlineClient -> Announce -> NodeInfo -> IO (Maybe Announced)
-announce client msg addr = do
-    mainlineSend (Method "announce_peer") id (\() -> msg) client () addr
+announce client msg addr =
+    resultToMaybe <$> mainlineSend (Method "announce_peer") id (\() -> msg) client () addr