{-# LANGUAGE OverloadedStrings #-}

module Arbiter.Worker.Heartbeat
  ( withJobsHeartbeat
  ) where

import Arbiter.Core.Exceptions (JobDeadlineExceeded (..), JobForceCancelled (..), displayEx, throwJobGoneIds)
import Arbiter.Core.HighLevel (JobOperation)
import Arbiter.Core.HighLevel qualified as Arb
import Arbiter.Core.Job.Types (JobRead, ObservabilityHooks (..), primaryKey)
import Arbiter.Core.Trace (capturingContext)
import Control.Concurrent (ThreadId, myThreadId)
import Control.Exception (Exception (..), asyncExceptionFromException, asyncExceptionToException, throwIO, throwTo)
import Control.Monad (unless, void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Foldable (traverse_)
import Data.List.NonEmpty (NonEmpty)
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Time (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime)
import Data.Void (Void, absurd)
import GHC.Clock (getMonotonicTime)
import UnliftIO.Async (race)
import UnliftIO.Concurrent (forkIO, threadDelay)
import UnliftIO.Exception (catchSyncOrAsync, tryAny)
import UnliftIO.STM (TMVar, TVar, atomically)
import UnliftIO.STM qualified as STM
import UnliftIO.Timeout (timeout)

import Arbiter.Worker.Logger (LogConfig, LogLevel (..), tryLog)
import Arbiter.Worker.Logger.Internal (runHook, withJobContext, withJobContextOne)
import Arbiter.Worker.Retry (isJobSignal)
import Arbiter.Worker.Settle (hasIdIn)

-- | Thrown into the handler thread at the duration deadline. Async, so sync catches cannot swallow it.
newtype DeadlineSignal = DeadlineSignal T.Text
  deriving stock (Int -> DeadlineSignal -> ShowS
[DeadlineSignal] -> ShowS
DeadlineSignal -> String
(Int -> DeadlineSignal -> ShowS)
-> (DeadlineSignal -> String)
-> ([DeadlineSignal] -> ShowS)
-> Show DeadlineSignal
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DeadlineSignal -> ShowS
showsPrec :: Int -> DeadlineSignal -> ShowS
$cshow :: DeadlineSignal -> String
show :: DeadlineSignal -> String
$cshowList :: [DeadlineSignal] -> ShowS
showList :: [DeadlineSignal] -> ShowS
Show)

instance Exception DeadlineSignal where
  backtraceDesired :: DeadlineSignal -> Bool
backtraceDesired DeadlineSignal
_ = Bool
False
  toException :: DeadlineSignal -> SomeException
toException = DeadlineSignal -> SomeException
forall e. Exception e => e -> SomeException
asyncExceptionToException
  fromException :: SomeException -> Maybe DeadlineSignal
fromException = SomeException -> Maybe DeadlineSignal
forall e. Exception e => SomeException -> Maybe e
asyncExceptionFromException

-- | Shortest gap between failed extends.
minRetryPause :: Double
minRetryPause :: Double
minRetryPause = Double
0.25

-- | Wait before the next extend. At most a beat and at least 'minRetryPause'.
heartbeatWait :: NominalDiffTime -> Bool -> Double -> Double
heartbeatWait :: NominalDiffTime -> Bool -> Double -> Double
heartbeatWait NominalDiffTime
intervalSecs Bool
extended Double
remaining
  | Bool
extended, Double
remaining Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
beat = Double
beat
  | Bool
otherwise = Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
beat (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
minRetryPause (Double
remaining Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2))
  where
    beat :: Double
beat = NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
intervalSecs

-- | Run an action with a heartbeat thread that extends job visibility at each
-- interval. Skip jobs that the handler finalized. An absent job is a normal
-- condition. A reclaimed job causes an exception. Call the heartbeat hook after
-- each successful extension.
withJobsHeartbeat
  :: forall payload m a
   . (JobOperation m payload)
  => ObservabilityHooks m payload
  -- ^ Observability hooks (for heartbeat hook)
  -> NominalDiffTime
  -- ^ Heartbeat interval
  -> NominalDiffTime
  -- ^ Visibility timeout
  -> Maybe NominalDiffTime
  -- ^ Longest the handler may run before the fence interrupts it
  -> UTCTime
  -- ^ Start time (for calculating elapsed time in heartbeat hook)
  -> NonEmpty (JobRead payload)
  -- ^ The job(s) being processed
  -> m [JobRead payload]
  -- ^ Read each tick for the job(s) still awaiting an outcome.
  -> LogConfig
  -- ^ Log configuration
  -> TMVar ()
  -- ^ Proof-of-work signal pulsed after each successful heartbeat.
  -> m a
  -- ^ Action to run with heartbeat protection
  -> m a
withJobsHeartbeat :: forall payload (m :: * -> *) a.
JobOperation m payload =>
ObservabilityHooks m payload
-> NominalDiffTime
-> NominalDiffTime
-> Maybe NominalDiffTime
-> UTCTime
-> NonEmpty (JobRead payload)
-> m [JobRead payload]
-> LogConfig
-> TMVar ()
-> m a
-> m a
withJobsHeartbeat ObservabilityHooks m payload
hooks NominalDiffTime
intervalSecs NominalDiffTime
timeoutSecs Maybe NominalDiffTime
maxDuration UTCTime
startTime NonEmpty (JobRead payload)
jobs m [JobRead payload]
pending LogConfig
logCfg TMVar ()
signal m a
action = do
  -- 'race' forks each side. The handler and the guard reattach the job span.
  inherited <- m (m a -> m a)
forall (m :: * -> *) a. MonadUnliftIO m => m (m a -> m a)
capturingContext
  wallNow <- liftIO getCurrentTime
  monoNow <- liftIO getMonotonicTime
  let elapsed = UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
wallNow UTCTime
startTime
      durationDeadline = (\NominalDiffTime
limit -> Double
monoNow Double -> Double -> Double
forall a. Num a => a -> a -> a
+ NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
limit) (NominalDiffTime -> Double)
-> Maybe NominalDiffTime -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe NominalDiffTime
maxDuration
  lease <- STM.newTVarIO (monoNow + realToFrac (timeoutSecs - elapsed))
  handlerId <- STM.newTVarIO Nothing
  outcome <-
    race
      (inherited (absurd <$> guardThread lease durationDeadline handlerId))
      (inherited (liftIO myThreadId >>= atomically . STM.writeTVar handlerId . Just >> deadlineAsSync action))
  pure (either id id outcome)
  where
    -- The signal leaves the handler thread as the sync exception settlement classifies.
    deadlineAsSync :: m a -> m a
deadlineAsSync = (m a -> (DeadlineSignal -> m a) -> m a)
-> (DeadlineSignal -> m a) -> m a -> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip m a -> (DeadlineSignal -> m a) -> m a
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
catchSyncOrAsync (\(DeadlineSignal Text
msg) -> IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (JobDeadlineExceeded -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (Text -> JobDeadlineExceeded
JobDeadlineExceeded Text
msg)))

    -- Sleeps until the earliest of the next beat, the lease end and the duration deadline.
    guardThread :: TVar Double -> Maybe Double -> TVar (Maybe ThreadId) -> m Void
    guardThread :: TVar Double -> Maybe Double -> TVar (Maybe ThreadId) -> m Void
guardThread TVar Double
lease Maybe Double
durationDeadline TVar (Maybe ThreadId)
handlerId = Bool -> Bool -> m Void
loop Bool
True Bool
False
      where
        loop :: Bool -> Bool -> m Void
loop Bool
extended Bool
fenced = do
          leaseUntil <- TVar Double -> m Double
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
STM.readTVarIO TVar Double
lease
          now <- liftIO getMonotonicTime
          let beatAt = Double
now Double -> Double -> Double
forall a. Num a => a -> a -> a
+ NominalDiffTime -> Bool -> Double -> Double
heartbeatWait NominalDiffTime
intervalSecs Bool
extended (Double
leaseUntil Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
now)
              fenceAt = if Bool
fenced then Maybe Double
forall a. Maybe a
Nothing else Maybe Double
durationDeadline
              due = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (Double
beatAt Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Double
leaseUntil Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double] -> (Double -> [Double]) -> Maybe Double -> [Double]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] Double -> [Double]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Double
fenceAt)
          threadDelay (max 0 (ceiling ((due - now) * 1_000_000)))
          woke <- liftIO getMonotonicTime
          leaseNow <- STM.readTVarIO lease
          live <- if woke >= leaseNow then pending else pure []
          if woke >= leaseNow && not (null live)
            then throwJobGoneIds "lease expired without renewal" (map primaryKey live)
            else case fenceAt of
              Just Double
