More TCP work.

1 files changed, 26 insertions, 36 deletions

diff --git a/src/Network/Kademlia/Search.hs b/src/Network/Kademlia/Search.hs
index d3aaae28..e87a8618 100644
--- a/src/Network/Kademlia/Search.hs
+++ b/src/Network/Kademlia/Search.hs
@@ -32,36 +32,38 @@ data Search nid addr tok ni r = Search
     { searchSpace        :: KademliaSpace nid ni
     , searchNodeAddress  :: ni -> addr
     , searchQuery        :: nid -> ni -> IO (Maybe ([ni], [r], Maybe tok))
+    , searchAlpha        :: Int -- α = 8
+      -- | 'searchK' should be larger than 'searchAlpha'.  How much larger depends on
+      -- how fast the queries are.  For Tox's much slower onion-routed queries, we
+      -- need to ensure that closer non-responding queries don't completely push out
+      -- farther away queries.
+      --
+      -- For BitTorrent, setting them both 8 was not an issue, but that is no longer
+      -- supported because now the number of remembered informants is now the
+      -- difference between these two numbers.  So, if searchK = 16 and searchAlpha =
+      -- 4, then the number of remembered query responses is 12.
+    , searchK            :: Int -- K = 16
     }
 
 data SearchState nid addr tok ni r = SearchState
-    {-
-    { searchParams :: Search nid addr ni r
-
-    , searchTarget       :: nid
-      -- | This action will be performed at least once on each search result.
-      -- It may be invoked multiple times since different nodes may report the
-      -- same result.  If the action returns 'False', the search will be
-      -- aborted, otherwise it will continue until it is decided that we've
-      -- asked the closest K nodes to the target.
-    , searchResult      :: r -> STM Bool
-
-    -}
-
     { -- | The number of pending queries.  Incremented before any query is sent
       -- and decremented when we get a reply.
       searchPendingCount :: TVar Int
-      -- | Nodes scheduled to be queried.
+      -- | Nodes scheduled to be queried (roughly at most K).
     , searchQueued       :: TVar (MinMaxPSQ ni nid)
       -- | The nearest (K - α) nodes that issued a reply.
+      --
+      -- α is the maximum number of simultaneous queries.
     , searchInformant    :: TVar (MinMaxPSQ' ni nid (Maybe tok))
       -- | This tracks already-queried addresses so we avoid bothering them
       -- again. XXX: We could probably keep only the pending queries in this
       -- set.  It also can be a bounded 'MinMaxPSQ', although searchAlpha
       -- should limit the number of outstanding queries.
     , searchVisited      :: TVar (Set addr)
+    , searchSpec         :: Search nid addr tok ni r
     }
 
+
 newSearch :: ( Ord addr
              , PSQKey nid
              , PSQKey ni
@@ -77,7 +79,7 @@ newSearch :: ( Ord addr
              -> nid
              -> [ni]                   -- Initial nodes to query.
              -> STM (SearchState nid addr tok ni r)
-newSearch (Search space nAddr qry) target ns = do
+newSearch s@(Search space nAddr qry _ _) target ns = do
     c <- newTVar 0
     q <- newTVar $ MM.fromList
                  $ map (\n -> n :-> kademliaXor space target (kademliaLocation space n))
@@ -85,7 +87,7 @@ newSearch (Search space nAddr qry) target ns = do
     i <- newTVar MM.empty
     v <- newTVar Set.empty
     return -- (Search space nAddr qry) , r , target
-           ( SearchState c q i v )
+           ( SearchState c q i v s )
 
 -- | Discard a value from a key-priority-value tuple. This is useful for
 -- swaping items from a "MinMaxPSQ'" to a "MinMaxPSQ".
@@ -110,21 +112,6 @@ reset nearestNodes qsearch target st = do
     writeTVar (searchPendingCount st) 0
     return st
 
-searchAlpha :: Int
-searchAlpha = 8
-
--- | 'searchK' should be larger than 'searchAlpha'.  How much larger depends on
--- how fast the queries are.  For Tox's much slower onion-routed queries, we
--- need to ensure that closer non-responding queries don't completely push out
--- farther away queries.
---
--- For BitTorrent, setting them both 8 was not an issue, but that is no longer
--- supported because now the number of remembered informants is now the
--- difference between these two numbers.  So, if searchK = 16 and searchAlpha =
--- 4, then the number of remembered query responses is 12.
-searchK :: Int
-searchK = 16
-
 sendQuery :: forall addr nid tok ni r.
             ( Ord addr
             , PSQKey nid
@@ -159,8 +146,11 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
              | otherwise = MM.insertTake k n ( kademliaXor searchSpace searchTarget
                                                $ kademliaLocation searchSpace n )
                                                    q
+
         qsize0 <- MM.size <$> readTVar searchQueued
-        let qsize = if qsize0 < searchK then searchK else qsize0
+        let qsize = if qsize0 < searchK then searchK else qsize0 -- Allow searchQueued to grow
+                                                                 -- only when there's fewer than
+                                                                 -- K elements.
         modifyTVar searchQueued    $ \q -> foldr (insertFoundNode qsize) q ns
         modifyTVar searchInformant $ MM.insertTake' (searchK - searchAlpha) ni tok d
         flip fix rs $ \loop -> \case
@@ -174,13 +164,13 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
 searchIsFinished :: ( PSQKey nid
                     , PSQKey ni
                     ) => SearchState nid addr tok ni r -> STM Bool
-searchIsFinished SearchState{ ..} = do
+searchIsFinished SearchState{..} = do
     q <- readTVar searchQueued
     cnt <- readTVar searchPendingCount
     informants <- readTVar searchInformant
     return $ cnt == 0
              && ( MM.null q
-                  || ( MM.size informants >= (searchK - searchAlpha)
+                  || ( MM.size informants >= (searchK searchSpec - searchAlpha searchSpec)
                        && ( PSQ.prio (fromJust $ MM.findMax informants)
                             <= PSQ.prio (fromJust $ MM.findMin q))))
 
@@ -197,7 +187,7 @@ search ::
     , Show nid
     ) => Search nid addr tok ni r -> R.BucketList ni -> nid -> (r -> STM Bool) -> IO (SearchState nid addr tok ni r)
 search sch buckets target result = do
-    let ns = R.kclosest (searchSpace sch) searchK target buckets
+    let ns = R.kclosest (searchSpace sch) (searchK sch) target buckets
     st <- atomically $ newSearch sch target ns
     forkIO $ searchLoop sch target result st
     return st
@@ -218,7 +208,7 @@ searchLoop sch@Search{..} target result s@SearchState{..} = do
             found <- MM.minView <$> readTVar searchQueued
             case found of
                 Just (ni :-> d, q)
-                  | -- If there's fewer than /k/ informants and there's any
+                  | -- If there's fewer than /k - α/ informants and there's any
                     -- node we haven't yet got a response from.
                     (MM.size informants < searchK - searchAlpha) && (cnt > 0 || not (MM.null q))
                     -- Or there's no informants yet at all.