4 files changed, 42 insertions, 16 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index f6559752..a8858abd 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -89,10 +89,10 @@ library
                        Network.BitTorrent.DHT.Token
                        Network.StreamServer
                        Network.SocketLike
-  other-modules:       Paths_bittorrent
-                       Data.Wrapper.PSQ
-                       Data.MinMaxPSQ
                        Network.BitTorrent.DHT.Search
+                       Data.MinMaxPSQ
+                       Data.Wrapper.PSQ
+  other-modules:       Paths_bittorrent
   if !flag(dht-only)
      exposed-modules:  Network.BitTorrent
                        Network.BitTorrent.Client
@@ -377,6 +377,7 @@ executable dhtd
                      , bittorrent
                      , unix
                      , containers
+                     , stm
   if flag(thread-debug)
     build-depends: time
     cpp-options:   -DTHREAD_DEBUG
diff --git a/examples/dht.hs b/examples/dht.hs
index feeee9ff..16b12678 100644
--- a/examples/dht.hs
+++ b/examples/dht.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE NondecreasingIndentation #-}
 import Control.Monad
-import Control.Monad.Fix
+import Data.Function
 import Control.Monad.IO.Class
 import Data.Char
 import Network.Socket as Socket
@@ -34,16 +34,21 @@ hReadInt h = do
 -- and /.../ is the sequence of characters
 --
 -- Note: The first byte after the count is ignored and discarded.
-readResponse :: Handle -> IO String
+readResponse :: Handle -> IO (Char, String)
 readResponse h = do
+    c <- hGetChar h
     n <- hReadInt h
-    sequence $ replicate n (hGetChar h)
+    s <- sequence $ replicate n (hGetChar h)
+    return (c,s)
 
 -- | Send a command to the dhtd daemon and then print the response.
 sendCommand :: Handle -> String -> InputT IO ()
-sendCommand h cmd = do resp <- liftIO $ do hPutStrLn h cmd
-                                           readResponse h
-                       outputStrLn resp
+sendCommand h cmd = do liftIO $ hPutStrLn h cmd
+                       fix $ \again -> do
+                           (c, resp) <- liftIO $ readResponse h
+                           if c /= '.'
+                            then outputStr resp >> again
+                            else outputStrLn resp
 
 -- | Get one line of input and send it to the daemon, then run the
 -- passed continuation if it wasn't "quit".
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."
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 39ef9604..5345f8b1 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -209,13 +209,17 @@ ioFindNode ih = do
 isearch :: (Ord r, Ord ip) =>
             (InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [r])))
             -> InfoHash
-            -> DHT ip (Set r)
+            -> DHT ip (ThreadId, Search.IterativeSearch ip r)
 isearch f ih = do
     qry <- f ih
     ns <- kclosest 8 ih <$> getTable
     liftIO $ do s <- Search.newSearch qry (toNodeId ih) ns
-                Search.search s
-                atomically $ readTVar (Search.searchResults s)
+                a <- fork $ do
+                    tid <- myThreadId
+                    labelThread tid ("search."++show ih)
+                    Search.search s
+                    -- atomically $ readTVar (Search.searchResults s)
+                return (a, s)
 
 
 type Search    ip o = Conduit [NodeInfo ip] (DHT ip) [o ip]