summaryrefslogtreecommitdiff
path: root/src/Network/BitTorrent/DHT/Search.hs
diff options
context:
space:
mode:
authorjoe <joe@jerkface.net>2017-06-29 10:37:07 -0400
committerjoe <joe@jerkface.net>2017-06-29 13:00:16 -0400
commit3195c0877b443e5ccd4d489f03944fc059d4d7aa (patch)
tree2a05c35a9b43d8f0725c52fc860b30ae191f3871 /src/Network/BitTorrent/DHT/Search.hs
parent05e70386c2248d87a61a8e8267e0211597f2fa88 (diff)
WIP: Generalizing DHT monad.
Diffstat (limited to 'src/Network/BitTorrent/DHT/Search.hs')
-rw-r--r--src/Network/BitTorrent/DHT/Search.hs52
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 #-}
5module Network.BitTorrent.DHT.Search where 6module Network.BitTorrent.DHT.Search where
6 7
7import Control.Concurrent 8import 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)
26import Network.Address hiding (NodeId) 27import Network.Address hiding (NodeId)
27import Network.DatagramServer.Types 28import Network.DatagramServer.Types
28#ifdef VERSION_bencoding 29import Data.Bits
29import Network.DatagramServer.Mainline (KMessageOf)
30type Ann = ()
31#else
32import Network.DatagramServer.Tox as Tox
33type KMessageOf = Tox.Message
34type Ann = Bool
35#endif
36 30
37data IterativeSearch ip r = IterativeSearch 31data 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
47newSearch :: Eq ip => (NodeInfo KMessageOf ip Ann -> IO ([NodeInfo KMessageOf ip Ann], [r])) 41newSearch :: ( 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)
49newSearch qry target ns = atomically $ do 46newSearch 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
60searchK :: Int 57searchK :: Int
61searchK = 8 58searchK = 8
62 59
63sendQuery :: forall a ip. (Ord a, Ord ip) => 60sendQuery :: 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 ()
67sendQuery IterativeSearch{..} (ni :-> d) = do 69sendQuery 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
85searchIsFinished :: Ord ip => IterativeSearch ip r -> STM Bool 87searchIsFinished :: ( Ord ip
88 , Ord (NodeId dht)
89 ) => IterativeSearch dht u ip r -> STM Bool
86searchIsFinished IterativeSearch{..} = do 90searchIsFinished 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
96search :: 100search ::
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 ()
99search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do 103search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do
100 fix $ \again -> do 104 fix $ \again -> do
101 join $ atomically $ do 105 join $ atomically $ do