summaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/Network/Kademlia/Search.hs45
1 files changed, 37 insertions, 8 deletions
diff --git a/src/Network/Kademlia/Search.hs b/src/Network/Kademlia/Search.hs
index 770d2f13..5f024cd0 100644
--- a/src/Network/Kademlia/Search.hs
+++ b/src/Network/Kademlia/Search.hs
@@ -25,7 +25,7 @@ import System.IO.Error
25import qualified Data.MinMaxPSQ as MM 25import qualified Data.MinMaxPSQ as MM
26 ;import Data.MinMaxPSQ (MinMaxPSQ, MinMaxPSQ') 26 ;import Data.MinMaxPSQ (MinMaxPSQ, MinMaxPSQ')
27import qualified Data.Wrapper.PSQ as PSQ 27import qualified Data.Wrapper.PSQ as PSQ
28 ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ, PSQKey) 28 ;import Data.Wrapper.PSQ (pattern (:->), Binding, pattern Binding, Binding', PSQ, PSQKey)
29import Network.Address hiding (NodeId) 29import Network.Address hiding (NodeId)
30import Network.Kademlia.Routing as R 30import Network.Kademlia.Routing as R
31#ifdef THREAD_DEBUG 31#ifdef THREAD_DEBUG
@@ -94,6 +94,29 @@ newSearch (Search space nAddr qry) target ns = do
94 return -- (Search space nAddr qry) , r , target 94 return -- (Search space nAddr qry) , r , target
95 ( SearchState c q i v ) 95 ( SearchState c q i v )
96 96
97-- | Discard a value from a key-priority-value tuple. This is useful for
98-- swaping items from a "MinMaxPSQ'" to a "MinMaxPSQ".
99stripValue :: Binding' k p v -> Binding k p
100stripValue (Binding ni _ nid) = (ni :-> nid)
101
102-- | Reset a 'SearchState' object to ready it for a repeated search.
103reset :: (Ord ni, Ord nid, Hashable ni, Hashable nid) =>
104 TVar (BucketList ni)
105 -> Search nid addr1 tok1 ni r1
106 -> nid
107 -> SearchState nid addr tok ni r
108 -> STM ()
109reset bkts qsearch target st = do
110 searchIsFinished st >>= check -- Wait for a search to finish before resetting.
111 bktNodes <- map (\ni -> ni :-> kademliaLocation (searchSpace qsearch) ni)
112 . R.kclosest (searchSpace qsearch) searchK target
113 <$> readTVar bkts
114 priorInformants <- map stripValue . MM.toList <$> readTVar (searchInformant st)
115 writeTVar (searchQueued st) $ MM.fromList $ priorInformants ++ bktNodes
116 writeTVar (searchInformant st) MM.empty
117 writeTVar (searchVisited st) Set.empty
118 writeTVar (searchPendingCount st) 0
119
97searchAlpha :: Int 120searchAlpha :: Int
98searchAlpha = 8 121searchAlpha = 8
99 122
@@ -108,7 +131,7 @@ sendQuery :: forall addr nid tok ni r.
108 ) => 131 ) =>
109 Search nid addr tok ni r 132 Search nid addr tok ni r
110 -> nid 133 -> nid
111 -> (r -> STM Bool) 134 -> (r -> STM Bool) -- ^ return False to terminate the search.
112 -> SearchState nid addr tok ni r 135 -> SearchState nid addr tok ni r
113 -> Binding ni nid 136 -> Binding ni nid
114 -> IO () 137 -> IO ()
@@ -124,16 +147,18 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
124 go (ns,rs,tok) = do 147 go (ns,rs,tok) = do
125 vs <- readTVar searchVisited 148 vs <- readTVar searchVisited
126 -- We only queue a node if it is not yet visited 149 -- We only queue a node if it is not yet visited
127 let insertFoundNode :: ni 150 let insertFoundNode :: Int
151 -> ni
128 -> MinMaxPSQ ni nid 152 -> MinMaxPSQ ni nid
129 -> MinMaxPSQ ni nid 153 -> MinMaxPSQ ni nid
130 insertFoundNode n q 154 insertFoundNode k n q
131 | searchNodeAddress n `Set.member` vs 155 | searchNodeAddress n `Set.member` vs
132 = q 156 = q
133 | otherwise = MM.insertTake searchK n ( kademliaXor searchSpace searchTarget 157 | otherwise = MM.insertTake k n ( kademliaXor searchSpace searchTarget
134 $ kademliaLocation searchSpace n ) 158 $ kademliaLocation searchSpace n )
135 q 159 q
136 modifyTVar searchQueued $ \q -> foldr insertFoundNode q ns 160 qsize <- MM.size <$> readTVar searchQueued
161 modifyTVar searchQueued $ \q -> foldr (insertFoundNode qsize) q ns
137 modifyTVar searchInformant $ MM.insertTake' searchK ni tok d 162 modifyTVar searchInformant $ MM.insertTake' searchK ni tok d
138 flip fix rs $ \loop -> \case 163 flip fix rs $ \loop -> \case
139 r:rs' -> do 164 r:rs' -> do
@@ -175,7 +200,11 @@ search sch buckets target result = do
175 return st 200 return st
176 201
177searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni ) 202searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni )
178 => Search nid addr tok ni r -> nid -> (r -> STM Bool) -> SearchState nid addr tok ni r -> IO () 203 => Search nid addr tok ni r -- ^ Query and distance methods.
204 -> nid -- ^ The target we are searching for.
205 -> (r -> STM Bool) -- ^ Invoked on each result. Return False to quit searching.
206 -> SearchState nid addr tok ni r -- ^ Search-related state.
207 -> IO ()
179searchLoop sch@Search{..} target result s@SearchState{..} = do 208searchLoop sch@Search{..} target result s@SearchState{..} = do
180 myThreadId >>= flip labelThread ("search."++show target) 209 myThreadId >>= flip labelThread ("search."++show target)
181 withTaskGroup ("search.g."++show target) searchAlpha $ \g -> fix $ \again -> do 210 withTaskGroup ("search.g."++show target) searchAlpha $ \g -> fix $ \again -> do