Async search.

1 files changed, 64 insertions, 12 deletions

diff --git a/dht/src/Network/BitTorrent/MainlineDHT.hs b/dht/src/Network/BitTorrent/MainlineDHT.hs
index 8532b492..d3904c40 100644
--- a/dht/src/Network/BitTorrent/MainlineDHT.hs
+++ b/dht/src/Network/BitTorrent/MainlineDHT.hs
@@ -512,8 +512,8 @@ data Routing = Routing
     { tentativeId :: NodeInfo
     , committee4  :: TriadCommittee NodeId SockAddr
     , committee6  :: TriadCommittee NodeId SockAddr
-    , refresher4  :: BucketRefresher NodeId NodeInfo
-    , refresher6  :: BucketRefresher NodeId NodeInfo
+    , refresher4  :: BucketRefresher NodeId NodeInfo TransactionId
+    , refresher6  :: BucketRefresher NodeId NodeInfo TransactionId
     }
 
 sched4 :: Routing -> TVar (Int.PSQ POSIXTime)
@@ -569,7 +569,6 @@ 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 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
@@ -730,7 +729,7 @@ defaultHandler meth = MethodHandler decodePayload errorPayload returnError
 
 mainlineKademlia :: MainlineClient
                     -> TriadCommittee NodeId SockAddr
-                    -> BucketRefresher NodeId NodeInfo
+                    -> BucketRefresher NodeId NodeInfo TransactionId
                     -> Kademlia NodeId NodeInfo
 mainlineKademlia client committee refresher
     = Kademlia quietInsertions
@@ -1037,6 +1036,35 @@ announceH (SwarmsDatabase peers toks _) naddr announcement = do
 isReadonlyClient :: MainlineClient -> Bool
 isReadonlyClient client = False -- TODO
 
+mainlineAsync :: ( BEncode xqry
+                 , BEncode xrsp
+                 ) => Method
+                  -> (xrsp -> rsp)
+                  -> (qry -> xqry)
+                  -> MainlineClient
+                  -> qry
+                  -> NodeInfo
+                  -> (TransactionId -> QR.Result rsp -> IO ())
+                  -> IO TransactionId
+mainlineAsync meth unwrap msg client nid addr withResult = do
+    asyncQuery client serializer (msg nid) addr $ \qid reply -> do
+        withResult qid $ case reply of
+            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)
+        , method         = meth
+        , wrapQuery      = encodeQueryPayload meth (isReadonlyClient client)
+        , unwrapResponse = (>>= either (Left . Error GenericError . C8.pack)
+                                       (Right . unwrap)
+                                       . BE.fromBEncode)
+                           . rspPayload
+        }
+
+
 mainlineSend :: ( BEncode xqry
                 , BEncode xrsp
                 ) => Method
@@ -1073,30 +1101,54 @@ ping client addr =
 getNodes :: MainlineClient -> NodeId -> NodeInfo -> IO (QR.Result ([NodeInfo],[NodeInfo],Maybe ()))
 getNodes = mainlineSend (Method "find_node") unwrapNodes $ flip FindNode (Just Want_Both)
 
+asyncGetNodes :: MainlineClient -> NodeId -> NodeInfo -> (TransactionId -> QR.Result ([NodeInfo],[NodeInfo],Maybe ()) -> IO ())
+                                -> IO TransactionId
+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 (QR.Result ([NodeInfo],[PeerAddr],Maybe Token))
 getPeers = mainlineSend (Method "get_peers") unwrapPeers $ flip GetPeers (Just Want_Both) . coerce
 
+asyncGetPeers :: MainlineClient -> NodeId -> NodeInfo -> (TransactionId -> QR.Result ([NodeInfo],[PeerAddr],Maybe Token) -> IO ())
+                                -> IO TransactionId
+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 (QR.Result ([NodeInfo], [r], Maybe tok)))
-                  -> Search NodeId (IP, PortNumber) tok NodeInfo r
-mainlineSearch qry = Search
+nullTransactionId :: TransactionId
+nullTransactionId = TransactionId B.empty
+
+nullSearch :: Search NodeId (IP, PortNumber) tok NodeInfo r TransactionId
+nullSearch = Search
+    { searchSpace       = mainlineSpace
+    , searchNodeAddress = nodeIP &&& nodePort
+    , searchQuery       = \_ _ f -> f nullTransactionId Canceled >> return nullTransactionId
+    , searchQueryCancel = \_ _ -> return ()
+    , searchAlpha       = 8
+    , searchK           = 16
+    }
+
+mainlineSearch :: MainlineClient
+                  -> (MainlineClient -> NodeId -> NodeInfo
+                      -> (TransactionId -> QR.Result ([NodeInfo], [r], Maybe tok) -> IO ()) -> IO TransactionId)
+                  -> Search NodeId (IP, PortNumber) tok NodeInfo r TransactionId
+mainlineSearch client qry = Search
     { searchSpace       = mainlineSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = qry
+    , searchQuery       = qry client
+    , searchQueryCancel = cancelQuery client
     , searchAlpha       = 8
     , searchK           = 16
     }
 
-nodeSearch :: MainlineClient -> Search NodeId (IP, PortNumber) () NodeInfo NodeInfo
-nodeSearch client = mainlineSearch (getNodes client)
+nodeSearch :: MainlineClient -> Search NodeId (IP, PortNumber) () NodeInfo NodeInfo TransactionId
+nodeSearch client = mainlineSearch client asyncGetNodes
 
-peerSearch :: MainlineClient -> Search NodeId (IP, PortNumber) Token NodeInfo PeerAddr
-peerSearch client = mainlineSearch (getPeers client)
+peerSearch :: MainlineClient -> Search NodeId (IP, PortNumber) Token NodeInfo PeerAddr TransactionId
+peerSearch client = mainlineSearch client asyncGetPeers
 
 -- | List of bootstrap nodes maintained by different bittorrent
 -- software authors.