{-# LANGUAGE DeriveGeneric #-}

module Prod.Healthcheck where

import Control.Concurrent (threadDelay)
import Control.Monad (void, (>=>))
import Data.Aeson (FromJSON (..), ToJSON (..))
import qualified Data.Either as Either
import Data.Foldable (traverse_)
import Data.IORef
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Proxy (Proxy (..))
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Time.Clock
import GHC.Generics (Generic)

import Prod.Background (BackgroundVal)
import qualified Prod.Background as Background
import qualified Prod.Discovery as Discovery
import Prod.Health (GetReadinessApi, Readiness (..))
import Prod.Tracer (Tracer, contramap)
import qualified Prometheus as Prometheus

import Network.HTTP.Client (Manager, defaultManagerSettings, newManager)
import qualified Servant.Client as ServantClient

type Host = Text
type Port = Int
type Error = Text

data Track
    = HealthCheckStarted Host Port
    | HealthCheckFinished Host Port Check
    | BackgroundTrack Host Port (Background.Track CheckSummary)
    deriving (Port -> Track -> ShowS
[Track] -> ShowS
Track -> String
(Port -> Track -> ShowS)
-> (Track -> String) -> ([Track] -> ShowS) -> Show Track
forall a.
(Port -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Port -> Track -> ShowS
showsPrec :: Port -> Track -> ShowS
$cshow :: Track -> String
show :: Track -> String
$cshowList :: [Track] -> ShowS
showList :: [Track] -> ShowS
Show)

data Check
    = Success UTCTime Readiness
    | Failed UTCTime Error
    deriving (Port -> Check -> ShowS
[Check] -> ShowS
Check -> String
(Port -> Check -> ShowS)
-> (Check -> String) -> ([Check] -> ShowS) -> Show Check
forall a.
(Port -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Port -> Check -> ShowS
showsPrec :: Port -> Check -> ShowS
$cshow :: Check -> String
show :: Check -> String
$cshowList :: [Check] -> ShowS
showList :: [Check] -> ShowS
Show, (forall x. Check -> Rep Check x)
-> (forall x. Rep Check x -> Check) -> Generic Check
forall x. Rep Check x -> Check
forall x. Check -> Rep Check x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Check -> Rep Check x
from :: forall x. Check -> Rep Check x
$cto :: forall x. Rep Check x -> Check
to :: forall x. Rep Check x -> Check
Generic)
instance ToJSON Check
instance FromJSON Check

resultTime :: Check -> UTCTime
resultTime :: Check -> UTCTime
resultTime (Success UTCTime
t Readiness
_) = UTCTime
t
resultTime (Failed UTCTime
t Host
_) = UTCTime
t

isSuccess :: Check -> Bool
isSuccess :: Check -> Bool
isSuccess (Success UTCTime
_ Readiness
Ready) = Bool
True
isSuccess Check
_ = Bool
False

getReadiness :: ServantClient.ClientM Readiness
getReadiness :: ClientM Readiness
getReadiness = Proxy GetReadinessApi -> Client ClientM GetReadinessApi
forall api.
HasClient ClientM api =>
Proxy api -> Client ClientM api
ServantClient.client (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @GetReadinessApi)

check :: Manager -> Host -> Port -> IO (Either Error Check)
check :: Manager -> Host -> Port -> IO (Either Host Check)
check Manager
httpManager Host
host Port
port = do
    UTCTime
now <- IO UTCTime
getCurrentTime
    let env :: ClientEnv
env = Manager -> BaseUrl -> ClientEnv
ServantClient.mkClientEnv Manager
httpManager (Scheme -> String -> Port -> String -> BaseUrl
ServantClient.BaseUrl Scheme
ServantClient.Http (Host -> String
Text.unpack Host
host) Port
port String
"")
    Either ClientError Readiness
r <- ClientM Readiness -> ClientEnv -> IO (Either ClientError Readiness)
forall a. ClientM a -> ClientEnv -> IO (Either ClientError a)
ServantClient.runClientM ClientM Readiness
getReadiness ClientEnv
env
    case Either ClientError Readiness
r of
        Left ClientError
err -> Either Host Check -> IO (Either Host Check)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Host Check -> IO (Either Host Check))
-> Either Host Check -> IO (Either Host Check)
forall a b. (a -> b) -> a -> b
$ Host -> Either Host Check
forall a b. a -> Either a b
Left (Host -> Either Host Check) -> Host -> Either Host Check
forall a b. (a -> b) -> a -> b
$ String -> Host
Text.pack (String -> Host) -> String -> Host
forall a b. (a -> b) -> a -> b
$ ClientError -> String
forall a. Show a => a -> String
show ClientError
err
        Right Readiness
v -> Either Host Check -> IO (Either Host Check)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Host Check -> IO (Either Host Check))
-> Either Host Check -> IO (Either Host Check)
forall a b. (a -> b) -> a -> b
$ Check -> Either Host Check
forall a b. b -> Either a b
Right (UTCTime -> Readiness -> Check
Success UTCTime
now Readiness
v)

data CheckSummary
    = CheckSummary
    { CheckSummary -> Maybe Check
lastReady :: Maybe Check
    , CheckSummary -> [Either Host Check]
recentChecks :: [Either Error Check]
    }
    deriving (Port -> CheckSummary -> ShowS
[CheckSummary] -> ShowS
CheckSummary -> String
(Port -> CheckSummary -> ShowS)
-> (CheckSummary -> String)
-> ([CheckSummary] -> ShowS)
-> Show CheckSummary
forall a.
(Port -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Port -> CheckSummary -> ShowS
showsPrec :: Port -> CheckSummary -> ShowS
$cshow :: CheckSummary -> String
show :: CheckSummary -> String
$cshowList :: [CheckSummary] -> ShowS
showList :: [CheckSummary] -> ShowS
Show, (forall x. CheckSummary -> Rep CheckSummary x)
-> (forall x. Rep CheckSummary x -> CheckSummary)
-> Generic CheckSummary
forall x. Rep CheckSummary x -> CheckSummary
forall x. CheckSummary -> Rep CheckSummary x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CheckSummary -> Rep CheckSummary x
from :: forall x. CheckSummary -> Rep CheckSummary x
$cto :: forall x. Rep CheckSummary x -> CheckSummary
to :: forall x. Rep CheckSummary x -> CheckSummary
Generic)
instance ToJSON CheckSummary
instance FromJSON CheckSummary

-- | Predicate to tell if a Summary contains a long-enough check history to be considered.
healthChecked :: CheckSummary -> Bool
healthChecked :: CheckSummary -> Bool
healthChecked CheckSummary
c =
    [Either Host Check] -> Port
forall a. [a] -> Port
forall (t :: * -> *) a. Foldable t => t a -> Port
length (CheckSummary -> [Either Host Check]
recentChecks CheckSummary
c) Port -> Port -> Bool
forall a. Ord a => a -> a -> Bool
>= Port
3

-- | Predicate to tell if a Summary contains no recent successful healthcheck.
neverHealthy :: CheckSummary -> Bool
neverHealthy :: CheckSummary -> Bool
neverHealthy CheckSummary
c =
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
        (Check -> Bool) -> [Check] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Check -> Bool
isSuccess ([Check] -> Bool) -> [Check] -> Bool
forall a b. (a -> b) -> a -> b
$
            [Either Host Check] -> [Check]
forall a b. [Either a b] -> [b]
Either.rights ([Either Host Check] -> [Check]) -> [Either Host Check] -> [Check]
forall a b. (a -> b) -> a -> b
$
                CheckSummary -> [Either Host Check]
recentChecks CheckSummary
c

-- | Predicate to tell if the most recent summary exists and is successful.
recentlyHealthy :: CheckSummary -> Bool
recentlyHealthy :: CheckSummary -> Bool
recentlyHealthy CheckSummary
c =
    Bool -> (Check -> Bool) -> Maybe Check -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Check -> Bool
isSuccess (Maybe Check -> Bool) -> Maybe Check -> Bool
forall a b. (a -> b) -> a -> b
$
        [Check] -> Maybe Check
forall a. [a] -> Maybe a
safeHead ([Check] -> Maybe Check) -> [Check] -> Maybe Check
forall a b. (a -> b) -> a -> b
$
            [Either Host Check] -> [Check]
forall a b. [Either a b] -> [b]
Either.rights ([Either Host Check] -> [Check]) -> [Either Host Check] -> [Check]
forall a b. (a -> b) -> a -> b
$
                CheckSummary -> [Either Host Check]
recentChecks CheckSummary
c

emptyCheckSummary :: CheckSummary
emptyCheckSummary :: CheckSummary
emptyCheckSummary = Maybe Check -> [Either Host Check] -> CheckSummary
CheckSummary Maybe Check
forall a. Maybe a
Nothing []

updateSummary :: Either Error Check -> CheckSummary -> CheckSummary
updateSummary :: Either Host Check -> CheckSummary -> CheckSummary
updateSummary v :: Either Host Check
v@(Right Check
c) CheckSummary
s
    | Check -> Bool
isSuccess Check
c = Maybe Check -> [Either Host Check] -> CheckSummary
CheckSummary (Check -> Maybe Check
forall a. a -> Maybe a
Just Check
c) (Either Host Check
v Either Host Check -> [Either Host Check] -> [Either Host Check]
forall a. a -> [a] -> [a]
: (Port -> [Either Host Check] -> [Either Host Check]
forall a. Port -> [a] -> [a]
take Port
2 (CheckSummary -> [Either Host Check]
recentChecks CheckSummary
s)))
    | Bool
otherwise = Maybe Check -> [Either Host Check] -> CheckSummary
CheckSummary (CheckSummary -> Maybe Check
lastReady CheckSummary
s) (Either Host Check
v Either Host Check -> [Either Host Check] -> [Either Host Check]
forall a. a -> [a] -> [a]
: (Port -> [Either Host Check] -> [Either Host Check]
forall a. Port -> [a] -> [a]
take Port
2 (CheckSummary -> [Either Host Check]
recentChecks CheckSummary
s)))
updateSummary v :: Either Host Check
v@(Left Host
_) CheckSummary
s =
    Maybe Check -> [Either Host Check] -> CheckSummary
CheckSummary (CheckSummary -> Maybe Check
lastReady CheckSummary
s) (Either Host Check
v Either Host Check -> [Either Host Check] -> [Either Host Check]
forall a. a -> [a] -> [a]
: (Port -> [Either Host Check] -> [Either Host Check]
forall a. Port -> [a] -> [a]
take Port
2 (CheckSummary -> [Either Host Check]
recentChecks CheckSummary
s)))

type CheckMap = Map (Host, Port) (BackgroundVal CheckSummary)

emptyCheckMap :: CheckMap
emptyCheckMap :: CheckMap
emptyCheckMap = CheckMap
forall k a. Map k a
Map.empty

initBackgroundCheck ::
    SpaceCounters ->
    Manager ->
    Tracer IO (Background.Track CheckSummary) ->
    (Host, Port) ->
    IO (BackgroundVal CheckSummary)
initBackgroundCheck :: SpaceCounters
-> Manager
-> Tracer IO (Track CheckSummary)
-> (Host, Port)
-> IO (BackgroundVal CheckSummary)
initBackgroundCheck SpaceCounters
cntrs Manager
manager Tracer IO (Track CheckSummary)
tracer (Host
h, Port
p) =
    Tracer IO (Track CheckSummary)
-> CheckSummary
-> CheckSummary
-> (CheckSummary -> IO (CheckSummary, CheckSummary))
-> IO (BackgroundVal CheckSummary)
forall a b.
Tracer IO (Track a)
-> b -> a -> (b -> IO (a, b)) -> IO (BackgroundVal a)
Background.background Tracer IO (Track CheckSummary)
tracer CheckSummary
emptyCheckSummary CheckSummary
emptyCheckSummary CheckSummary -> IO (CheckSummary, CheckSummary)
step
  where
    step :: CheckSummary -> IO (CheckSummary, CheckSummary)
    step :: CheckSummary -> IO (CheckSummary, CheckSummary)
step CheckSummary
st0 = do
        SpaceCounters -> WithSpaceCounter
ns_healthcheck_count SpaceCounters
cntrs WithSpaceCounter -> WithSpaceCounter
forall a b. (a -> b) -> a -> b
$ Counter -> IO ()
forall (m :: * -> *). MonadMonitor m => Counter -> m ()
Prometheus.incCounter
        Either Host Check
res <- Manager -> Host -> Port -> IO (Either Host Check)
check Manager
manager Host
h Port
p
        Port -> IO ()
threadDelay Port
5000000
        let st1 :: CheckSummary
st1 = Either Host Check -> CheckSummary -> CheckSummary
updateSummary Either Host Check
res CheckSummary
st0
        (CheckSummary, CheckSummary) -> IO (CheckSummary, CheckSummary)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CheckSummary
st1, CheckSummary
st1)

terminateBackgroundCheck :: BackgroundVal CheckSummary -> IO ()
terminateBackgroundCheck :: BackgroundVal CheckSummary -> IO ()
terminateBackgroundCheck = BackgroundVal CheckSummary -> IO ()
forall (m :: * -> *) a.
(HasCallStack, MonadIO m) =>
BackgroundVal a -> m ()
Background.kill

data Space
    = Space
    { Space -> Manager
spacehttpManager :: Manager
    , Space -> IORef CheckMap
backgroundChecks :: IORef CheckMap
    , Space -> (Host, Port) -> IO (BackgroundVal CheckSummary)
requestCheck :: (Host, Port) -> IO (BackgroundVal CheckSummary)
    , Space -> (Host, Port) -> IO ()
cancelCheck :: (Host, Port) -> IO ()
    }

clearSpace :: Space -> IO ()
clearSpace :: Space -> IO ()
clearSpace Space
sp = do
    CheckMap
v <- IORef CheckMap -> (CheckMap -> (CheckMap, CheckMap)) -> IO CheckMap
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Space -> IORef CheckMap
backgroundChecks Space
sp) (\CheckMap
old -> (CheckMap
forall k a. Map k a
Map.empty, CheckMap
old))
    (BackgroundVal CheckSummary -> IO ()) -> CheckMap -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ BackgroundVal CheckSummary -> IO ()
