Health.hs

-- | Kubernetes liveness and readiness probe endpoints for servant.
--
-- A service mounts 'HealthApi' under @"health"@ on its umbrella route record
-- and supplies two 'ProbeCheck' actions. It then answers @GET \/health\/live@
-- and @GET \/health\/ready@ with HTTP 200 and a 'ProbeStatus' body when the
-- check passes, and HTTP 503 with the same body shape when it fails.
module Servant.Health
  ( -- * Wire contract
    ProbeStatus (..),
    probeJsonOptions,

    -- * Check seam
    ProbeVerdict (..),
    ProbeCheck,

    -- * Routes
    ProbeResponses,
    ProbeResult (..),
    HealthApi (..),

    -- * Server
    healthServer,
    runProbe,
  )
where

import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson (FromJSON (..), Options, ToJSON (..), defaultOptions, genericParseJSON, genericToJSON)
import Data.OpenApi (ToSchema (..), fromAesonOptions, genericDeclareNamedSchema)
import Data.SOP (I (..), NS (S, Z))
import Data.Text (Text)
import Data.Time (UTCTime)
import GHC.Generics (Generic)
import Servant.API (JSON, StdMethod (GET), (:>))
import Servant.API.Generic ((:-))
import Servant.API.MultiVerb (AsUnion (..), MultiVerb, Respond)
import Servant.Server.Generic (AsServerT)

-- | The probe report every fleet service answers with. The key set is a wire
-- contract consumed by dashboards and runbooks: exactly @status@, @check@,
-- @failingSince@, with @failingSince@ null while healthy.
data ProbeStatus = ProbeStatus
  { status :: !Text,
    check :: !Text,
    failingSince :: !(Maybe UTCTime)
  }
  deriving stock (Generic, Eq, Show)

-- | The one Options value driving the codec (and, in a later plan, the
-- schema). Field names are used verbatim.
probeJsonOptions :: Options
probeJsonOptions = defaultOptions

instance ToJSON ProbeStatus where
  toJSON = genericToJSON probeJsonOptions

instance FromJSON ProbeStatus where
  parseJSON = genericParseJSON probeJsonOptions

-- | Derived from the same 'Options' value as the codec above, so the published
-- schema and the bytes on the wire cannot disagree. Non-orphan on purpose: an
-- instance next to its type can never be forgotten or duplicated by a
-- consumer generating a document.
instance ToSchema ProbeStatus where
  declareNamedSchema = genericDeclareNamedSchema (fromAesonOptions probeJsonOptions)

-- | What a check reports. @since@ is when the failure began, which feeds the
-- wire type's @failingSince@.
data ProbeVerdict
  = Healthy
  | Unhealthy {failedCheck :: !Text, since :: !UTCTime}
  deriving stock (Generic, Eq, Show)

-- | The injectable seam: services supply these, this library never runs
-- anything it invented itself.
type ProbeCheck = IO ProbeVerdict

-- | The two alternatives a probe route can answer with.
type ProbeResponses =
  '[ Respond 200 "Probe passed" ProbeStatus,
     Respond 503 "Probe failed" ProbeStatus
   ]

-- | The handler-side sum matching 'ProbeResponses'.
data ProbeResult
  = ProbePassed !ProbeStatus
  | ProbeFailed !ProbeStatus
  deriving stock (Generic, Eq, Show)

-- | Hand-written on purpose. Both alternatives carry the same body type, so
-- this instance is the only place that decides passed=200, failed=503.
instance AsUnion ProbeResponses ProbeResult where
  toUnion = \case
    ProbePassed body -> Z (I body)
    ProbeFailed body -> S (Z (I body))
  fromUnion = \case
    Z (I body) -> ProbePassed body
    S (Z (I body)) -> ProbeFailed body
    S (S impossible) -> case impossible of {}

-- | Mount under @"health"@ on the service's umbrella NamedRoutes record.
data HealthApi mode = HealthApi
  { live :: mode :- "live" :> MultiVerb 'GET '[JSON] ProbeResponses ProbeResult,
    ready :: mode :- "ready" :> MultiVerb 'GET '[JSON] ProbeResponses ProbeResult
  }
  deriving stock (Generic)

-- | Build the probe server from a liveness check and a readiness check, in
-- that order.
--
-- Both checks are run afresh on every request; nothing here caches a verdict.
--
-- Wrap every check you mount in @Servant.Health.Check.safeCheck@. An exception
-- escaping a check becomes a 500 rather than a probe report, which pages an
-- operator with a stack trace instead of telling them which dependency is
-- down. The same module has combinators for sequencing several readiness
-- checks, bounding a liveness check with a timeout, and reporting when a
-- failure began.
healthServer :: (MonadIO m) => ProbeCheck -> ProbeCheck -> HealthApi (AsServerT m)
healthServer liveCheck readyCheck =
  HealthApi
    { live = runProbe liveCheck,
      ready = runProbe readyCheck
    }

-- | Run one check and turn its verdict into the response sum.
runProbe :: (MonadIO m) => ProbeCheck -> m ProbeResult
runProbe checkAction =
  liftIO checkAction >>= \case
    Healthy ->
      pure (ProbePassed (ProbeStatus "ok" "all" Nothing))
    Unhealthy failed began ->
      pure (ProbeFailed (ProbeStatus "failed" failed (Just began)))