Check.hs

-- | Combinators for building the 'ProbeCheck' actions a service hands to
-- 'Servant.Health.healthServer'.
--
-- This module contains only the generic glue every service otherwise
-- reinvents: running several checks in order, bounding a check with a timeout,
-- turning a thrown exception into a failed probe rather than a 500, and
-- remembering when a failure began. The checks themselves — pinging a
-- database, measuring subscription lag — belong to the service.
module Servant.Health.Check
  ( -- * Building checks
    boolCheck,

    -- * Composing checks
    sequenceChecks,

    -- * Hardening checks
    safeCheck,
    withProbeTimeout,

    -- * Reporting failure onset
    newFailureTracker,
  )
where

import Control.Exception (SomeAsyncException, SomeException, fromException, throwIO, try)
import Data.IORef (atomicModifyIORef', atomicWriteIORef, newIORef)
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Data.Time (getCurrentTime)
import Servant.Health (ProbeCheck, ProbeVerdict (Healthy, Unhealthy))
import System.Timeout (timeout)

-- | Report a failure of the named check, observed now.
unhealthyNow :: Text -> IO ProbeVerdict
unhealthyNow name = Unhealthy name <$> getCurrentTime

-- | Name a boolean test. 'True' is 'Healthy'; 'False' is 'Unhealthy' under the
-- given name, with the failure onset recorded as the moment of observation.
--
-- To report the /original/ onset across repeated probe calls rather than the
-- most recent observation, wrap the result with a tracker from
-- 'newFailureTracker'.
boolCheck :: Text -> IO Bool -> ProbeCheck
boolCheck name test = do
  passed <- test
  if passed
    then pure Healthy
    else unhealthyNow name

-- | Run checks left to right and report the first failure, skipping the rest.
--
-- An empty list is 'Healthy', which makes this a sensible default for a
-- service that has no readiness dependencies yet. Short-circuiting is part of
-- the contract: a readiness check that has already failed should not spend a
-- database round trip confirming the next dependency is also unhappy.
sequenceChecks :: [ProbeCheck] -> ProbeCheck
sequenceChecks = go
  where
    go [] = pure Healthy
    go (action : rest) =
      action >>= \case
        Healthy -> go rest
        failure -> pure failure

-- | Convert a thrown synchronous exception into a failed probe under the given
-- name. Every check a service mounts should be wrapped in this: an exception
-- escaping a probe handler becomes a 500, which Kubernetes reads as a probe
-- failure anyway but which pages an operator with a stack trace instead of a
-- status report.
--
-- Asynchronous exceptions are re-thrown rather than swallowed, so a probe never
-- delays process shutdown. This is also what lets 'withProbeTimeout' work from
-- the outside: 'System.Timeout.timeout' interrupts by throwing an asynchronous
-- exception, and a 'safeCheck' nested inside it must let that through.
--
-- The distinction is the standard one drawn by 'SomeAsyncException': an
-- exception type whose 'Control.Exception.toException' goes through
-- 'Control.Exception.asyncExceptionToException' is asynchronous. A synchronous
-- exception type delivered by 'Control.Concurrent.throwTo' is indistinguishable
-- from an ordinary failure and will be reported as one.
safeCheck :: Text -> ProbeCheck -> ProbeCheck
safeCheck name action = do
  outcome <- try action
  case outcome of
    Right verdict -> pure verdict
    Left err
      | isAsync err -> throwIO err
      | otherwise -> unhealthyNow name
  where
    isAsync :: SomeException -> Bool
    isAsync err = case fromException err of
      Just (_ :: SomeAsyncException) -> True
      Nothing -> False

-- | Bound a check with a timeout, in microseconds. On expiry the check reports
-- 'Unhealthy' under the given name.
--
-- This is what the fleet standard asks for on liveness: a deadlocked process
-- should /fail/ its liveness probe so Kubernetes restarts it, rather than hang
-- the probe and leave the pod wedged until the kubelet's own timeout fires.
--
-- The microsecond unit matches kiroku's @livenessTimeoutUs@ convention.
--
-- Note that 'System.Timeout.timeout' cannot interrupt a foreign call unless the
-- program is built with the threaded runtime. Fleet executables already are —
-- the OpenTelemetry standard requires it — but a test suite that exercises this
-- combinator needs @-threaded@ in its @ghc-options@.
withProbeTimeout :: Int -> Text -> ProbeCheck -> ProbeCheck
withProbeTimeout micros name action =
  timeout micros action >>= \case
    Just verdict -> pure verdict
    Nothing -> unhealthyNow name

-- | Build a wrapper that makes the wire type's @failingSince@ field mean what
-- it says.
--
-- A bare check reports the onset it observed on /this/ call, so a probe that
-- has been failing for an hour reports a fresh onset every time Kubernetes
-- asks. The wrapper returned here remembers the first failure of a consecutive
-- run and rewrites later failures to carry that original onset, so a dashboard
-- reads @failingSince: 12:03:41Z@ rather than @failingSince: now@. A 'Healthy'
-- result clears the memory, so the next failure starts a new run.
--
-- Each call to 'newFailureTracker' creates one independent memory, so liveness
-- and readiness need their own:
--
-- > trackLiveness <- newFailureTracker
-- > trackReadiness <- newFailureTracker
-- > pure (healthServer (trackLiveness livenessCheck) (trackReadiness readinessCheck))
newFailureTracker :: IO (ProbeCheck -> ProbeCheck)
newFailureTracker = do
  onsetRef <- newIORef Nothing
  pure $ \action ->
    action >>= \case
      Healthy -> do
        atomicWriteIORef onsetRef Nothing
        pure Healthy
      Unhealthy failed observed -> do
        onset <- atomicModifyIORef' onsetRef $ \stored ->
          let effective = fromMaybe observed stored
           in (Just effective, effective)
        pure (Unhealthy failed onset)