terminateBackgroundCheck CheckMap
v

data Counters
    = Counters
    { Counters -> Vector Host Counter
healthcheck_added :: !(Prometheus.Vector Text Prometheus.Counter)
    , Counters -> Vector Host Counter
healthcheck_removed :: !(Prometheus.Vector Text Prometheus.Counter)
    , Counters -> Vector Host Counter
healthcheck_count :: !(Prometheus.Vector Text Prometheus.Counter)
    }

newCounters :: IO Counters
newCounters :: IO Counters
newCounters =
    Vector Host Counter
-> Vector Host Counter -> Vector Host Counter -> Counters
Counters
        (Vector Host Counter
 -> Vector Host Counter -> Vector Host Counter -> Counters)
-> IO (Vector Host Counter)
-> IO (Vector Host Counter -> Vector Host Counter -> Counters)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Host -> IO (Vector Host Counter)
forall {m :: * -> *} {l}.
(MonadIO m, Label l, IsString l) =>
Host -> m (Vector l Counter)
counts Host
"healthcheck_added"
        IO (Vector Host Counter -> Vector Host Counter -> Counters)
-> IO (Vector Host Counter) -> IO (Vector Host Counter -> Counters)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Host -> IO (Vector Host Counter)
forall {m :: * -> *} {l}.
(MonadIO m, Label l, IsString l) =>
Host -> m (Vector l Counter)
counts Host
"healthcheck_removed"
        IO (Vector Host Counter -> Counters)
