{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE StandaloneDeriving #-}

module Prod.Health (
    HealthApi,
    GetReadinessApi,
    handleHealth,
    Liveness (..),
    Reason (..),
    Readiness (..),
    completeReadiness,
    Runtime (..),
    alwaysReadyRuntime,
    withLiveness,
    withReadiness,
    Track (..),
)
where

import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson (FromJSON, ToJSON (..), Value (String))
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Text (Text)
import GHC.Generics (Generic)
import GHC.Stack (CallStack, HasCallStack, callStack)
import Prod.Tracer
import Servant

data Runtime
    = Runtime
    { Runtime -> IO Liveness
liveness :: IO Liveness
    , Runtime -> IO Readiness
readiness :: IO Readiness
    , Runtime -> IORef (Set Reason)
conditions :: IORef (Set Reason)
    , Runtime -> Tracer IO Track
tracer :: Tracer IO Track
    }

alwaysReadyRuntime :: Tracer IO Track -> IO Runtime
alwaysReadyRuntime :: Tracer IO Track -> IO Runtime
alwaysReadyRuntime Tracer IO Track
tracer = IO Liveness
-> IO Readiness -> IORef (Set Reason) -> Tracer IO Track -> Runtime
Runtime (Liveness -> IO Liveness
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Liveness
Alive) (Readiness -> IO Readiness
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Readiness
Ready) (IORef (Set Reason) -> Tracer IO Track -> Runtime)
-> IO (IORef (Set Reason)) -> IO (Tracer IO Track -> Runtime)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Set Reason -> IO (IORef (Set Reason))
forall a. a -> IO (IORef a)
newIORef Set Reason
forall a. Monoid a => a
mempty) IO (Tracer IO Track -> Runtime)
-> IO (Tracer IO Track) -> IO Runtime
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tracer IO Track -> IO (Tracer IO Track)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tracer IO Track
tracer

withReadiness :: IO Readiness -> Runtime -> Runtime
withReadiness :: IO Readiness -> Runtime -> Runtime
withReadiness IO Readiness
io Runtime
rt = Runtime
rt{readiness = io}

withLiveness :: IO Liveness -> Runtime -> Runtime
withLiveness :: IO Liveness -> Runtime -> Runtime
withLiveness IO Liveness
io Runtime
rt = Runtime
rt{liveness = io}

data Liveness = Alive

instance ToJSON Liveness where
    toJSON :: Liveness -> Value
toJSON = Value -> Liveness -> Value
forall a b. a -> b -> a
const (Value -> Liveness -> Value) -> Value -> Liveness -> Value
forall a b. (a -> b) -> a -> b
$ Text -> Value
String Text
"alive"

newtype Reason = Reason Text
    deriving stock (Reason -> Reason -> Bool
(Reason -> Reason -> Bool)
-> (Reason -> Reason -> Bool) -> Eq Reason
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Reason -> Reason -> Bool
== :: Reason -> Reason -> Bool
$c/= :: Reason -> Reason -> Bool
/= :: Reason -> Reason -> Bool
Eq, Eq Reason
Eq Reason =>
(Reason -> Reason -> Ordering)
-> (Reason -> Reason -> Bool)
-> (Reason -> Reason -> Bool)
-> (Reason -> Reason -> Bool)
-> (Reason -> Reason -> Bool)
-> (Reason -> Reason -> Reason)
-> (Reason -> Reason -> Reason)
-> Ord Reason
Reason -> Reason -> Bool
Reason -> Reason -> Ordering
Reason -> Reason -> Reason
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Reason -> Reason -> Ordering
compare :: Reason -> Reason -> Ordering
$c< :: Reason -> Reason -> Bool
< :: Reason -> Reason -> Bool
$c<= :: Reason -> Reason -> Bool
<= :: Reason -> Reason -> Bool
$c> :: Reason -> Reason -> Bool
> :: Reason -> Reason -> Bool
$c>= :: Reason -> Reason -> Bool
>= :: Reason -> Reason -> Bool
$cmax :: Reason -> Reason -> Reason
max :: Reason -> Reason -> Reason
$cmin :: Reason -> Reason -> Reason
min :: Reason -> Reason -> Reason
Ord, Int -> Reason -> ShowS
[Reason] -> ShowS
Reason -> String
(Int -> Reason -> ShowS)
-> (Reason -> String) -> ([Reason] -> ShowS) -> Show Reason
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Reason -> ShowS
showsPrec :: Int -> Reason -> ShowS
$cshow :: Reason -> String
show :: Reason -> String
$cshowList :: [Reason] -> ShowS
showList :: [Reason] -> ShowS
Show)
    deriving
        ([Reason] -> Value
[Reason] -> Encoding
Reason -> Bool
Reason -> Value
Reason -> Encoding
(Reason -> Value)
-> (Reason -> Encoding)
-> ([Reason] -> Value)
-> ([Reason] -> Encoding)
-> (Reason -> Bool)
-> ToJSON Reason
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: Reason -> Value
toJSON :: Reason -> Value
$ctoEncoding :: Reason -> Encoding
toEncoding :: Reason -> Encoding
$ctoJSONList :: [Reason] -> Value
toJSONList :: [Reason] -> Value
$ctoEncodingList :: [Reason] -> Encoding
toEncodingList :: [Reason] -> Encoding
$comitField :: Reason -> Bool
omitField :: Reason -> Bool
ToJSON, Maybe Reason
Value -> Parser [Reason]
Value -> Parser Reason
(Value -> Parser Reason)
-> (Value -> Parser [Reason]) -> Maybe Reason -> FromJSON Reason
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser Reason
parseJSON :: Value -> Parser Reason
$cparseJSONList :: Value -> Parser [Reason]
parseJSONList :: Value -> Parser [Reason]
$comittedField :: Maybe Reason
omittedField :: Maybe Reason
FromJSON)
        via Text

