diff options
author | joe <joe@jerkface.net> | 2017-06-29 10:37:07 -0400 |
---|---|---|
committer | joe <joe@jerkface.net> | 2017-06-29 13:00:16 -0400 |
commit | 3195c0877b443e5ccd4d489f03944fc059d4d7aa (patch) | |
tree | 2a05c35a9b43d8f0725c52fc860b30ae191f3871 /src/Network/BitTorrent/DHT/Search.hs | |
parent | 05e70386c2248d87a61a8e8267e0211597f2fa88 (diff) |
WIP: Generalizing DHT monad.
Diffstat (limited to 'src/Network/BitTorrent/DHT/Search.hs')
-rw-r--r-- | src/Network/BitTorrent/DHT/Search.hs | 52 |
1 files changed, 28 insertions, 24 deletions
diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs index f5cd7834..356f6fd9 100644 --- a/src/Network/BitTorrent/DHT/Search.hs +++ b/src/Network/BitTorrent/DHT/Search.hs | |||
@@ -2,6 +2,7 @@ | |||
2 | {-# LANGUAGE PatternSynonyms #-} | 2 | {-# LANGUAGE PatternSynonyms #-} |
3 | {-# LANGUAGE RecordWildCards #-} | 3 | {-# LANGUAGE RecordWildCards #-} |
4 | {-# LANGUAGE ScopedTypeVariables #-} | 4 | {-# LANGUAGE ScopedTypeVariables #-} |
5 | {-# LANGUAGE FlexibleContexts #-} | ||
5 | module Network.BitTorrent.DHT.Search where | 6 | module Network.BitTorrent.DHT.Search where |
6 | 7 | ||
7 | import Control.Concurrent | 8 | import Control.Concurrent |
@@ -25,27 +26,23 @@ import qualified Data.Wrapper.PSQ as PSQ | |||
25 | ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ) | 26 | ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ) |
26 | import Network.Address hiding (NodeId) | 27 | import Network.Address hiding (NodeId) |
27 | import Network.DatagramServer.Types | 28 | import Network.DatagramServer.Types |
28 | #ifdef VERSION_bencoding | 29 | import Data.Bits |
29 | import Network.DatagramServer.Mainline (KMessageOf) | ||
30 | type Ann = () | ||
31 | #else | ||
32 | import Network.DatagramServer.Tox as Tox | ||
33 | type KMessageOf = Tox.Message | ||
34 | type Ann = Bool | ||
35 | #endif | ||
36 | 30 | ||
37 | data IterativeSearch ip r = IterativeSearch | 31 | data IterativeSearch dht u ip r = IterativeSearch |
38 | { searchTarget :: NodeId KMessageOf | 32 | { searchTarget :: NodeId dht |
39 | , searchQuery :: NodeInfo KMessageOf ip Ann -> IO ([NodeInfo KMessageOf ip Ann], [r]) | 33 | , searchQuery :: NodeInfo dht ip u -> IO ([NodeInfo dht ip u], [r]) |
40 | , searchPendingCount :: TVar Int | 34 | , searchPendingCount :: TVar Int |
41 | , searchQueued :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip Ann) (NodeDistance (NodeId KMessageOf))) | 35 | , searchQueued :: TVar (MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht))) |
42 | , searchInformant :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip Ann) (NodeDistance (NodeId KMessageOf))) | 36 | , searchInformant :: TVar (MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht))) |
43 | , searchVisited :: TVar (Set (NodeAddr ip)) | 37 | , searchVisited :: TVar (Set (NodeAddr ip)) |
44 | , searchResults :: TVar (Set r) | 38 | , searchResults :: TVar (Set r) |
45 | } | 39 | } |
46 | 40 | ||
47 | newSearch :: Eq ip => (NodeInfo KMessageOf ip Ann -> IO ([NodeInfo KMessageOf ip Ann], [r])) | 41 | newSearch :: ( Eq ip |
48 | -> NodeId KMessageOf -> [NodeInfo KMessageOf ip Ann] -> IO (IterativeSearch ip r) | 42 | , Ord (NodeId dht) |
43 | , FiniteBits (NodeId dht) | ||
44 | ) => (NodeInfo dht ip u -> IO ([NodeInfo dht ip u], [r])) | ||
45 | -> NodeId dht -> [NodeInfo dht ip u] -> IO (IterativeSearch dht u ip r) | ||
49 | newSearch qry target ns = atomically $ do | 46 | newSearch qry target ns = atomically $ do |
50 | c <- newTVar 0 | 47 | c <- newTVar 0 |
51 | q <- newTVar $ MM.fromList $ map (\n -> n :-> distance target (nodeId n)) ns | 48 | q <- newTVar $ MM.fromList $ map (\n -> n :-> distance target (nodeId n)) ns |
@@ -60,9 +57,14 @@ searchAlpha = 3 | |||
60 | searchK :: Int | 57 | searchK :: Int |
61 | searchK = 8 | 58 | searchK = 8 |
62 | 59 | ||
63 | sendQuery :: forall a ip. (Ord a, Ord ip) => | 60 | sendQuery :: forall a ip dht u. |
64 | IterativeSearch ip a | 61 | ( Ord a |
65 | -> Binding (NodeInfo KMessageOf ip Ann) (NodeDistance (NodeId KMessageOf)) | 62 | , Ord ip |
63 | , Ord (NodeId dht) | ||
64 | , FiniteBits (NodeId dht) | ||
65 | ) => | ||
66 | IterativeSearch dht u ip a | ||
67 | -> Binding (NodeInfo dht ip u) (NodeDistance (NodeId dht)) | ||
66 | -> IO () | 68 | -> IO () |
67 | sendQuery IterativeSearch{..} (ni :-> d) = do | 69 | sendQuery IterativeSearch{..} (ni :-> d) = do |
68 | (ns,rs) <- handle (\(SomeException e) -> return ([],[])) | 70 | (ns,rs) <- handle (\(SomeException e) -> return ([],[])) |
@@ -71,9 +73,9 @@ sendQuery IterativeSearch{..} (ni :-> d) = do | |||
71 | modifyTVar searchPendingCount pred | 73 | modifyTVar searchPendingCount pred |
72 | vs <- readTVar searchVisited | 74 | vs <- readTVar searchVisited |
73 | -- We only queue a node if it is not yet visited | 75 | -- We only queue a node if it is not yet visited |
74 | let insertFoundNode :: NodeInfo KMessageOf ip u | 76 | let insertFoundNode :: NodeInfo dht ip u |
75 | -> MinMaxPSQ (NodeInfo KMessageOf ip u) (NodeDistance (NodeId KMessageOf)) | 77 | -> MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht)) |
76 | -> MinMaxPSQ (NodeInfo KMessageOf ip u) (NodeDistance (NodeId KMessageOf)) | 78 | -> MinMaxPSQ (NodeInfo dht ip u) (NodeDistance (NodeId dht)) |
77 | insertFoundNode n q | 79 | insertFoundNode n q |
78 | | nodeAddr n `Set.member` vs = q | 80 | | nodeAddr n `Set.member` vs = q |
79 | | otherwise = MM.insertTake searchK n (distance searchTarget $ nodeId n) q | 81 | | otherwise = MM.insertTake searchK n (distance searchTarget $ nodeId n) q |
@@ -82,7 +84,9 @@ sendQuery IterativeSearch{..} (ni :-> d) = do | |||
82 | modifyTVar searchResults $ \s -> foldr Set.insert s rs | 84 | modifyTVar searchResults $ \s -> foldr Set.insert s rs |
83 | 85 | ||
84 | 86 | ||
85 | searchIsFinished :: Ord ip => IterativeSearch ip r -> STM Bool | 87 | searchIsFinished :: ( Ord ip |
88 | , Ord (NodeId dht) | ||
89 | ) => IterativeSearch dht u ip r -> STM Bool | ||
86 | searchIsFinished IterativeSearch{..} = do | 90 | searchIsFinished IterativeSearch{..} = do |
87 | q <- readTVar searchQueued | 91 | q <- readTVar searchQueued |
88 | cnt <- readTVar searchPendingCount | 92 | cnt <- readTVar searchPendingCount |
@@ -94,8 +98,8 @@ searchIsFinished IterativeSearch{..} = do | |||
94 | <= PSQ.prio (fromJust $ MM.findMin q)))) | 98 | <= PSQ.prio (fromJust $ MM.findMin q)))) |
95 | 99 | ||
96 | search :: | 100 | search :: |
97 | (Ord r, Ord ip) => | 101 | (Ord r, Ord ip, Ord (NodeId dht), FiniteBits (NodeId dht)) => |
98 | IterativeSearch ip r -> IO () | 102 | IterativeSearch dht u ip r -> IO () |
99 | search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do | 103 | search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do |
100 | fix $ \again -> do | 104 | fix $ \again -> do |
101 | join $ atomically $ do | 105 | join $ atomically $ do |