-> IO (Vector Host Counter) -> IO Counters
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Host -> IO (Vector Host Counter)
forall {m :: * -> *} {l}.
(MonadIO m, Label l, IsString l) =>
Host -> m (Vector l Counter)
counts Host
"healthchecks"
  where
    counts :: Host -> m (Vector l Counter)
counts Host
k =
        Metric (Vector l Counter) -> m (Vector l Counter)
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
Prometheus.register (Metric (Vector l Counter) -> m (Vector l Counter))
-> Metric (Vector l Counter) -> m (Vector l Counter)
forall a b. (a -> b) -> a -> b
$
            l -> Metric Counter -> Metric (Vector l Counter)
forall l m. Label l => l -> Metric m -> Metric (Vector l m)
Prometheus.vector l
"ns" (Metric Counter -> Metric (Vector l Counter))
-> Metric Counter -> Metric (Vector l Counter)
forall a b. (a -> b) -> a -> b
$
                Info -> Metric Counter
Prometheus.counter (Host -> Host -> Info
Prometheus.Info Host
k Host
"")

type WithSpaceCounter = (Prometheus.Counter -> IO ()) -> IO ()

data SpaceCounters
    = SpaceCounters
    { SpaceCounters -> WithSpaceCounter
ns_healthcheck_added :: WithSpaceCounter
    , SpaceCounters -> WithSpaceCounter
ns_healthcheck_removed :: WithSpaceCounter
    , SpaceCounters -> WithSpaceCounter
ns_healthcheck_count :: WithSpaceCounter
    }