data Readiness = Ready | Ill (Set Reason)
    deriving stock (Readiness -> Readiness -> Bool
(Readiness -> Readiness -> Bool)
-> (Readiness -> Readiness -> Bool) -> Eq Readiness
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Readiness -> Readiness -> Bool
== :: Readiness -> Readiness -> Bool
$c/= :: Readiness -> Readiness -> Bool
/= :: Readiness -> Readiness -> Bool
Eq, Eq Readiness
Eq Readiness =>
(Readiness -> Readiness -> Ordering)
-> (Readiness -> Readiness -> Bool)
-> (Readiness -> Readiness -> Bool)
-> (Readiness -> Readiness -> Bool)
-> (Readiness -> Readiness -> Bool)
-> (Readiness -> Readiness -> Readiness)
-> (Readiness -> Readiness -> Readiness)
-> Ord Readiness
Readiness -> Readiness -> Bool
Readiness -> Readiness -> Ordering
Readiness -> Readiness -> Readiness
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Readiness -> Readiness -> Ordering
compare :: Readiness -> Readiness -> Ordering
$c< :: Readiness -> Readiness -> Bool
< :: Readiness -> Readiness -> Bool
$c<= :: Readiness -> Readiness -> Bool
<= :: Readiness -> Readiness -> Bool
$c> :: Readiness -> Readiness -> Bool
> :: Readiness -> Readiness -> Bool
$c>= :: Readiness -> Readiness -> Bool
>= :: Readiness -> Readiness -> Bool
$cmax :: Readiness -> Readiness -> Readiness
max :: Readiness -> Readiness -> Readiness
$cmin :: Readiness -> Readiness -> Readiness
min :: Readiness -> Readiness -> Readiness
Ord, Int -> Readiness -> ShowS
[Readiness] -> ShowS
Readiness -> String
(Int -> Readiness -> ShowS)
-> (Readiness -> String)
-> ([Readiness] -> ShowS)
-> Show Readiness
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Readiness -> ShowS
showsPrec :: Int -> Readiness -> ShowS
$cshow :: Readiness -> String
show :: Readiness -> String
$cshowList :: [Readiness] -> ShowS
showList :: [Readiness] -> ShowS
Show)
    deriving ((forall x. Readiness -> Rep Readiness x)
-> (forall x. Rep Readiness x -> Readiness) -> Generic Readiness
forall x. Rep Readiness x -> Readiness
forall x. Readiness -> Rep Readiness x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Readiness -> Rep Readiness x
from :: forall x. Readiness -> Rep Readiness x
$cto :: forall x. Rep Readiness x -> Readiness
to :: forall x. Rep Readiness x -> Readiness
Generic)

instance ToJSON Readiness
instance FromJSON Readiness

data Track = Afflict CallStack Reason | Cure CallStack Reason
    deriving (Int -> Track -> ShowS
[Track] -> ShowS
Track -> String
(Int -> Track -> ShowS)
-> (Track -> String) -> ([Track] -> ShowS) -> Show Track
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Track -> ShowS
showsPrec :: Int -> Track -> ShowS
$cshow :: Track -> String
show :: Track -> String
$cshowList :: [Track] -> ShowS
showList :: [Track] -> ShowS
Show)

trace :: (HasCallStack, MonadIO m) => Runtime -> Track -> m ()
trace :: forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
Runtime -> Track -> m ()
trace Runtime
rt = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (Track -> IO ()) -> Track -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tracer IO Track -> Track -> IO ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
runTracer (Runtime -> Tracer IO Track
tracer Runtime
rt)

combineReasons :: Readiness -> Set Reason -> Readiness
combineReasons :: Readiness -> Set Reason -> Readiness
combineReasons Readiness
Ready Set Reason
rs
    | Set Reason -> Bool
forall a. Set a -> Bool
Set.null Set Reason
rs = Readiness
Ready
    | Bool
otherwise = Set Reason -> Readiness
Ill Set Reason
rs
combineReasons (Ill Set Reason
rs1) Set Reason
rs2 = Set Reason -> Readiness
Ill (Set Reason
rs1 Set Reason -> Set Reason -> Set Reason
forall a. Semigroup a => a -> a -> a
<> Set Reason
rs2)

