Fix query timeouts

2 files changed, 40 insertions, 13 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index 5517e20f..a098d339 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -91,7 +91,8 @@ dht = runDHT handlers
 {-# INLINE dht #-}
 
 -- | One good node may be sufficient. The list of bootstrapping nodes
--- usually obtained from 'Data.Torrent.tNodes' field.
+-- usually obtained from 'Data.Torrent.tNodes' field. Bootstrapping
+-- process can take up to 5 minutes.
 --
 -- (TODO) This operation is asynchronous and do not block.
 --
@@ -109,12 +110,12 @@ bootstrap startNodes = do
           $(logWarnS) "bootstrap" $ T.pack $ show (e :: IOError)
 
         Right (NodeFound closest) -> do
-          $(logDebug) ("Get a list of closest nodes: " <>
-                       T.pack (PP.render (pretty closest)))
-          forM_ (L.take 2 closest) $ \ info @ NodeInfo {..} -> do
+          $(logDebug) $ "Get a list of closest nodes: " <>
+                         T.pack (PP.render (pretty closest))
+          forM_ closest $ \ info @ NodeInfo {..} -> do
             let prettyAddr = T.pack (show (pretty nodeAddr))
             $(logInfoS) "bootstrap" $ "table detalization" <> prettyAddr
-            fork $ insertClosest nodeAddr
+            insertClosest nodeAddr
 
 -- | Get list of peers which downloading this torrent.
 --
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index a6c56d70..55e66a7e 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -40,17 +40,25 @@ module Network.BitTorrent.DHT.Session
          -- ** Accept
        , NodeHandler
        , nodeHandler
+
+         -- ** Iterate
+       , Iteration
+       , Search
+       , search
        ) where
 
+import Prelude hiding (ioError)
+
 import Control.Applicative
 import Control.Concurrent.STM
-import Control.Concurrent.Lifted
+import Control.Concurrent.Lifted hiding (yield)
 import Control.Exception.Lifted hiding (Handler)
 import Control.Monad.Base
 import Control.Monad.Logger
 import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
+import Data.Conduit
 import Data.Default
 import Data.Fixed
 import Data.Hashable
@@ -107,10 +115,10 @@ instance Default Options where
     , optTimeout = 5 -- seconds
     }
 
-milliseconds :: NominalDiffTime -> Int
-milliseconds dt = fromEnum millis
+microseconds :: NominalDiffTime -> Int
+microseconds dt = fromEnum millis
   where
-    millis = realToFrac dt :: Milli
+    millis = realToFrac dt :: Micro
 
 {-----------------------------------------------------------------------
 -- Tokens policy
@@ -325,13 +333,15 @@ queryNode addr q = do
   let signature = T.decodeUtf8 name <> " @ " <> T.pack (render (pretty addr))
   $(logDebugS) "queryNode" $ "Query sent | " <> signature
 
-  interval <- asks (milliseconds . optTimeout . options)
-  result   <- timeout interval  $ query (toSockAddr addr) (Query nid q)
+  interval <- asks (optTimeout . options)
+  result   <- timeout (microseconds interval) $ do
+                query (toSockAddr addr) (Query nid q)
   case result of
     Nothing -> do
       $(logWarnS) "queryNode" $ "not responding @ "
-                   <> T.pack (show (pretty  addr))
-      throwIO $ KError GenericError "timeout expired" ""
+                   <> T.pack (show (pretty  addr)) <> " for "
+                   <> T.pack (show interval)
+      ioError $ userError "timeout expired"
 
     Just (Response remoteId r) -> do
       $(logDebugS) "queryNode" $ "Query recv | " <> signature
@@ -354,3 +364,19 @@ nodeHandler action = handler $ \ sockAddr (Query remoteId q) -> do
     Just naddr -> do
       insertNode (NodeInfo remoteId naddr)
       Response <$> getNodeId <*> action naddr q
+
+{-----------------------------------------------------------------------
+--  Search
+-----------------------------------------------------------------------}
+
+type Iteration ip i o = i ip -> DHT ip (Either [i ip] [o ip])
+type Search    ip i o = Conduit [i ip] (DHT ip) [o ip]
+
+-- TODO: use all inputs
+search :: Address ip => Iteration ip i o -> Search ip i o
+search action = do
+  alpha <- lift $ asks (optAlpha . options)
+  awaitForever $ \ inputs -> do
+    forM_ (L.take alpha inputs) $ \ input -> do
+      result <- lift $ action input
+      either leftover yield result