namespaceCounters :: Namespace -> Counters -> SpaceCounters
namespaceCounters :: Host -> Counters -> SpaceCounters
namespaceCounters Host
ns Counters
cntrs =
    WithSpaceCounter
-> WithSpaceCounter -> WithSpaceCounter -> SpaceCounters
SpaceCounters
        (Vector Host Counter -> WithSpaceCounter
forall {m :: * -> *} {metric}.
MonadMonitor m =>
Vector Host metric -> (metric -> IO ()) -> m ()
withNamespace (Counters -> Vector Host Counter
healthcheck_added Counters
cntrs))
        (Vector Host Counter -> WithSpaceCounter
forall {m :: * -> *} {metric}.
MonadMonitor m =>
Vector Host metric -> (metric -> IO ()) -> m ()
withNamespace (Counters -> Vector Host Counter
healthcheck_removed Counters
cntrs))
        (Vector Host Counter -> WithSpaceCounter
forall {m :: * -> *} {metric}.
MonadMonitor m =>
Vector Host metric -> (metric -> IO ()) -> m ()
withNamespace (Counters -> Vector Host Counter
healthcheck_count Counters
cntrs))
  where
    withNamespace :: Vector Host metric -> (metric -> IO ()) -> m ()
withNamespace Vector Host metric
v metric -> IO ()
f = Vector Host metric -> Host -> (metric -> IO ()) -> m ()
forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
Prometheus.withLabel Vector Host metric
v Host
ns metric -> IO ()
f

initSpace :: SpaceCounters -> Manager -> Tracer IO Track -> IO Space
initSpace :: SpaceCounters -> Manager -> Tracer IO Track -> IO Space
initSpace SpaceCounters
cntrs Manager
manager Tracer IO Track
tracer = do
    IORef CheckMap
r <- CheckMap -> IO (IORef CheckMap)
forall a. a -> IO (IORef a)
newIORef CheckMap
emptyCheckMap
    Space -> IO Space
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Space -> IO Space) -> Space -> IO Space
forall a b. (a -> b) -> a -> b
$ Manager
-> IORef CheckMap
-> ((Host, Port) -> IO (BackgroundVal CheckSummary))
-> ((Host, Port) -> IO ())
-> Space
Space Manager
manager IORef CheckMap
r (SpaceCounters
-> IORef CheckMap
-> Manager
-> (Host, Port)
-> IO (BackgroundVal CheckSummary)
add SpaceCounters
cntrs IORef CheckMap
r Manager
manager) (SpaceCounters -> IORef CheckMap -> (Host, Port) -> IO ()
del SpaceCounters
cntrs IORef CheckMap
r)
  where
    add :: SpaceCounters -> IORef CheckMap -> Manager -> (Host, Port) -> IO (BackgroundVal CheckSummary)
    add :: SpaceCounters
