{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Network.Legion.Distribution (
ParticipationDefaults,
Peer,
empty,
modify,
findPartition,
rebalanceAction,
RebalanceAction(..),
newPeer,
minimumCompleteServiceSet,
) where
import Prelude hiding (null)
import Data.Aeson (ToJSON, toJSON, object, (.=))
import Data.Binary (Binary)
import Data.Function (on)
import Data.List (sort, sortBy)
import Data.Set (Set, toList)
import Data.Text (pack)
import Data.UUID (UUID)
import GHC.Generics (Generic)
import Network.Legion.KeySet (KeySet, member, (\\), null)
import Network.Legion.LIO (LIO)
import Network.Legion.PartitionKey (PartitionKey)
import Network.Legion.UUID (getUUID)
import Text.Read (readPrec)
import qualified Data.Set as Set
import qualified Network.Legion.KeySet as KS
newtype Peer = Peer UUID deriving (Show, Binary, Eq, Ord)
instance Read Peer where
readPrec = Peer <$> readPrec
newtype ParticipationDefaults = D {
unD :: [(KeySet, Set Peer)]
} deriving (Show, Binary)
instance ToJSON ParticipationDefaults where
toJSON (D dist) = object [
pack (show ks) .= Set.map show peers
| (ks, peers) <- dist
]
empty :: ParticipationDefaults
empty = D []
findPartition :: PartitionKey -> ParticipationDefaults -> Set Peer
findPartition k d =
case [ps | (ks, ps) <- unD d, k `member` ks] of
[ps] -> ps
_ -> error
$ "No exact mach for key in distribution. This means there is a bug in "
++ "the module `Network.Legion.Distribution`. Please report this bug "
++ "via github: " ++ show (k, d)
minimumCompleteServiceSet :: ParticipationDefaults -> Set Peer
minimumCompleteServiceSet defs = Set.fromList [
p
| (_, peers) <- unD defs
, Just (p, _) <- [Set.minView peers]
]
modify
:: (Set Peer -> Set Peer)
-> KeySet
-> ParticipationDefaults
-> ParticipationDefaults
modify fun keyset =
D . filter (not . null . fst) . doModify keyset . unD
where
doModify ks [] = [(ks, fun Set.empty)]
doModify ks ((r, ps):dist) =
let {
unaffected = r \\ ks;
affected = r \\ unaffected;
remaining = ks \\ affected;
} in
(unaffected, ps):(affected, fun ps):doModify remaining dist
rebalanceAction
:: Peer
-> Set Peer
-> ParticipationDefaults
-> Maybe RebalanceAction
rebalanceAction self allPeers (D dist) =
rebuild
where
_rebalance :: a
_rebalance = error "rebalance undefined"
rebuild =
let
underserved = [
(ks, ps)
| (ks, ps) <- dist
, Set.size ps < 3
, not (self `Set.member` ps)
]
mostUnderserved = sortBy (compare `on` Set.size . snd) underserved
in case mostUnderserved of
[] -> Nothing
(ks, ps):_ ->
let
candidateHosts = toList (allPeers Set.\\ ps)
bestHosts = sort [(weightOf p, p) | p <- candidateHosts]
in case bestHosts of
(_, candidate):_ | candidate == self -> Just (Invite ks)
_ -> Nothing
weightOf p = sum [KS.size ks | (ks, ps) <- dist, p `Set.member` ps]
data RebalanceAction
= Invite KeySet
deriving (Show, Generic)
instance Binary RebalanceAction
newPeer :: LIO Peer
newPeer = Peer <$> getUUID