diff options
author | joe <joe@jerkface.net> | 2017-07-20 01:11:19 -0400 |
---|---|---|
committer | joe <joe@jerkface.net> | 2017-07-20 01:11:19 -0400 |
commit | eab23d20425d43445a9fdab9d604344172f092f7 (patch) | |
tree | 3a2f166bcca841e50b1c5ef4f5af55f6cb55e2be /src/Network/BitTorrent | |
parent | 0641750f1e360a0a04f1e2f035a2ac8ed36919f0 (diff) |
Updated Search module for rewrite.
Diffstat (limited to 'src/Network/BitTorrent')
-rw-r--r-- | src/Network/BitTorrent/DHT/Search.hs | 88 |
1 files changed, 52 insertions, 36 deletions
diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs index 07240755..91a1079b 100644 --- a/src/Network/BitTorrent/DHT/Search.hs +++ b/src/Network/BitTorrent/DHT/Search.hs | |||
@@ -26,31 +26,37 @@ import qualified Data.Wrapper.PSQ as PSQ | |||
26 | ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ, PSQKey) | 26 | ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ, PSQKey) |
27 | import Network.Address hiding (NodeId) | 27 | import Network.Address hiding (NodeId) |
28 | import Network.DatagramServer.Types | 28 | import Network.DatagramServer.Types |
29 | import Data.Bits | 29 | import Network.DHT.Routing as R |
30 | 30 | ||
31 | data IterativeSearch dht u ip r = IterativeSearch | 31 | data IterativeSearch nid addr ni r = IterativeSearch |
32 | { searchTarget :: NodeId dht | 32 | { searchTarget :: nid |
33 | , searchQuery :: NodeInfo dht ip u -> IO ([NodeInfo dht ip u], [r]) | 33 | , searchSpace :: KademliaSpace nid ni |
34 | , searchNodeAddress :: ni -> addr | ||
35 | , searchQuery :: ni -> IO ([ni], [r]) | ||
34 | , searchPendingCount :: TVar Int | 36 | , searchPendingCount :: TVar Int |
35 | , searchQueued :: TVar (MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht))) | 37 | , searchQueued :: TVar (MinMaxPSQ ni nid) |
36 | , searchInformant :: TVar (MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht))) | 38 | , searchInformant :: TVar (MinMaxPSQ ni nid) |
37 | , searchVisited :: TVar (Set (NodeAddr ip)) | 39 | , searchVisited :: TVar (Set addr) |
38 | , searchResults :: TVar (Set r) | 40 | , searchResults :: TVar (Set r) |
39 | } | 41 | } |
40 | 42 | ||
41 | newSearch :: ( Eq ip | 43 | newSearch :: ( Ord addr |
42 | , PSQKey (NodeId dht) | 44 | , PSQKey nid |
43 | , PSQKey (NodeInfo dht ip u) | 45 | , PSQKey ni |
44 | , FiniteBits (NodeId dht) | 46 | ) => |
45 | ) => (NodeInfo dht ip u -> IO ([NodeInfo dht ip u], [r])) | 47 | KademliaSpace nid ni |
46 | -> NodeId dht -> [NodeInfo dht ip u] -> IO (IterativeSearch dht u ip r) | 48 | -> (ni -> addr) |
47 | newSearch qry target ns = atomically $ do | 49 | -> (ni -> IO ([ni], [r])) |
50 | -> nid -> [ni] -> IO (IterativeSearch nid addr ni r) | ||
51 | newSearch space nAddr qry target ns = atomically $ do | ||
48 | c <- newTVar 0 | 52 | c <- newTVar 0 |
49 | q <- newTVar $ MM.fromList $ map (\n -> n :-> distance target (nodeId n)) ns | 53 | q <- newTVar $ MM.fromList |
54 | $ map (\n -> n :-> kademliaXor space target (kademliaLocation space n)) | ||
55 | $ ns | ||
50 | i <- newTVar MM.empty | 56 | i <- newTVar MM.empty |
51 | v <- newTVar Set.empty | 57 | v <- newTVar Set.empty |
52 | r <- newTVar Set.empty | 58 | r <- newTVar Set.empty |
53 | return $ IterativeSearch target qry c q i v r | 59 | return $ IterativeSearch target space nAddr qry c q i v r |
54 | 60 | ||
55 | searchAlpha :: Int | 61 | searchAlpha :: Int |
56 | searchAlpha = 3 | 62 | searchAlpha = 3 |
@@ -58,15 +64,14 @@ searchAlpha = 3 | |||
58 | searchK :: Int | 64 | searchK :: Int |
59 | searchK = 8 | 65 | searchK = 8 |
60 | 66 | ||
61 | sendQuery :: forall a ip dht u. | 67 | sendQuery :: forall addr nid ni r. |
62 | ( Ord a | 68 | ( Ord addr |
63 | , Ord ip | 69 | , Ord r |
64 | , PSQKey (NodeId dht) | 70 | , PSQKey nid |
65 | , PSQKey (NodeInfo dht ip u) | 71 | , PSQKey ni |
66 | , FiniteBits (NodeId dht) | ||
67 | ) => | 72 | ) => |
68 | IterativeSearch dht u ip a | 73 | IterativeSearch nid addr ni r |
69 | -> Binding (NodeInfo dht ip u) (NodeDistance (NodeId dht)) | 74 | -> Binding ni nid |
70 | -> IO () | 75 | -> IO () |
71 | sendQuery IterativeSearch{..} (ni :-> d) = do | 76 | sendQuery IterativeSearch{..} (ni :-> d) = do |
72 | (ns,rs) <- handle (\(SomeException e) -> return ([],[])) | 77 | (ns,rs) <- handle (\(SomeException e) -> return ([],[])) |
@@ -75,21 +80,24 @@ sendQuery IterativeSearch{..} (ni :-> d) = do | |||
75 | modifyTVar searchPendingCount pred | 80 | modifyTVar searchPendingCount pred |
76 | vs <- readTVar searchVisited | 81 | vs <- readTVar searchVisited |
77 | -- We only queue a node if it is not yet visited | 82 | -- We only queue a node if it is not yet visited |
78 | let insertFoundNode :: NodeInfo dht ip u | 83 | let insertFoundNode :: ni |
79 | -> MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht)) | 84 | -> MinMaxPSQ ni nid |
80 | -> MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht)) | 85 | -> MinMaxPSQ ni nid |
81 | insertFoundNode n q | 86 | insertFoundNode n q |
82 | | nodeAddr n `Set.member` vs = q | 87 | | searchNodeAddress n `Set.member` vs |
83 | | otherwise = MM.insertTake searchK n (distance searchTarget $ nodeId n) q | 88 | = q |
89 | | otherwise = MM.insertTake searchK n ( kademliaXor searchSpace searchTarget | ||
90 | $ kademliaLocation searchSpace n ) | ||
91 | q | ||
84 | modifyTVar searchQueued $ \q -> foldr insertFoundNode q ns | 92 | modifyTVar searchQueued $ \q -> foldr insertFoundNode q ns |
85 | modifyTVar searchInformant $ MM.insertTake searchK ni d | 93 | modifyTVar searchInformant $ MM.insertTake searchK ni d |
86 | modifyTVar searchResults $ \s -> foldr Set.insert s rs | 94 | modifyTVar searchResults $ \s -> foldr Set.insert s rs |
87 | 95 | ||
88 | 96 | ||
89 | searchIsFinished :: ( Ord ip | 97 | searchIsFinished :: ( Ord addr |
90 | , PSQKey (NodeId dht) | 98 | , PSQKey nid |
91 | , PSQKey (NodeInfo dht ip u) | 99 | , PSQKey ni |
92 | ) => IterativeSearch dht u ip r -> STM Bool | 100 | ) => IterativeSearch nid addr ni r -> STM Bool |
93 | searchIsFinished IterativeSearch{..} = do | 101 | searchIsFinished IterativeSearch{..} = do |
94 | q <- readTVar searchQueued | 102 | q <- readTVar searchQueued |
95 | cnt <- readTVar searchPendingCount | 103 | cnt <- readTVar searchPendingCount |
@@ -100,9 +108,17 @@ searchIsFinished IterativeSearch{..} = do | |||
100 | && ( PSQ.prio (fromJust $ MM.findMax informants) | 108 | && ( PSQ.prio (fromJust $ MM.findMax informants) |
101 | <= PSQ.prio (fromJust $ MM.findMin q)))) | 109 | <= PSQ.prio (fromJust $ MM.findMin q)))) |
102 | 110 | ||
111 | searchCancel :: IterativeSearch nid addr ni r -> IO () | ||
112 | searchCancel IterativeSearch{..} = atomically $ do | ||
113 | writeTVar searchPendingCount 0 | ||
114 | writeTVar searchQueued MM.empty | ||
115 | |||
103 | search :: | 116 | search :: |
104 | (Ord r, Ord ip, PSQKey (NodeId dht), PSQKey (NodeInfo dht ip u), FiniteBits (NodeId dht)) => | 117 | ( Ord r |
105 | IterativeSearch dht u ip r -> IO () | 118 | , Ord addr |
119 | , PSQKey nid | ||
120 | , PSQKey ni | ||
121 | ) => IterativeSearch nid addr ni r -> IO () | ||
106 | search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do | 122 | search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do |
107 | fix $ \again -> do | 123 | fix $ \again -> do |
108 | join $ atomically $ do | 124 | join $ atomically $ do |
@@ -114,7 +130,7 @@ search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do | |||
114 | | (MM.size informants < searchK) && (cnt > 0 || not (MM.null q)) | 130 | | (MM.size informants < searchK) && (cnt > 0 || not (MM.null q)) |
115 | || (PSQ.prio (fromJust $ MM.findMax informants) > d) | 131 | || (PSQ.prio (fromJust $ MM.findMax informants) > d) |
116 | -> do writeTVar searchQueued q | 132 | -> do writeTVar searchQueued q |
117 | modifyTVar searchVisited $ Set.insert (nodeAddr ni) | 133 | modifyTVar searchVisited $ Set.insert (searchNodeAddress ni) |
118 | modifyTVar searchPendingCount succ | 134 | modifyTVar searchPendingCount succ |
119 | return $ withAsync g (sendQuery s (ni :-> d)) (const again) | 135 | return $ withAsync g (sendQuery s (ni :-> d)) (const again) |
120 | _ -> do check (cnt == 0) | 136 | _ -> do check (cnt == 0) |