completeReadiness :: Runtime -> IO Readiness
completeReadiness :: Runtime -> IO Readiness
completeReadiness Runtime
rt =
    Readiness -> Set Reason -> Readiness
combineReasons (Readiness -> Set Reason -> Readiness)
-> IO Readiness -> IO (Set Reason -> Readiness)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Runtime -> IO Readiness
readiness Runtime
rt IO (Set Reason -> Readiness) -> IO (Set Reason) -> IO Readiness
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IORef (Set Reason) -> IO (Set Reason)
forall a. IORef a -> IO a
readIORef (Runtime -> IORef (Set Reason)
conditions Runtime
rt)

-- | Add some illness reason.
afflict :: (MonadIO m, HasCallStack) => Runtime -> Reason -> m ()
afflict :: forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Runtime -> Reason -> m ()
afflict Runtime
rt Reason
r = do
    Runtime -> Track -> m ()
forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
Runtime -> Track -> m ()
trace Runtime
rt (Track -> m ()) -> Track -> m ()
forall a b. (a -> b) -> a -> b
$ CallStack -> Reason -> Track
Afflict CallStack
HasCallStack => CallStack
callStack Reason
r
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef (Set Reason) -> (Set Reason -> (Set Reason, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Runtime -> IORef (Set Reason)
conditions Runtime
rt) (\Set Reason
rs -> (Reason -> Set Reason -> Set Reason
forall a. Ord a => a -> Set a -> Set a
Set.insert Reason
r Set Reason
rs, ()))

-- | Remove some illness reason.
cure :: (MonadIO m, HasCallStack) => Runtime -> Reason -> m ()
cure :: forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Runtime -> Reason -> m ()
cure Runtime
rt Reason
r = do
    Runtime -> Track -> m ()
forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
Runtime -> Track -> m ()
trace Runtime
rt (Track -> m ()) -> Track -> m ()
forall a b. (a -> b) -> a -> b
$ CallStack -> Reason -> Track
Cure CallStack
HasCallStack => CallStack
callStack Reason
r
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef (Set Reason) -> (Set Reason -> (Set Reason, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Runtime -> IORef (Set Reason)
conditions Runtime
rt) (\Set Reason
rs -> (Reason -> Set Reason -> Set Reason
forall a. Ord a => a -> Set a -> Set a
Set.delete Reason
r Set Reason
rs, ()))

type GetLivenessApi =
    Summary "Health liveness probe."
        :> "health"
        :> "alive"
        :> Get '[JSON] Liveness

type GetReadinessApi =
    Summary "Health readiness probe."
        :> "health"
        :> "ready"
        :> Get '[JSON] Readiness

type DrainApi =
    Summary "Set a specific 'drained' condition."
        :> "health"
        :> "drain"
        :> Post '[JSON] Readiness

type HealthApi =
    GetLivenessApi
        :<|> GetReadinessApi
        :<|> DrainApi

handleHealth :: Runtime -> Server HealthApi
handleHealth :: Runtime -> Server HealthApi
handleHealth Runtime
runtime =
    Runtime -> Handler Liveness
handleLiveness Runtime
runtime
        Handler Liveness
-> (Handler Readiness :<|> Handler Readiness)
-> Handler Liveness :<|> (Handler Readiness :<|> Handler Readiness)
forall a b. a -> b -> a :<|> b
:<|> Runtime -> Handler Readiness
handleReadiness Runtime
runtime
        Handler Readiness
-> Handler Readiness -> Handler Readiness :<|> Handler Readiness
forall a b. a -> b -> a :<|> b
:<|> Runtime -> Handler Readiness
handleDrain Runtime
runtime
  where
    handleLiveness :: Runtime -> Handler Liveness
    handleLiveness :: Runtime -> Handler Liveness
handleLiveness = IO Liveness -> Handler Liveness
forall a. IO a -> Handler a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Liveness -> Handler Liveness)
-> (Runtime -> IO Liveness) -> Runtime -> Handler Liveness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Runtime -> IO Liveness
liveness
    handleReadiness :: Runtime -> Handler Readiness
    handleReadiness :: Runtime -> Handler Readiness
handleReadiness Runtime
rt = IO Readiness -> Handler Readiness
forall a. IO a -> Handler a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Readiness -> Handler Readiness)
-> IO Readiness -> Handler Readiness
forall a b. (a -> b) -> a -> b
$ do
        Runtime -> IO Readiness
completeReadiness Runtime
rt
    handleDrain :: Runtime -> Handler Readiness
    handleDrain :: Runtime -> Handler Readiness
handleDrain Runtime
rt = do
        Runtime -> Reason -> Handler ()
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Runtime -> Reason -> m ()
afflict Runtime
rt (Text -> Reason
Reason Text
"drained")
        Runtime -> Handler Readiness
handleReadiness Runtime
rt