-> IORef CheckMap
-> Manager
-> (Host, Port)
-> IO (BackgroundVal CheckSummary)
add SpaceCounters
cntrs IORef CheckMap
r Manager
manager = \(Host, Port)
hp -> do
        Maybe (BackgroundVal CheckSummary)
c <- (Host, Port) -> CheckMap -> Maybe (BackgroundVal CheckSummary)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Host, Port)
hp (CheckMap -> Maybe (BackgroundVal CheckSummary))
-> IO CheckMap -> IO (Maybe (BackgroundVal CheckSummary))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef CheckMap -> IO CheckMap
forall a. IORef a -> IO a
readIORef IORef CheckMap
r
        case Maybe (BackgroundVal CheckSummary)
c of
            Maybe (BackgroundVal CheckSummary)
Nothing -> SpaceCounters
-> IORef CheckMap
-> Manager
-> (Host, Port)
-> IO (BackgroundVal CheckSummary)
doadd SpaceCounters
cntrs IORef CheckMap
r Manager
manager (Host, Port)
hp
            Just BackgroundVal CheckSummary
v -> BackgroundVal CheckSummary -> IO (BackgroundVal CheckSummary)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BackgroundVal CheckSummary
v

    doadd :: SpaceCounters -> IORef CheckMap -> Manager -> (Host, Port) -> IO (BackgroundVal CheckSummary)
    doadd :: SpaceCounters
-> IORef CheckMap
-> Manager
-> (Host, Port)
-> IO (BackgroundVal CheckSummary)
doadd SpaceCounters
cntrs IORef CheckMap
r Manager
manager = \hp :: (Host, Port)
hp@(Host
h, Port
p) -> do
        SpaceCounters -> WithSpaceCounter
ns_healthcheck_added SpaceCounters
cntrs WithSpaceCounter -> WithSpaceCounter
forall a b. (a -> b) -> a -> b
$ Counter -> IO ()
forall (m :: * -> *). MonadMonitor m => Counter -> m ()
Prometheus.incCounter
        BackgroundVal CheckSummary
c <- SpaceCounters
-> Manager
-> Tracer IO (Track CheckSummary)
-> (Host, Port)
-> IO (BackgroundVal CheckSummary)
initBackgroundCheck SpaceCounters
cntrs Manager
manager ((Track CheckSummary -> Track)
-> Tracer IO Track -> Tracer IO (Track CheckSummary)
forall a' a. (a' -> a) -> Tracer IO a -> Tracer IO a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (Host -> Port -> Track CheckSummary -> Track
BackgroundTrack Host
h Port
p) Tracer IO Track
tracer) (Host, Port)
hp
        Maybe (BackgroundVal CheckSummary)
concurrentlyAdded <- IORef CheckMap
-> (CheckMap -> (CheckMap, Maybe (BackgroundVal CheckSummary)))
-> IO (Maybe (BackgroundVal CheckSummary))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef CheckMap
r (\CheckMap
st0 -> ((BackgroundVal CheckSummary
 -> BackgroundVal CheckSummary -> BackgroundVal CheckSummary)
-> (Host, Port)
-> BackgroundVal CheckSummary
-> CheckMap
-> CheckMap
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith (\BackgroundVal CheckSummary
_ BackgroundVal CheckSummary
old -> BackgroundVal CheckSummary
old) (Host, Port)
hp BackgroundVal CheckSummary
c CheckMap
st0, (Host, Port) -> CheckMap -> Maybe (BackgroundVal CheckSummary)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Host, Port)
hp CheckMap
st0))
        case Maybe (BackgroundVal CheckSummary)
concurrentlyAdded of
            Maybe (BackgroundVal CheckSummary)
Nothing -> BackgroundVal CheckSummary -> IO (BackgroundVal CheckSummary)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BackgroundVal CheckSummary
c
            Just BackgroundVal CheckSummary
leader -> BackgroundVal CheckSummary -> IO ()
terminateBackgroundCheck BackgroundVal CheckSummary
c IO ()
-> IO (BackgroundVal CheckSummary)
-> IO (BackgroundVal CheckSummary)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> BackgroundVal CheckSummary -> IO (BackgroundVal CheckSummary)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure BackgroundVal CheckSummary
leader

    del :: SpaceCounters -> IORef CheckMap -> (Host, Port) -> IO ()
    del :: SpaceCounters -> IORef CheckMap -> (Host, Port) -> IO ()
del SpaceCounters
cntrs IORef CheckMap
r = \(Host, Port)
hp -> do
        (String, (Host, Port)) -> IO ()
forall a. Show a => a -> IO ()
print (String
"removing", (Host, Port)
hp)
        SpaceCounters -> WithSpaceCounter
ns_healthcheck_removed SpaceCounters
cntrs WithSpaceCounter -> WithSpaceCounter
forall a b. (a -> b) -> a -> b
$ Counter -> IO ()
forall (m :: * -> *). MonadMonitor m => Counter -> m ()
Prometheus.incCounter
        Maybe (BackgroundVal CheckSummary)
