|
1 | 1 | module DMQ.Diffusion.PeerSelection where |
2 | 2 |
|
3 | | -import Data.Set (Set) |
| 3 | +import Control.Concurrent.Class.MonadSTM.Strict |
| 4 | +import Data.List (sortOn, unfoldr) |
| 5 | +import Data.Map.Strict qualified as Map |
4 | 6 | import Data.Set qualified as Set |
5 | | -import Network.Socket (SockAddr) |
6 | | -import Ouroboros.Network.PeerSelection.Governor.Types |
7 | | -import System.Random (Random (..), StdGen) |
| 7 | +import Data.Word (Word32) |
| 8 | +import Ouroboros.Network.PeerSelection |
| 9 | +import System.Random (Random (..), StdGen, split) |
8 | 10 |
|
9 | 11 | -- | Trivial peer selection policy used as dummy value |
10 | 12 | -- |
11 | | -policy :: StdGen -> PeerSelectionPolicy SockAddr IO |
12 | | -policy gen = |
| 13 | +policy :: forall peerAddr m. |
| 14 | + ( MonadSTM m |
| 15 | + , Ord peerAddr |
| 16 | + ) |
| 17 | + => StrictTVar m StdGen |
| 18 | + -> PeerSelectionPolicy peerAddr m |
| 19 | +policy rngVar = |
13 | 20 | PeerSelectionPolicy { |
14 | | - policyPickKnownPeersForPeerShare = \_ _ _ -> pickTrivially |
15 | | - , policyPickColdPeersToForget = \_ _ _ -> pickTrivially |
16 | | - , policyPickColdPeersToPromote = \_ _ _ -> pickTrivially |
17 | | - , policyPickWarmPeersToPromote = \_ _ _ -> pickTrivially |
18 | | - , policyPickHotPeersToDemote = \_ _ _ -> pickTrivially |
19 | | - , policyPickWarmPeersToDemote = \_ _ _ -> pickTrivially |
20 | | - , policyPickInboundPeers = \_ _ _ -> pickTrivially |
21 | | - , policyFindPublicRootTimeout = 5 |
22 | | - , policyMaxInProgressPeerShareReqs = 0 |
23 | | - , policyPeerShareRetryTime = 0 -- seconds |
24 | | - , policyPeerShareBatchWaitTime = 0 -- seconds |
25 | | - , policyPeerShareOverallTimeout = 0 -- seconds |
26 | | - , policyPeerShareActivationDelay = 2 -- seconds |
| 21 | + policyPickKnownPeersForPeerShare = simplePromotionPolicy, |
| 22 | + policyPickColdPeersToPromote = simplePromotionPolicy, |
| 23 | + policyPickWarmPeersToPromote = simplePromotionPolicy, |
| 24 | + policyPickInboundPeers = simplePromotionPolicy, |
| 25 | + |
| 26 | + policyPickHotPeersToDemote = hotDemotionPolicy, |
| 27 | + policyPickWarmPeersToDemote = warmDemotionPolicy, |
| 28 | + policyPickColdPeersToForget = coldForgetPolicy, |
| 29 | + |
| 30 | + policyFindPublicRootTimeout = 5, |
| 31 | + policyMaxInProgressPeerShareReqs = 0, |
| 32 | + policyPeerShareRetryTime = 0, -- seconds |
| 33 | + policyPeerShareBatchWaitTime = 0, -- seconds |
| 34 | + policyPeerShareOverallTimeout = 0, -- seconds |
| 35 | + policyPeerShareActivationDelay = 2 -- seconds |
27 | 36 | } |
28 | 37 | where |
29 | | - pickTrivially :: Applicative m => Set SockAddr -> Int -> m (Set SockAddr) |
30 | | - pickTrivially set n = pure |
31 | | - . fst |
32 | | - $ go gen (Set.toList set) n [] |
33 | | - where |
34 | | - go g _ 0 acc = (Set.fromList acc, g) |
35 | | - go g [] _ acc = (Set.fromList acc, g) |
36 | | - go g xs k acc = |
37 | | - let (idx, g') = randomR (0, length xs - 1) g |
38 | | - picked = xs !! idx |
39 | | - xs' = take idx xs ++ drop (idx + 1) xs |
40 | | - in go g' xs' (k - 1) (picked : acc) |
| 38 | + hotDemotionPolicy :: PickPolicy peerAddr (STM m) |
| 39 | + hotDemotionPolicy _ _ _ available pickNum = do |
| 40 | + available' <- addRand rngVar available (,) |
| 41 | + return $ Set.fromList |
| 42 | + . map fst |
| 43 | + . take pickNum |
| 44 | + . sortOn snd |
| 45 | + . Map.assocs |
| 46 | + $ available' |
| 47 | + |
| 48 | + -- Randomly pick peers to demote, peers with knownPeerTepid set are twice |
| 49 | + -- as likely to be demoted. |
| 50 | + warmDemotionPolicy :: PickPolicy peerAddr (STM m) |
| 51 | + warmDemotionPolicy _ _ isTepid available pickNum = do |
| 52 | + available' <- addRand rngVar available (tepidWeight isTepid) |
| 53 | + return $ Set.fromList |
| 54 | + . map fst |
| 55 | + . take pickNum |
| 56 | + . sortOn snd |
| 57 | + . Map.assocs |
| 58 | + $ available' |
| 59 | + |
| 60 | + simplePromotionPolicy :: PickPolicy peerAddr (STM m) |
| 61 | + simplePromotionPolicy _ _ _ available pickNum = do |
| 62 | + available' <- addRand rngVar available (,) |
| 63 | + return $ Set.fromList |
| 64 | + . map fst |
| 65 | + . take pickNum |
| 66 | + . sortOn snd |
| 67 | + . Map.assocs |
| 68 | + $ available' |
| 69 | + |
| 70 | + -- Randomly pick peers to forget, peers with failures are more likely to |
| 71 | + -- be forgotten. |
| 72 | + coldForgetPolicy :: PickPolicy peerAddr (STM m) |
| 73 | + coldForgetPolicy _ failCnt _ available pickNum = do |
| 74 | + available' <- addRand rngVar available (failWeight failCnt) |
| 75 | + return $ Set.fromList |
| 76 | + . map fst |
| 77 | + . take pickNum |
| 78 | + . sortOn snd |
| 79 | + . Map.assocs |
| 80 | + $ available' |
| 81 | + |
| 82 | + -- Failures lowers r |
| 83 | + failWeight :: (peerAddr -> Int) |
| 84 | + -> peerAddr |
| 85 | + -> Word32 |
| 86 | + -> (peerAddr, Word32) |
| 87 | + failWeight failCnt peer r = |
| 88 | + (peer, r `div` fromIntegral (failCnt peer + 1)) |
| 89 | + |
| 90 | + -- Tepid flag cuts r in half |
| 91 | + tepidWeight :: (peerAddr -> Bool) |
| 92 | + -> peerAddr |
| 93 | + -> Word32 |
| 94 | + -> (peerAddr, Word32) |
| 95 | + tepidWeight isTepid peer r = |
| 96 | + if isTepid peer then (peer, r `div` 2) |
| 97 | + else (peer, r) |
| 98 | + |
| 99 | + |
| 100 | + -- Add scaled random number in order to prevent ordering based on SockAddr |
| 101 | +addRand :: ( MonadSTM m |
| 102 | + , Ord peerAddr |
| 103 | + ) |
| 104 | + => StrictTVar m StdGen |
| 105 | + -> Set.Set peerAddr |
| 106 | + -> (peerAddr -> Word32 -> (peerAddr, Word32)) |
| 107 | + -> STM m (Map.Map peerAddr Word32) |
| 108 | +addRand rngVar available scaleFn = do |
| 109 | + inRng <- readTVar rngVar |
| 110 | + |
| 111 | + let (rng, rng') = split inRng |
| 112 | + rns = take (Set.size available) $ unfoldr (Just . random) rng :: [Word32] |
| 113 | + available' = Map.fromList $ zipWith scaleFn (Set.toList available) rns |
| 114 | + writeTVar rngVar rng' |
| 115 | + return available' |
| 116 | + |
0 commit comments