Use async queries for all UDP kademlia searches.

1 files changed, 46 insertions, 8 deletions

diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index 83865c98..89851e88 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -563,7 +563,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 $ Left $ \_ _ -> return Nothing
         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
@@ -1035,13 +1035,36 @@ mainlineSend :: ( BEncode a
                   -> NodeInfo
                   -> IO (Maybe b)
 mainlineSend meth unwrap msg client nid addr = do
-    reply <- sendQuery client serializer (msg nid) addr
+    reply <- sendQuery client (mainlineSerializeer meth unwrap client) (msg nid) addr
     -- sendQuery will return (Just (Left _)) on a parse error.  We're going to
     -- blow it away with the join-either sequence.
     -- TODO: Do something with parse errors.
     return $ join $ either (const Nothing) Just <$> reply
- where
-    serializer = MethodSerializer
+
+mainlineAsync :: (BEncode a1, BEncode a2) =>
+                       Method
+                       -> (a2 -> a3)
+                       -> (t -> a1)
+                       -> Client String Method TransactionId NodeInfo (Message BValue)
+                       -> t
+                       -> NodeInfo
+                       -> (Maybe a3 -> IO ())
+                       -> IO ()
+mainlineAsync meth unwrap msg client nid addr onresult = do
+    asyncQuery client (mainlineSerializeer meth unwrap client) (msg nid) addr
+        $ \reply ->
+            -- sendQuery will return (Just (Left _)) on a parse error.  We're going to
+            -- blow it away with the join-either sequence.
+            -- TODO: Do something with parse errors.
+            onresult $ join $ either (const Nothing) Just <$> reply
+
+mainlineSerializeer :: (BEncode a2, BEncode a1) =>
+                        Method
+                        -> (a2 -> b)
+                        -> MainlineClient
+                        -> MethodSerializer
+                             TransactionId NodeInfo (Message BValue) Method a1 (Either Error b)
+mainlineSerializeer meth unwrap client = MethodSerializer
         { methodTimeout  = \_ ni -> return (ni, 5000000)
         , method         = meth
         , wrapQuery      = encodeQueryPayload meth (isReadonlyClient client)
@@ -1060,30 +1083,45 @@ ping client addr =
 getNodes :: MainlineClient -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
 getNodes = mainlineSend (Method "find_node") unwrapNodes $ flip FindNode (Just Want_Both)
 
+asyncGetNodes :: Client String Method TransactionId NodeInfo (Message BValue)
+                   -> NodeId
+                   -> NodeInfo
+                   -> (Maybe ([NodeInfo], [NodeInfo], Maybe ()) -> IO ())
+                   -> IO ()
+asyncGetNodes = mainlineAsync (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 = mainlineSend (Method "get_peers") unwrapPeers $ flip GetPeers (Just Want_Both) . coerce
 
+asyncGetPeers :: Client String Method TransactionId NodeInfo (Message BValue)
+                 -> NodeId
+                 -> NodeInfo
+                 -> (Maybe ([NodeInfo], [PeerAddr], Maybe Token) -> IO ())
+                 -> IO ()
+asyncGetPeers = mainlineAsync (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 :: Either (NodeId -> NodeInfo -> IO (Maybe ([NodeInfo], [r], Maybe tok)))
+                         (NodeId -> NodeInfo -> (Maybe ([NodeInfo], [r], Maybe tok) -> IO ()) -> IO ())
                   -> Search NodeId (IP, PortNumber) tok NodeInfo r
 mainlineSearch qry = Search
     { searchSpace       = mainlineSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = Left qry
+    , searchQuery       = qry
     , searchAlpha       = 8
     , searchK           = 16
     }
 
 nodeSearch :: MainlineClient -> Search NodeId (IP, PortNumber) () NodeInfo NodeInfo
-nodeSearch client = mainlineSearch (getNodes client)
+nodeSearch client = mainlineSearch (Right $ asyncGetNodes client)
 
 peerSearch :: MainlineClient -> Search NodeId (IP, PortNumber) Token NodeInfo PeerAddr
-peerSearch client = mainlineSearch (getPeers client)
+peerSearch client = mainlineSearch (Right $ asyncGetPeers client)
 
 -- | List of bootstrap nodes maintained by different bittorrent
 -- software authors.