1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
|
{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
module Network.Kademlia.Search where
import Control.Concurrent.Tasks
import Control.Concurrent.STM
import Control.Monad
import Data.Function
import Data.Maybe
import qualified Data.Set as Set
;import Data.Set (Set)
import Data.Hashable (Hashable(..)) -- for type sigs
import System.IO.Error
import qualified Data.MinMaxPSQ as MM
;import Data.MinMaxPSQ (MinMaxPSQ, MinMaxPSQ')
import qualified Data.Wrapper.PSQ as PSQ
;import Data.Wrapper.PSQ (pattern (:->), Binding, pattern Binding, Binding', PSQKey)
import Network.Kademlia.Routing as R
#ifdef THREAD_DEBUG
import Control.Concurrent.Lifted.Instrument
#else
import Control.Concurrent.Lifted
import GHC.Conc (labelThread)
#endif
data Search nid addr tok ni r = Search
{ searchSpace :: KademliaSpace nid ni
, searchNodeAddress :: ni -> addr
, searchQuery :: Either (nid -> ni -> IO (Maybe ([ni], [r], Maybe tok)))
(nid -> ni -> (Maybe ([ni],[r],Maybe tok) -> IO ()) -> IO ())
, 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
{ -- | 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 (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
) =>
{-
KademliaSpace nid ni
-> (ni -> addr)
-> (ni -> IO ([ni], [r])) -- the query action.
-> (r -> STM Bool) -- receives search results.
-> nid -- target of search
-}
Search nid addr tok ni r
-> nid
-> [ni] -- Initial nodes to query.
-> STM (SearchState nid addr tok ni r)
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))
$ ns
i <- newTVar MM.empty
v <- newTVar Set.empty
return -- (Search space nAddr qry) , r , target
( 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".
stripValue :: Binding' k p v -> Binding k p
stripValue (Binding ni _ nid) = (ni :-> nid)
-- | Reset a 'SearchState' object to ready it for a repeated search.
reset :: (Ord ni, Ord nid, Hashable ni, Hashable nid) =>
(nid -> STM [ni])
-> Search nid addr1 tok1 ni r1
-> nid
-> SearchState nid addr tok ni r
-> STM (SearchState nid addr tok ni r)
reset nearestNodes qsearch target st = do
searchIsFinished st >>= check -- Wait for a search to finish before resetting.
bktNodes <- map (\ni -> ni :-> kademliaLocation (searchSpace qsearch) ni)
<$> nearestNodes target
priorInformants <- map stripValue . MM.toList <$> readTVar (searchInformant st)
writeTVar (searchQueued st) $ MM.fromList $ priorInformants ++ bktNodes
writeTVar (searchInformant st) MM.empty
writeTVar (searchVisited st) Set.empty
writeTVar (searchPendingCount st) 0
return st
sendAsyncQuery :: forall addr nid tok ni r.
( Ord addr
, PSQKey nid
, PSQKey ni
, Show nid
) =>
Search nid addr tok ni r
-> nid
-> (r -> STM Bool) -- ^ return False to terminate the search.
-> SearchState nid addr tok ni r
-> Binding ni nid
-> TaskGroup
-> IO ()
sendAsyncQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) g =
case searchQuery of
Left blockingQuery ->
forkTask g "searchQuery" $ do
myThreadId >>= flip labelThread ("searchQuery." ++ show searchTarget)
reply <- blockingQuery searchTarget ni `catchIOError` const (return Nothing)
atomically $ do
modifyTVar searchPendingCount pred
maybe (return ()) go reply
Right nonblockingQuery -> do
nonblockingQuery searchTarget ni $ \reply ->
atomically $ do
modifyTVar searchPendingCount pred
maybe (return ()) go reply
where
go (ns,rs,tok) = do
vs <- readTVar searchVisited
-- We only queue a node if it is not yet visited
let insertFoundNode :: Int
-> ni
-> MinMaxPSQ ni nid
-> MinMaxPSQ ni nid
insertFoundNode k n q
| searchNodeAddress n `Set.member` vs
= q
| 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 -- 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
r:rs' -> do
wanting <- searchResult r
if wanting then loop rs'
else searchCancel sch
[] -> return ()
searchIsFinished :: ( PSQKey nid
, PSQKey ni
) => SearchState nid addr tok ni r -> STM Bool
searchIsFinished SearchState{..} = do
q <- readTVar searchQueued
cnt <- readTVar searchPendingCount
informants <- readTVar searchInformant
return $ cnt == 0
&& ( MM.null q
|| ( MM.size informants >= (searchK searchSpec - searchAlpha searchSpec)
&& ( PSQ.prio (fromJust $ MM.findMax informants)
<= PSQ.prio (fromJust $ MM.findMin q))))
searchCancel :: SearchState nid addr tok ni r -> STM ()
searchCancel SearchState{..} = do
writeTVar searchPendingCount 0
writeTVar searchQueued MM.empty
search ::
( Ord r
, Ord addr
, PSQKey nid
, PSQKey ni
, 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 sch) target buckets
st <- atomically $ newSearch sch target ns
forkIO $ searchLoop sch target result st
return st
searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni )
=> Search nid addr tok ni r -- ^ Query and distance methods.
-> nid -- ^ The target we are searching for.
-> (r -> STM Bool) -- ^ Invoked on each result. Return False to quit searching.
-> SearchState nid addr tok ni r -- ^ Search-related state.
-> IO ()
searchLoop sch@Search{..} target result s@SearchState{..} = do
myThreadId >>= flip labelThread ("search."++show target)
withTaskGroup ("search.g."++show target) searchAlpha $ \g -> fix $ \again -> do
join $ atomically $ do
cnt <- readTVar $ searchPendingCount
check (cnt <= 8) -- Only 8 pending queries at a time.
informants <- readTVar searchInformant
found <- MM.minView <$> readTVar searchQueued
case found of
Just (ni :-> d, q)
| -- 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.
|| MM.null informants
-- Or if the closest scheduled node is nearer than the
-- nearest /k/ informants.
|| (d < PSQ.prio (fromJust $ MM.findMax informants))
-> -- Then the search continues, send a query.
do writeTVar searchQueued q
modifyTVar searchVisited $ Set.insert (searchNodeAddress ni)
modifyTVar searchPendingCount succ
return $ do
sendAsyncQuery sch target result s (ni :-> d) g
again
_ -> -- Otherwise, we are finished.
do check (cnt == 0)
return $ return ()
|
