-- SPDX-FileCopyrightText: 2026 Alexandra de Wit
--
-- SPDX-License-Identifier: MIT

{- | The mirror worker heartbeat behind @\/livez@, shared by embedded and dedicated workers.
Startup and later stalls use the same allowance. Roles without a worker use 'alwaysLive'.
-}
module Ecluse.Core.Worker.Liveness (
    WorkerHeartbeat,
    newWorkerHeartbeat,
    newWorkerHeartbeatWithClock,
    recordPoll,
    lastPoll,
    workerJobStepAllowance,
    workerHeartbeatStaleAfter,
    heartbeatHealthy,
    Liveness (..),
    alwaysLive,
    heartbeatLivenessNow,
) where

import Data.Time (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime)

-- | Worker progress and its startup allowance, separate from HTTP readiness.
data WorkerHeartbeat = WorkerHeartbeat
    { WorkerHeartbeat -> UTCTime
whStartedAt :: UTCTime
    , WorkerHeartbeat -> IO UTCTime
whNow :: IO UTCTime
    , WorkerHeartbeat -> TVar (Maybe UTCTime)
whLastPoll :: TVar (Maybe UTCTime)
    }

{- | Build a fresh 'WorkerHeartbeat' with no poll yet recorded ('lastPoll' is
'Nothing' until the worker's first successful @receive@).
-}
newWorkerHeartbeat :: IO WorkerHeartbeat
newWorkerHeartbeat :: IO WorkerHeartbeat
newWorkerHeartbeat = IO UTCTime -> IO WorkerHeartbeat
newWorkerHeartbeatWithClock IO UTCTime
getCurrentTime

-- | Start the allowance at the supplied clock, which also drives liveness probes.
newWorkerHeartbeatWithClock :: IO UTCTime -> IO WorkerHeartbeat
newWorkerHeartbeatWithClock :: IO UTCTime -> IO WorkerHeartbeat
newWorkerHeartbeatWithClock IO UTCTime
clock = do
    startedAt <- IO UTCTime
clock
    var <- newTVarIO Nothing
    pure WorkerHeartbeat{whStartedAt = startedAt, whNow = clock, whLastPoll = var}

{- | Stamp the heartbeat with the given instant, recording a unit of worker progress.
The worker calls it through 'Ecluse.Core.Worker.Types.recordWorkerProgress'.
-}
recordPoll :: WorkerHeartbeat -> UTCTime -> IO ()
recordPoll :: WorkerHeartbeat -> UTCTime -> IO ()
recordPoll WorkerHeartbeat
heartbeat UTCTime
now = STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar (Maybe UTCTime) -> Maybe UTCTime -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (WorkerHeartbeat -> TVar (Maybe UTCTime)
whLastPoll WorkerHeartbeat
heartbeat) (UTCTime -> Maybe UTCTime
forall a. a -> Maybe a
Just UTCTime
now))

{- | The instant of the worker's last recorded progress, a successful poll or a completed
job, or 'Nothing' before its first.
-}
lastPoll :: WorkerHeartbeat -> IO (Maybe UTCTime)
lastPoll :: WorkerHeartbeat -> IO (Maybe UTCTime)
lastPoll = TVar (Maybe UTCTime) -> IO (Maybe UTCTime)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO (TVar (Maybe UTCTime) -> IO (Maybe UTCTime))
-> (WorkerHeartbeat -> TVar (Maybe UTCTime))
-> WorkerHeartbeat
-> IO (Maybe UTCTime)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkerHeartbeat -> TVar (Maybe UTCTime)
whLastPoll

{- | How long one long job step may run: uploading the largest artifact the memory plan admits
(512 MiB) over a 2 MiB-per-second link. Renewing a receipt's visibility never extends it.
-}
workerJobStepAllowance :: NominalDiffTime
workerJobStepAllowance :: NominalDiffTime
workerJobStepAllowance = NominalDiffTime
300

{- | Startup and progress allowance: a fetch and a publish of that largest artifact plus a
minute, so a slow job never restarts the process.
-}
workerHeartbeatStaleAfter :: NominalDiffTime
workerHeartbeatStaleAfter :: NominalDiffTime
workerHeartbeatStaleAfter = NominalDiffTime
2 NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
workerJobStepAllowance NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
+ NominalDiffTime
60

-- | Judge progress at @now@, using startup time only until the first successful progress.
heartbeatHealthy :: UTCTime -> UTCTime -> Maybe UTCTime -> Bool
heartbeatHealthy :: UTCTime -> UTCTime -> Maybe UTCTime -> Bool
heartbeatHealthy UTCTime
now UTCTime
startedAt Maybe UTCTime
polledAt =
    UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
now (UTCTime -> Maybe UTCTime -> UTCTime
forall a. a -> Maybe a -> a
fromMaybe UTCTime
startedAt Maybe UTCTime
polledAt) NominalDiffTime -> NominalDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
<= NominalDiffTime
workerHeartbeatStaleAfter

{- | What @\/livez@ answers from: the health verdict, plus the instant the checked loop last
recorded progress so an orchestrator can judge staleness rather than only pass or fail.
-}
data Liveness = Liveness
    { Liveness -> Bool
liveHealthy :: Bool
    , Liveness -> Maybe UTCTime
liveLastPoll :: Maybe UTCTime
    -- ^ 'Nothing' before the loop's first poll, and for a role that runs no such loop.
    }
    deriving stock (Liveness -> Liveness -> Bool
(Liveness -> Liveness -> Bool)
-> (Liveness -> Liveness -> Bool) -> Eq Liveness
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Liveness -> Liveness -> Bool
== :: Liveness -> Liveness -> Bool
$c/= :: Liveness -> Liveness -> Bool
/= :: Liveness -> Liveness -> Bool
Eq, Int -> Liveness -> ShowS
[Liveness] -> ShowS
Liveness -> String
(Int -> Liveness -> ShowS)
-> (Liveness -> String) -> ([Liveness] -> ShowS) -> Show Liveness
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Liveness -> ShowS
showsPrec :: Int -> Liveness -> ShowS
$cshow :: Liveness -> String
show :: Liveness -> String
$cshowList :: [Liveness] -> ShowS
showList :: [Liveness] -> ShowS
Show)

-- | The verdict of a role with no background loop to stall: live, with nothing to report.
alwaysLive :: Liveness
alwaysLive :: Liveness
alwaysLive = Liveness{liveHealthy :: Bool
liveHealthy = Bool
True, liveLastPoll :: Maybe UTCTime
liveLastPoll = Maybe UTCTime
forall a. Maybe a
Nothing}

{- | Read the worker heartbeat and judge it against the current wall clock, keeping the
instant judged. Both the embedded and the dedicated worker answer @\/livez@ through this.
-}
heartbeatLivenessNow :: WorkerHeartbeat -> IO Liveness
heartbeatLivenessNow :: WorkerHeartbeat -> IO Liveness
heartbeatLivenessNow WorkerHeartbeat
heartbeat = do
    now <- WorkerHeartbeat -> IO UTCTime
whNow WorkerHeartbeat
heartbeat
    polledAt <- lastPoll heartbeat
    pure Liveness{liveHealthy = heartbeatHealthy now (whStartedAt heartbeat) polledAt, liveLastPoll = polledAt}