Health.hs

-- | Kubernetes-style health checks derived from the metrics snapshot.
--
-- * Liveness — can we read a snapshot within a time budget? (Is the process and
--   collector responsive?)
-- * Readiness — is the store ready to serve: no subscription overflow-stopped, no
--   subscription lagging beyond 'Kiroku.Metrics.Config.readinessMaxLag', and all
--   configured 'DependencyCheck's passing?
-- * Detailed health — the readiness verdict plus the full snapshot, for humans.
--
-- The built-in 'postgresPing' dependency check issues @SELECT 1@ through the
-- store's pool. Other dependency checks are arbitrary @IO 'DependencyStatus'@
-- actions the caller supplies.
--
-- The lag signal is an /upper bound/: the collector observes a subscription's
-- position only at lifecycle callback points (see the collector's limitation), so
-- readiness may briefly report a caught-up subscription as lagging until its next
-- lifecycle event.
module Kiroku.Metrics.Health
  ( DependencyCheck,
    DependencyStatus (..),
    LivenessStatus (..),
    ReadinessStatus (..),
    checkLiveness,
    checkReadiness,
    checkDetailedHealth,
    postgresPing,
  )
where

import Data.Aeson (ToJSON (..), object, (.=))
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Clock (getMonotonicTimeNSec)
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import System.Timeout (timeout)

import Kiroku.Metrics.Collector (KirokuMetrics, snapshotMetrics)
import Kiroku.Metrics.Config (MetricsServerConfig (..))
import Kiroku.Metrics.Types
  ( MetricsSnapshot (..),
    SubscriptionMetrics (..),
  )
import Kiroku.Store (KirokuStore (..))

-- | A dependency check is an IO action returning that dependency's status.
type DependencyCheck = IO DependencyStatus

-- | Status of one external dependency.
data DependencyStatus = DependencyStatus
  { name :: !Text,
    healthy :: !Bool,
    latencyMs :: !(Maybe Int),
    errorMsg :: !(Maybe Text)
  }
  deriving stock (Eq, Show)

instance ToJSON DependencyStatus where
  toJSON ds =
    object
      [ "name" .= ds.name,
        "healthy" .= ds.healthy,
        "latency_ms" .= ds.latencyMs,
        "error" .= ds.errorMsg
      ]

-- | Liveness verdict.
newtype LivenessStatus = LivenessStatus {alive :: Bool}
  deriving stock (Eq, Show)

instance ToJSON LivenessStatus where
  toJSON s = object ["alive" .= s.alive]

-- | Readiness verdict and its components.
data ReadinessStatus = ReadinessStatus
  { ready :: !Bool,
    lagOk :: !Bool,
    noOverflow :: !Bool,
    dependencies :: ![DependencyStatus]
  }
  deriving stock (Eq, Show)

instance ToJSON ReadinessStatus where
  toJSON s =
    object
      [ "ready" .= s.ready,
        "lag_ok" .= s.lagOk,
        "no_overflow" .= s.noOverflow,
        "dependencies" .= s.dependencies
      ]

-- | Liveness: can we take a snapshot within 'livenessTimeoutUs'? Proves the
-- collector and store @TVar@s are reachable within the budget.
checkLiveness :: MetricsServerConfig -> KirokuMetrics -> IO LivenessStatus
checkLiveness cfg m = do
  result <- timeout cfg.livenessTimeoutUs (snapshotMetrics m)
  pure (LivenessStatus {alive = isJust result})

-- | Readiness: no subscription overflow-stopped, none lagging beyond
-- 'readinessMaxLag', and all dependency checks healthy.
checkReadiness :: MetricsServerConfig -> KirokuMetrics -> [DependencyCheck] -> IO ReadinessStatus
checkReadiness cfg m depChecks = do
  snap <- snapshotMetrics m
  deps <- sequence depChecks
  let subs = Map.elems snap.subscriptions
      noOverflow = not (any ((== Just "overflow") . (.lastStopReason)) subs)
      lagOk = all ((<= cfg.readinessMaxLag) . (.lag)) subs
      depsOk = all (.healthy) deps
  pure
    ReadinessStatus
      { ready = noOverflow && lagOk && depsOk,
        lagOk,
        noOverflow,
        dependencies = deps
      }

-- | Detailed health: the readiness verdict plus the full snapshot.
checkDetailedHealth ::
  MetricsServerConfig ->
  KirokuMetrics ->
  [DependencyCheck] ->
  IO (ReadinessStatus, MetricsSnapshot)
checkDetailedHealth cfg m depChecks = do
  readiness <- checkReadiness cfg m depChecks
  snap <- snapshotMetrics m
  pure (readiness, snap)

-- | Built-in dependency check: time a @SELECT 1@ through the store's pool. A
-- @Right@ result is healthy; a @Left@ 'Pool.UsageError' is unhealthy with the
-- error rendered into 'errorMsg'.
postgresPing :: KirokuStore -> DependencyCheck
postgresPing store = do
  t0 <- getMonotonicTimeNSec
  result <- Pool.use store.pool (Session.script "SELECT 1")
  t1 <- getMonotonicTimeNSec
  let ms = fromIntegral ((t1 - t0) `div` 1_000_000)
  pure $ case result of
    Right () -> DependencyStatus "postgres" True (Just ms) Nothing
    Left e -> DependencyStatus "postgres" False (Just ms) (Just (T.pack (show e)))