Use async queries for all UDP kademlia searches.

1 files changed, 55 insertions, 14 deletions

diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index abd607c3..1eec93b9 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -400,20 +400,61 @@ unsendNodes _                    = Nothing
 unwrapNodes :: SendNodes -> ( [NodeInfo], [NodeInfo], Maybe () )
 unwrapNodes (SendNodes ns) = (ns,ns,Just ())
 
+data SendableQuery x a b = SendableQuery
+    { sendableSerializer :: MethodSerializer TransactionId NodeInfo Message PacketKind a (Maybe x)
+    , sendableQuery      :: NodeId -> a
+    , sendableResult     :: Maybe (Maybe x) -> IO b
+    }
+
+sendQ :: SendableQuery x a b
+               -> QR.Client err PacketKind TransactionId NodeInfo Message
+               -> NodeId
+               -> NodeInfo
+               -> IO b
+sendQ s client nid addr = do
+    reply <- QR.sendQuery client (sendableSerializer s) (sendableQuery s nid) addr
+    sendableResult s reply
+
+asyncQ :: SendableQuery x a b
+            -> QR.Client err PacketKind TransactionId NodeInfo Message
+            -> NodeId
+            -> NodeInfo
+            -> (b -> IO ())
+            -> IO ()
+asyncQ s client nid addr go = do
+    QR.asyncQuery client (sendableSerializer s) (sendableQuery s nid) addr
+        $ sendableResult s >=> go
+
+getNodesSendable :: TVar (HashMap NodeId [NodeInfoCallback])
+                    -> NodeInfo
+                    -> SendableQuery SendNodes GetNodes (Maybe ([NodeInfo], [NodeInfo], Maybe ()))
+getNodesSendable cbvar addr = SendableQuery (serializer GetNodesType DHTGetNodes unsendNodes)
+                                            GetNodes
+                                            go
+  where
+    go reply = do
+        forM_ (join reply) $ \(SendNodes ns) ->
+            forM_ ns $ \n -> do
+                now <- getPOSIXTime
+                atomically $ do
+                    mcbs <- HashMap.lookup (nodeId n) <$> readTVar cbvar
+                    forM_ mcbs $ \cbs -> do
+                        forM_ cbs $ \cb -> do
+                            rumoredAddress cb now (nodeAddr addr) n
+        return $ fmap unwrapNodes $ join reply
+
 getNodes :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
-getNodes client cbvar nid addr = do
-    -- dput XMisc $ show addr ++ " <-- getnodes " ++ show nid
-    reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
-    -- dput XMisc $ show addr ++ " -sendnodes-> " ++ show reply
-    forM_ (join reply) $ \(SendNodes ns) -> 
-        forM_ ns $ \n -> do
-            now <- getPOSIXTime
-            atomically $ do
-                mcbs <- HashMap.lookup (nodeId n) <$> readTVar cbvar
-                forM_ mcbs $ \cbs -> do
-                    forM_ cbs $ \cb -> do
-                        rumoredAddress cb now (nodeAddr addr) n
-    return $ fmap unwrapNodes $ join reply
+getNodes client cbvar nid addr =
+    sendQ (getNodesSendable cbvar addr) client nid addr
+
+asyncGetNodes :: QR.Client err PacketKind TransactionId NodeInfo Message
+                 -> TVar (HashMap NodeId [NodeInfoCallback])
+                 -> NodeId
+                 -> NodeInfo
+                 -> (Maybe ([NodeInfo], [NodeInfo], Maybe ()) -> IO ())
+                 -> IO ()
+asyncGetNodes client cbvar nid addr go =
+    asyncQ (getNodesSendable cbvar addr) client nid addr go
 
 updateRouting :: Client -> Routing
                         -> (TVar (R.BucketList NodeInfo) -> RoutingTransition NodeInfo -> STM ())
@@ -525,7 +566,7 @@ nodeSearch :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> Search NodeI
 nodeSearch client cbvar = Search
     { searchSpace       = toxSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = Left $ getNodes client cbvar
+    , searchQuery       = Right $ asyncGetNodes client cbvar
     , searchAlpha       = 8
     , searchK           = 16