deadline | Double
woke Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
deadline -> do
                -- Delivered off-thread so a masked handler cannot stall the beat. The handler unwinds under a live lease.
                TVar (Maybe ThreadId) -> m (Maybe ThreadId)
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
STM.readTVarIO TVar (Maybe ThreadId)
handlerId m (Maybe ThreadId) -> (Maybe ThreadId -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ThreadId -> m ()) -> Maybe ThreadId -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (\ThreadId
tid -> m ThreadId -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m () -> m ThreadId
forall (m :: * -> *). MonadUnliftIO m => m () -> m ThreadId
forkIO (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ThreadId -> DeadlineSignal -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
tid (Text -> DeadlineSignal
DeadlineSignal Text
durationMessage)))))
                Bool -> Bool -> m Void
loop Bool
extended Bool
True
              Maybe Double
_
                | Double
woke Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
beatAt -> Bool -> m Void
attempt Bool
fenced
                | Bool
otherwise -> Bool -> Bool -> m Void
loop Bool
extended Bool
fenced
        -- An extend that hangs past the lease must not hold up the lease check.
        attempt :: Bool -> m Void
attempt Bool
fenced = do
          leaseUntil <- TVar Double -> m Double
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
STM.readTVarIO TVar Double
lease
          now <- liftIO getMonotonicTime
          outcome <- timeout (ceiling (max minRetryPause (leaseUntil - now) * 1_000_000)) (tryAny (tick lease))
          case outcome of
            Maybe (Either SomeException ())