c <- IORef CheckMap
-> (CheckMap -> (CheckMap, Maybe (BackgroundVal CheckSummary)))
-> IO (Maybe (BackgroundVal CheckSummary))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef CheckMap
r (\CheckMap
st0 -> ((Host, Port) -> CheckMap -> CheckMap
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (Host, Port)
hp CheckMap
st0, (Host, Port) -> CheckMap -> Maybe (BackgroundVal CheckSummary)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Host, Port)
hp CheckMap
st0))
        case Maybe (BackgroundVal CheckSummary)
c of
            Maybe (BackgroundVal CheckSummary)
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just BackgroundVal CheckSummary
b -> BackgroundVal CheckSummary -> IO ()
terminateBackgroundCheck BackgroundVal CheckSummary
b

setChecks :: Space -> [(Host, Port)] -> IO ()
setChecks :: Space -> [(Host, Port)] -> IO ()
setChecks Space
space [(Host, Port)]
hps = do
    let wantedSet :: Set (Host, Port)
wantedSet = [(Host, Port)] -> Set (Host, Port)
forall a. Ord a => [a] -> Set a
Set.fromList [(Host, Port)]
hps
    Set (Host, Port)
currentSet <- CheckMap -> Set (Host, Port)
forall k a. Map k a -> Set k
Map.keysSet (CheckMap -> Set (Host, Port))
-> IO CheckMap -> IO (Set (Host, Port))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef CheckMap -> IO CheckMap
forall a. IORef a -> IO a
readIORef (Space -> IORef CheckMap
backgroundChecks Space
space)
    let spurious :: Set (Host, Port)
spurious = Set (Host, Port)
currentSet Set (Host, Port) -> Set (Host, Port) -> Set (Host, Port)
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set (Host, Port)
wantedSet
    let missing :: Set (Host, Port)
missing = Set (Host, Port)
wantedSet Set (Host, Port) -> Set (Host, Port) -> Set (Host, Port)
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set (Host, Port)
currentSet
    ((Host, Port) -> IO ()) -> Set (Host, Port) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Space -> (Host, Port) -> IO ()
cancelCheck Space
space) Set (Host, Port)
spurious
    ((Host, Port) -> IO (BackgroundVal CheckSummary))
-> Set (Host, Port) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Space -> (Host, Port) -> IO (BackgroundVal CheckSummary)
requestCheck Space
space) Set (Host, Port)
missing

cancelDeadChecks :: Space -> IO ()
cancelDeadChecks :: Space -> IO ()
cancelDeadChecks Space
space = do
    SummaryMap
summary <- Space -> IO SummaryMap
readBackgroundChecks Space
space
    ((Host, Port) -> IO ()) -> [(Host, Port)] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Space -> (Host, Port) -> IO ()
cancelCheck Space
space) (SummaryMap -> [(Host, Port)]
deadKeys SummaryMap
summary)

{- | Helper to build a Tracer to update hosts to check based on DNS-discovered answers.
Note that the DNSTrack only gives Host, so you need to fmap the port.
-}
setChecksFromDNSDiscovery :: Space -> Discovery.DNSTrack [(Host, Port)] -> IO ()
setChecksFromDNSDiscovery :: Space -> DNSTrack [(Host, Port)] -> IO ()
setChecksFromDNSDiscovery Space
space (Discovery.DNSTrack Host
_ Host
_ (Discovery.BackgroundTrack (Background.RunDone Result [(Host, Port)]
_ Result [(Host, Port)]
newDNSResult))) =
    case Result [(Host, Port)] -> Maybe [(Host, Port)]
forall a. Result a -> Maybe a
Discovery.toMaybe Result [(Host, Port)]
newDNSResult of
        Just [(Host, Port)]
xs -> ((Host, Port) -> IO (BackgroundVal CheckSummary))
-> [(Host, Port)] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Space -> (Host, Port) -> IO (BackgroundVal CheckSummary)
requestCheck Space
space) [(Host, Port)]
xs
        Maybe [(Host, Port)]
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
setChecksFromDNSDiscovery Space
hcrt DNSTrack [(Host, Port)]
_ = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

{- | Same as 'setChecksFromDNSDiscovery' but only adding new checks.
You should clear checks of permanently invalid backends.
-}
addChecksFromDNSDiscovery :: Space -> Discovery.DNSTrack [(Host, Port)] -> IO ()
addChecksFromDNSDiscovery :: Space -> DNSTrack [(Host, Port)] -> IO ()
addChecksFromDNSDiscovery Space
space (Discovery.DNSTrack Host
_ Host
_ (Discovery.BackgroundTrack (Background.RunDone Result [(Host, Port)]
_ Result [(Host, Port)]
newDNSResult))) =
    case Result [(Host, Port)] -> Maybe [(Host, Port)]
forall a. Result a -> Maybe a
Discovery.toMaybe Result [(Host, Port)]
newDNSResult of
        Just [(Host, Port)]
