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
|
{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
module Network.BitTorrent.DHT.Search where
import Control.Concurrent
import Control.Concurrent.Async.Pool
import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.Bool
import Data.Function
import Data.List
import qualified Data.Map.Strict as Map
;import Data.Map.Strict (Map)
import Data.Maybe
import qualified Data.Set as Set
;import Data.Set (Set)
import System.IO
import qualified Data.MinMaxPSQ as MM
;import Data.MinMaxPSQ (MinMaxPSQ)
import qualified Data.Wrapper.PSQ as PSQ
;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ, PSQKey)
import Network.Address hiding (NodeId)
import Network.DatagramServer.Types
import Network.DHT.Routing as R
data Search nid addr ni r = Search
{ searchSpace :: KademliaSpace nid ni
, searchNodeAddress :: ni -> addr
, searchQuery :: ni -> IO ([ni], [r])
}
data SearchState nid addr 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.
, searchQueued :: TVar (MinMaxPSQ ni nid)
-- | The nearest K nodes that issued a reply.
, searchInformant :: TVar (MinMaxPSQ ni nid)
-- | 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)
}
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 ni r
-> nid
-> [ni] -- Initial nodes to query.
-> IO (SearchState nid addr ni r)
newSearch (Search space nAddr qry) target ns = atomically $ 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 )
searchAlpha :: Int
searchAlpha = 3
searchK :: Int
searchK = 8
sendQuery :: forall addr nid ni r.
( Ord addr
, Ord r
, PSQKey nid
, PSQKey ni
) =>
Search nid addr ni r
-> nid
-> (r -> STM Bool)
-> SearchState nid addr ni r
-> Binding ni nid
-> IO ()
sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) = do
(ns,rs) <- handle (\(SomeException e) -> return ([],[]))
(searchQuery ni)
atomically $ do
modifyTVar searchPendingCount pred
vs <- readTVar searchVisited
-- We only queue a node if it is not yet visited
let insertFoundNode :: ni
-> MinMaxPSQ ni nid
-> MinMaxPSQ ni nid
insertFoundNode n q
| searchNodeAddress n `Set.member` vs
= q
| otherwise = MM.insertTake searchK n ( kademliaXor searchSpace searchTarget
$ kademliaLocation searchSpace n )
q
modifyTVar searchQueued $ \q -> foldr insertFoundNode q ns
modifyTVar searchInformant $ MM.insertTake searchK ni d
flip fix rs $ \loop -> \case
r:rs' -> do
wanting <- searchResult r
if wanting then loop rs'
else searchCancel sch
[] -> return ()
searchIsFinished :: ( Ord addr
, PSQKey nid
, PSQKey ni
) => SearchState nid addr 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
&& ( PSQ.prio (fromJust $ MM.findMax informants)
<= PSQ.prio (fromJust $ MM.findMin q))))
searchCancel :: SearchState nid addr ni r -> STM ()
searchCancel SearchState{..} = do
writeTVar searchPendingCount 0
writeTVar searchQueued MM.empty
search ::
( Ord r
, Ord addr
, PSQKey nid
, PSQKey ni
) => Search nid addr ni r -> R.BucketList ni -> nid -> (r -> STM Bool) -> IO ()
search sch@Search{..} buckets target result = withTaskGroup searchAlpha $ \g -> do
let ns = R.kclosest searchSpace searchK target buckets
s@SearchState{..} <- newSearch sch target ns
fix $ \again -> do
join $ atomically $ do
cnt <- readTVar $ searchPendingCount
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) && (cnt > 0 || not (MM.null q))
-- Or if the closest scheduled node is nearer than the
-- nearest /k/ informants.
|| (PSQ.prio (fromJust $ MM.findMax informants) > d)
-> -- Then the search continues, send a query.
do writeTVar searchQueued q
modifyTVar searchVisited $ Set.insert (searchNodeAddress ni)
modifyTVar searchPendingCount succ
return $ withAsync g (sendQuery sch target result s (ni :-> d)) (const again)
_ -> -- Otherwise, we are finished.
do check (cnt == 0)
return $ return ()
|