Nothing -> Bool -> Bool -> m Void
loop Bool
False Bool
fenced
            Just (Right ()) -> Bool -> Bool -> m Void
loop Bool
True Bool
fenced
            Just (Left SomeException
exception)
              | SomeException -> Bool
isJobSignal SomeException
exception -> IO Void -> m Void
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (SomeException -> IO Void
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
exception)
              | Bool
otherwise -> do
                  LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog (LogConfig -> NonEmpty (JobRead payload) -> LogConfig
forall payload.
LogConfig -> NonEmpty (JobRead payload) -> LogConfig
withJobContext LogConfig
logCfg NonEmpty (JobRead payload)
jobs) LogLevel
Error (Text
"Heartbeat error (retrying): " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SomeException -> Text
displayEx SomeException
exception)
                  Bool -> Bool -> m Void
loop Bool
False Bool
fenced

    tick :: TVar Double -> m ()
tick TVar Double
lease = do
      live <- m [JobRead payload]
pending
      -- Read before the extend. The tracked deadline stays inside the row's.
      issuedAt <- liftIO getMonotonicTime
      results <- Arb.setVisibilityTimeoutBatch timeoutSecs live
      -- Rows this worker settled during the statement do not count.
      stillPending <- Set.fromList . map primaryKey <$> pending
      let cancelledJobs = [Int64
jobId | Arb.JobCancelled Int64
jobId <- [SetVisibilityResult]
results]
          stolenJobs = [Int64
jobId | Arb.JobReclaimed Int64
jobId Int64
_ Int64
_ <- [SetVisibilityResult]
results]
          goneJobs = [Int64
jobId | Arb.JobGone Int64
jobId <- [SetVisibilityResult]
results]
          unmoved = [() | Arb.VisibilityUnchanged Int64
jobId <- [SetVisibilityResult]
results, Int64 -> Set Int64 -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int64
jobId Set Int64
stillPending]
      atomically $ do
        when (null unmoved) $ STM.writeTVar lease (issuedAt + realToFrac timeoutSecs)
        void $ STM.tryPutTMVar signal ()
      unless (null cancelledJobs) $
        liftIO (throwIO (JobForceCancelled cancelledJobs (stolenJobs <> goneJobs)))
      unless (null stolenJobs) $
        throwJobGoneIds "reclaimed by another worker" stolenJobs
      let activeJobIds = [Int64] -> Set Int64
forall a. Ord a => [a] -> Set a
Set.fromList [Int64
jobId | Arb.VisibilityExtended Int64
jobId <- [SetVisibilityResult]
results]
          activeJobs = (JobRead payload -> Bool) -> [JobRead payload] -> [JobRead payload]
forall a. (a -> Bool) -> [a] -> [a]
filter (Set Int64 -> JobRead payload -> Bool
forall payload. Set Int64 -> JobRead payload -> Bool
hasIdIn Set Int64
activeJobIds) [JobRead payload]
live
      currentTime <- liftIO getCurrentTime
      traverse_
        ( \JobRead payload
job ->
            LogConfig -> Text -> m () -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> Text -> m () -> m ()
runHook (LogConfig -> JobRead payload -> LogConfig
forall payload. LogConfig -> JobRead payload -> LogConfig
withJobContextOne LogConfig
logCfg JobRead payload
job) Text
"onJobHeartbeat" (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
              ObservabilityHooks m payload
-> JobPayload payload =>
   JobRead payload -> UTCTime -> UTCTime -> m ()
forall (m :: * -> *) payload.
ObservabilityHooks m payload
-> JobPayload payload =>
   JobRead payload -> UTCTime -> UTCTime -> m ()
onJobHeartbeat ObservabilityHooks m payload
hooks JobRead payload
job UTCTime
currentTime UTCTime
startTime
        )
        activeJobs

    durationMessage :: Text
durationMessage =
      Text
"handler ran past the maximum job duration"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> (NominalDiffTime -> Text) -> Maybe NominalDiffTime -> Text
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Text
" of " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text)
-> (NominalDiffTime -> Text) -> NominalDiffTime -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack (String -> Text)
-> (NominalDiffTime -> String) -> NominalDiffTime -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NominalDiffTime -> String
forall a. Show a => a -> String
show) Maybe NominalDiffTime
maxDuration