xs -> Space -> [(Host, Port)] -> IO ()
setChecks Space
space [(Host, Port)]
xs
        Maybe [(Host, Port)]
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
addChecksFromDNSDiscovery Space
hcrt DNSTrack [(Host, Port)]
_ = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

type SummaryMap = Map (Host, Port) (CheckSummary)

readCheckMap :: CheckMap -> IO SummaryMap
readCheckMap :: CheckMap -> IO SummaryMap
readCheckMap = (BackgroundVal CheckSummary -> IO CheckSummary)
-> CheckMap -> IO SummaryMap
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map (Host, Port) a -> f (Map (Host, Port) b)
traverse BackgroundVal CheckSummary -> IO CheckSummary
forall (m :: * -> *) a. MonadIO m => BackgroundVal a -> m a
Background.readBackgroundVal

readBackgroundChecks :: Space -> IO SummaryMap
readBackgroundChecks :: Space -> IO SummaryMap
readBackgroundChecks = IORef CheckMap -> IO CheckMap
forall a. IORef a -> IO a
readIORef (IORef CheckMap -> IO CheckMap)
-> (Space -> IORef CheckMap) -> Space -> IO CheckMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Space -> IORef CheckMap
backgroundChecks (Space -> IO CheckMap)
-> (CheckMap -> IO SummaryMap) -> Space -> IO SummaryMap
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> CheckMap -> IO SummaryMap
readCheckMap

{- | Returns the set of (Host,Port) that are healthy in a given SummaryMap.

Healthiness consists in having the latest healthcheck as healthy.
-}
healthyKeys :: SummaryMap -> [(Host, Port)]
healthyKeys :: SummaryMap -> [(Host, Port)]
healthyKeys SummaryMap
m =
    (((Host, Port), CheckSummary) -> (Host, Port))
-> [((Host, Port), CheckSummary)] -> [(Host, Port)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Host, Port), CheckSummary) -> (Host, Port)
forall a b. (a, b) -> a
fst ([((Host, Port), CheckSummary)] -> [(Host, Port)])
-> [((Host, Port), CheckSummary)] -> [(Host, Port)]
forall a b. (a -> b) -> a -> b
$
        (((Host, Port), CheckSummary) -> Bool)
-> [((Host, Port), CheckSummary)] -> [((Host, Port), CheckSummary)]
forall a. (a -> Bool) -> [a] -> [a]
filter (CheckSummary -> Bool
recentlyHealthy (CheckSummary -> Bool)
-> (((Host, Port), CheckSummary) -> CheckSummary)
-> ((Host, Port), CheckSummary)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Host, Port), CheckSummary) -> CheckSummary
forall a b. (a, b) -> b
snd) ([((Host, Port), CheckSummary)] -> [((Host, Port), CheckSummary)])
-> [((Host, Port), CheckSummary)] -> [((Host, Port), CheckSummary)]
forall a b. (a -> b) -> a -> b
$
            SummaryMap -> [((Host, Port), CheckSummary)]
forall k a. Map k a -> [(k, a)]
Map.toList SummaryMap
m

{- | Returns the set of (Host,Port) that have no recent successful activity
provided there is enough health-checking history.
-}
deadKeys :: SummaryMap -> [(Host, Port)]
deadKeys :: SummaryMap -> [(Host, Port)]
deadKeys SummaryMap
m =
    (((Host, Port), CheckSummary) -> (Host, Port))
-> [((Host, Port), CheckSummary)] -> [(Host, Port)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Host, Port), CheckSummary) -> (Host, Port)
forall a b. (a, b) -> a
fst ([((Host, Port), CheckSummary)] -> [(Host, Port)])
-> [((Host, Port), CheckSummary)] -> [(Host, Port)]
forall a b. (a -> b) -> a -> b
$
        (((Host, Port), CheckSummary) -> Bool)
-> [((Host, Port), CheckSummary)] -> [((Host, Port), CheckSummary)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\((Host, Port)
_, CheckSummary
x) -> CheckSummary -> Bool
neverHealthy CheckSummary
x Bool -> Bool -> Bool
&& CheckSummary -> Bool
healthChecked CheckSummary
x) ([((Host, Port), CheckSummary)] -> [((Host, Port), CheckSummary)])
-> [((Host, Port), CheckSummary)] -> [((Host, Port), CheckSummary)]
forall a b. (a -> b) -> a -> b
$
            SummaryMap -> [((Host, Port), CheckSummary)]
forall k a. Map k a -> [(k, a)]
Map.toList SummaryMap
m

