Distribution.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{- |
  This module defines the data structures and functions used for handling the
  key space distribution.
-}
module Network.Legion.Distribution (
  ParticipationDefaults,
  Peer,
  empty,
  modify,
  findPartition,
  rebalanceAction,
  RebalanceAction(..),
  newPeer,
) 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


{- |
  The way to identify a peer.
-}
newtype Peer = Peer UUID deriving (Show, Binary, Eq, Ord)
instance Read Peer where
  readPrec = Peer <$> readPrec
  


{- |
  The distribution of partitions and partition replicas among the cluster.
-}
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
    ]


{- |
  Constuct a distribution that contains no partitions.
-}
empty :: ParticipationDefaults

empty = D []


{- |
  Find the peers that own the specified partition.
-}
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)


{- |
  Modify the default participation for the key set.
-}
modify
  :: (Set Peer -> Set Peer)
  -> KeySet
  -> ParticipationDefaults
  -> ParticipationDefaults

modify fun keyset =
    {-
      doModify can produce key ranges that contain zero keys, which is
      why the `filter`.
    -}
    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


{- |
  Return the best action, if any, that the indicated peer should take to
  rebalance an unbalanced distribution.
-}
rebalanceAction
  :: Peer
  -> Set Peer
  -> ParticipationDefaults
  -> Maybe RebalanceAction
rebalanceAction self allPeers (D dist) =
    rebuild
    {- TODO rebalance -}
  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
            {- we are the best host -}
            (_, candidate):_ | candidate == self -> Just (Invite ks)
            _ -> Nothing

    weightOf p = sum [KS.size ks | (ks, ps) <- dist, p `Set.member` ps]



{- |
  The actions that are taken in order to build a balanced cluster.
-}
data RebalanceAction
  = Invite KeySet
  deriving (Show, Generic)
instance Binary RebalanceAction


{- |
  Create a new peer.
-}
newPeer :: LIO Peer
newPeer = Peer <$> getUUID


-- {- |
--   Trace helper
-- -}
-- t :: (Show a) => String -> a -> a
-- t msg a = trace (msg ++ ": " ++ show a) a