Show peers as soon as they're found.

1 files changed, 20 insertions, 4 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 1b02a37a..68074e2b 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -33,6 +33,7 @@ import System.Mem
 import Data.Torrent (InfoHash)
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT
+import Network.BitTorrent.DHT.Search
 import Network.BitTorrent.DHT.Query
 import Network.BitTorrent.DHT.Message (FindNode(..),NodeFound(..),GetPeers(..),GotPeers(..))
 import Network.KRPC.Manager (QueryFailure(..))
@@ -49,6 +50,7 @@ import Data.Time.Clock
 #else
 import Control.Concurrent
 #endif
+import Control.Concurrent.STM
 import System.Environment
 
 mkNodeAddr :: SockAddr -> NodeAddr IPv4
@@ -130,7 +132,10 @@ marshalForClient :: String -> String
 marshalForClient s = show (length s) ++ ":" ++ s
 
 hPutClient :: Handle -> String -> IO ()
-hPutClient h s = hPutStr h (marshalForClient s)
+hPutClient h s = hPutStr h ('.' : marshalForClient s)
+
+hPutClientChunk :: Handle -> String -> IO ()
+hPutClientChunk h s = hPutStr h (' ' : marshalForClient s)
 
 clientSession :: Node IPv4 -> MVar () -> RestrictedSocket -> Int -> Handle -> IO ()
 clientSession st signalQuit sock n h = do
@@ -280,9 +285,20 @@ clientSession st signalQuit sock n h = do
         ("search-peers", s) -> cmd $ do
                     case readEither s of
                         Right ih -> do
-                            ps <- Set.toList <$> isearch ioGetPeers ih
-                            return $ do
-                            hPutClient h $ unlines $ map (show . pPrint) ps
+                            (tid, s) <- isearch ioGetPeers ih
+                            flip fix Set.empty $ \again shown -> do
+                                (chk,fin) <- liftIO . atomically $ do
+                                    r <- (Set.\\ shown) <$> readTVar (searchResults s)
+                                    if not $ Set.null r
+                                        then (,) r <$> searchIsFinished s
+                                        else searchIsFinished s >>= check >> return (Set.empty,True)
+                                let ps = case Set.toList chk of
+                                            [] -> ""
+                                            _  -> unlines $ map (show . pPrint) $ Set.toList chk
+                                if fin then return $ hPutClient h ps
+                                       else do
+                                            liftIO $ hPutClientChunk h ps
+                                            again (shown `Set.union` chk)
                         Left er -> return $ hPutClient h er
 
         _    -> cmd0 $ hPutClient h "error."