safeHead :: [a] -> Maybe a
safeHead :: forall a. [a] -> Maybe a
safeHead (a
x : [a]
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
x
safeHead [a]
_ = Maybe a
forall a. Maybe a
Nothing

type Namespace = Text

type Namespaced a = (Namespace, a)

data Runtime
    = Runtime
    { Runtime -> Counters
counters :: Counters
    , Runtime -> Manager
httpManager :: Manager
    , Runtime -> Tracer IO (Namespaced Track)
tracer :: Tracer IO (Namespaced Track)
    , -- todo: split globals env values and dynamic-space storage
      Runtime -> IORef (Map Host Space)
spaces :: IORef (Map Namespace Space)
    }

initRuntime :: Tracer IO (Namespaced Track) -> IO Runtime
initRuntime :: Tracer IO (Namespaced Track) -> IO Runtime
initRuntime Tracer IO (Namespaced Track)
tracer = do
    IORef (Map Host Space)
r <- Map Host Space -> IO (IORef (Map Host Space))
forall a. a -> IO (IORef a)
newIORef Map Host Space
forall k a. Map k a
Map.empty
    Manager
manager <- ManagerSettings -> IO Manager
newManager ManagerSettings
defaultManagerSettings
    Counters
cntrs <- IO Counters
newCounters
    Runtime -> IO Runtime
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Runtime -> IO Runtime) -> Runtime -> IO Runtime
forall a b. (a -> b) -> a -> b
$ Counters
-> Manager
-> Tracer IO (Namespaced Track)
-> IORef (Map Host Space)
-> Runtime
Runtime Counters
cntrs Manager
manager Tracer IO (Namespaced Track)
tracer IORef (Map Host Space)
r

readSpaces :: Runtime -> IO (Map Namespace SummaryMap)
readSpaces :: Runtime -> IO (Map Host SummaryMap)
readSpaces Runtime
rt = do
    Map Host Space
r <- IORef (Map Host Space) -> IO (Map Host Space)
forall a. IORef a -> IO a
readIORef (IORef (Map Host Space) -> IO (Map Host Space))
-> (Runtime -> IORef (Map Host Space))
-> Runtime
-> IO (Map Host Space)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Runtime -> IORef (Map Host Space)
spaces (Runtime -> IO (Map Host Space)) -> Runtime -> IO (Map Host Space)
forall a b. (a -> b) -> a -> b
$ Runtime
rt
    (Space -> IO SummaryMap)
-> Map Host Space -> IO (Map Host SummaryMap)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Host a -> f (Map Host b)
traverse Space -> IO SummaryMap
readBackgroundChecks Map Host Space
r

registerSpace :: Runtime -> Namespace -> IO Space
registerSpace :: Runtime -> Host -> IO Space
registerSpace Runtime
rt Host
ns = Runtime -> Host -> (Space -> IO Space) -> IO Space
forall a. Runtime -> Host -> (Space -> IO a) -> IO a
withSpace Runtime
rt Host
ns Space -> IO Space
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

withSpace :: Runtime -> Namespace -> (Space -> IO a) -> IO a
withSpace :: forall a. Runtime -> Host -> (Space -> IO a) -> IO a
withSpace Runtime
rt Host
ns Space -> IO a
run = do
    let r :: IORef (Map Host Space)
r = Runtime -> IORef (Map Host Space)
spaces Runtime
rt
    Maybe Space
sp <- Host -> Map Host Space -> Maybe Space
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Host
ns (Map Host Space -> Maybe Space)
-> IO (Map Host Space) -> IO (Maybe Space)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Map Host Space) -> IO (Map Host Space)
forall a. IORef a -> IO a
readIORef IORef (Map Host Space)
r
    case Maybe Space
sp of
        Just Space
s -> Space -> IO a
run Space
s
        Maybe Space
Nothing -> do
            Space
s <- Runtime -> Host -> IO Space
initRuntimeSpace Runtime
rt Host
ns
            Maybe Space
concurrentlyAdded <- IORef (Map Host Space)
-> (Map Host Space -> (Map Host Space, Maybe Space))
-> IO (Maybe Space)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Map Host Space)
r (\Map Host Space
st0 -> ((Space -> Space -> Space)
-> Host -> Space -> Map Host Space -> Map Host Space
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith (\Space
_ Space
old -> Space
old) Host
ns Space
s Map Host Space
st0, Host -> Map Host Space -> Maybe Space
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Host
ns Map Host Space
st0))
            case Maybe Space
concurrentlyAdded of
                Maybe Space
Nothing -> Space -> IO a
run Space
s
                Just Space
leader -> Space -> IO ()
clearSpace Space
s IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Space -> IO a
run Space
leader

-- | Only create a space (no registration).
initRuntimeSpace :: Runtime -> Namespace -> IO Space
initRuntimeSpace :: Runtime -> Host -> IO Space
initRuntimeSpace Runtime
rt Host
ns =
    SpaceCounters -> Manager -> Tracer IO Track -> IO Space
initSpace
        (Host -> Counters -> SpaceCounters
namespaceCounters Host
ns (Counters -> SpaceCounters) -> Counters -> SpaceCounters
forall a b. (a -> b) -> a -> b
$ Runtime -> Counters
counters Runtime
rt)
        (Runtime -> Manager
httpManager Runtime
rt)
        ((Track -> Namespaced Track)
-> Tracer IO (Namespaced Track) -> Tracer IO Track
forall a' a. (a' -> a) -> Tracer IO a -> Tracer IO a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (\Track
x -> (Host
ns, Track
x)) (Runtime -> Tracer IO (Namespaced Track)
tracer Runtime
rt))