Policies.hs

-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at http://mozilla.org/MPL/2.0/.

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RankNTypes #-}

module Database.CQL.IO.Cluster.Policies
    ( Policy (..)
    , random
    , roundRobin
    ) where

import Control.Applicative
import Control.Concurrent.STM
import Control.Lens ((^.), view, over, makeLenses)
import Control.Monad
import Data.Map.Strict (Map)
import Data.Maybe (isNothing)
import Data.Word
import Database.CQL.IO.Cluster.Host
import Database.CQL.IO.PrepQuery (ServerPrepQuery)
import Database.CQL.Protocol (Tuple)
import System.Random.MWC
import Prelude

import qualified Data.Map.Strict as Map

-- | A policy defines a load-balancing strategy and generally
-- handles host visibility.
data Policy = Policy
    { setup :: [Host] -> [Host] -> IO ()
      -- ^ Initialise the policy with two sets of hosts. The first
      -- parameter are hosts known to be available, the second are other
      -- nodes. Note that a policy may be re-initialised at any point
      -- through this function.
    , onEvent :: HostEvent -> IO ()
      -- ^ Event handler. Policies will be informed about cluster changes
      -- through this function.
    , select :: forall k a b. Tuple a => Maybe (ServerPrepQuery k a b, a) -> IO (Maybe Host)
      -- ^ Host selection. The driver will ask for a host to use in a query
      -- through this function. A policy which has no available nodes may
      -- return Nothing.
    , current :: IO [Host]
      -- ^ Return all currently alive hosts.
    , acceptable :: Host -> IO Bool
      -- ^ During startup and node discovery, the driver will ask the
      -- policy if a dicovered host should be ignored.
    , priority :: Host -> IO Int
      -- ^ During startup and node discovery, the driver will ask the
      -- policy if a dicovered host should be ignored, and if not, for a
      -- priority value for a control connection to that host. Lower values are
      -- preferred before later values. Non-positive values are ignored.
    , hostCount :: IO Word
      -- ^ During query processing, the driver will ask the policy for
      -- a rough estimate of alive hosts. The number is used to repeatedly
      -- invoke 'select' (with the underlying assumption that the policy
      -- returns mostly different hosts).
    , display :: IO String
      -- ^ Like having an effectful 'Show' instance for this policy.
    }

type HostMap = TVar Hosts

data Hosts = Hosts
    { _alive :: !(Map InetAddr Host)
    , _other :: !(Map InetAddr Host)
    } deriving Show

makeLenses ''Hosts

-- | Iterate over hosts one by one.
roundRobin :: IO Policy
roundRobin = do
    h <- newTVarIO emptyHosts
    c <- newTVarIO 0
    return (defPolicy h)
      { select = pickHost h c }
  where
    pickHost h c _ = atomically $ do
        m <- view alive <$> readTVar h
        if Map.null m then
            return Nothing
        else do
            k <- readTVar c
            writeTVar c $ succ k `mod` Map.size m
            return . Just . snd $ Map.elemAt (k `mod` Map.size m) m

-- | Return hosts in random order.
random :: IO Policy
random = do
    h <- newTVarIO emptyHosts
    g <- createSystemRandom
    return (defPolicy h)
      { select = pickHost h g }
  where
    pickHost h g _ = do
        m <- view alive <$> readTVarIO h
        if Map.null m then
            return Nothing
        else do
            let i = uniformR (0, Map.size m - 1) g
            Just . snd . flip Map.elemAt m <$> i

-----------------------------------------------------------------------------
-- Defaults

defPolicy :: TVar Hosts -> Policy
defPolicy h = Policy
  { setup = defSetup h
  , onEvent = defOnEvent h
  , select = const $ pure Nothing
  , current = defCurrent h
  , acceptable = defAcceptable
  , priority = defPriority
  , hostCount = defHostCount h
  , display = defDisplay h
  }

emptyHosts :: Hosts
emptyHosts = Hosts Map.empty Map.empty

defDisplay :: HostMap -> IO String
defDisplay h = show <$> readTVarIO h

defAcceptable :: Host -> IO Bool
defAcceptable = fmap (/= 0) . defPriority

defPriority :: Host -> IO Int
defPriority = const $ pure 1

defSetup :: HostMap -> [Host] -> [Host] -> IO ()
defSetup r a b = do
    let ha = Map.fromList $ zip (map (view hostAddr) a) a
    let hb = Map.fromList $ zip (map (view hostAddr) b) b
    let hosts = Hosts ha hb
    atomically $ writeTVar r hosts

defHostCount :: HostMap -> IO Word
defHostCount r = fromIntegral . Map.size . view alive <$> readTVarIO r

defCurrent :: HostMap -> IO [Host]
defCurrent r = Map.elems . view alive <$> readTVarIO r

defOnEvent :: HostMap -> HostEvent -> IO ()
defOnEvent r (HostNew h) = atomically $ do
    m <- readTVar r
    when (isNothing $ get (h^.hostAddr) m) $
        writeTVar r (over alive (Map.insert (h^.hostAddr) h) m)
defOnEvent r (HostGone h) = atomically $ do
    m <- readTVar r
    let a = (h^.hostAddr)
    if Map.member a (m^.alive) then
        writeTVar r (over alive (Map.delete a) m)
    else
        writeTVar r (over other (Map.delete a) m)
defOnEvent r (HostUp h) = atomically $ do
    m <- readTVar r
    let a = (h^.hostAddr)
    case get a m of
        Nothing -> return ()
        Just _ -> writeTVar r
            $ over alive (Map.insert a h)
            . over other (Map.delete a)
            $ m
defOnEvent r (HostDown h) = atomically $ do
    m <- readTVar r
    let a = (h^.hostAddr)
    case get a m of
        Nothing -> return ()
        Just _ -> writeTVar r
            $ over other (Map.insert a h)
            . over alive (Map.delete a)
            $ m
defOnEvent r (AddrDown a) = atomically $ do
    m <- readTVar r
    case get a m of
        Nothing -> return ()
        Just h -> writeTVar r
            $ over other (Map.insert a h)
            . over alive (Map.delete a)
            $ m

get :: InetAddr -> Hosts -> Maybe Host
get a m = Map.lookup a (m^.alive) <|> Map.lookup a (m^.other)