{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- The pool's heartbeat guard, run in IO. These are the IO wrappers over
-- "Arbiter.Worker.Heartbeat.Guard".
module Arbiter.Worker.Heartbeat
  ( PoolGuard
  , newHeartbeatGuard
  , runHeartbeatGuard
  , recheckJob
  ) where

import Arbiter.Core.HighLevel (JobOperation)
import Arbiter.Core.HighLevel qualified as Arb
import Arbiter.Core.Job.Types (JobId, JobRead, ObservabilityHooks (..), notVisibleUntil, primaryKey)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Void (Void)
import UnliftIO (UnliftIO (..), askUnliftIO)
import UnliftIO.STM (atomically)

import Arbiter.Worker.Config (WorkerConfig (..), pulseHeartbeat)
import Arbiter.Worker.Heartbeat.Guard (GuardConfig (..), toDiffTime)
import Arbiter.Worker.Heartbeat.Guard qualified as Guard
import Arbiter.Worker.Logger.Internal (jobHook, poolLog)

-- | The pool's guard, run in IO and keyed on the job.
type PoolGuard payload = Guard.HeartbeatGuard IO (JobRead payload)

-- | Build the pool's guard from its config. IO wrapper over 'Arbiter.Worker.Heartbeat.Guard.newHeartbeatGuard'.
newHeartbeatGuard :: (JobOperation m payload) => WorkerConfig m payload -> m (PoolGuard payload)
newHeartbeatGuard :: forall (m :: * -> *) payload.
JobOperation m payload =>
WorkerConfig m payload -> m (PoolGuard payload)
newHeartbeatGuard WorkerConfig m payload
config = do
  UnliftIO run <- m (UnliftIO m)
forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  liftIO . Guard.newHeartbeatGuard $
    GuardConfig
      { configInterval = toDiffTime (jobHeartbeatInterval config)
      , configTimeout = toDiffTime (visibilityTimeout config)
      , configMaxDuration = toDiffTime <$> maxJobDuration config
      , configKey = primaryKey
      , configLease = notVisibleUntil
      , configExtend = run . Arb.setVisibilityTimeoutBatch (visibilityTimeout config)
      , configExtended = atomically (pulseHeartbeat config)
      , configLog = poolLog (logConfig config)
      , configHeartbeat = \JobRead payload
job UTCTime
now UTCTime
start ->
          m () -> IO ()
forall a. m a -> IO a
run (LogConfig -> JobRead payload -> Text -> m () -> m ()
forall (m :: * -> *) payload.
MonadUnliftIO m =>
LogConfig -> JobRead payload -> Text -> m () -> m ()
jobHook (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) JobRead payload
job Text
"onJobHeartbeat" (ObservabilityHooks m payload
-> JobPayload payload =>
   JobRead payload -> UTCTime -> UTCTime -> m ()
forall (m :: * -> *) payload.
ObservabilityHooks m payload
-> JobPayload payload =>
   JobRead payload -> UTCTime -> UTCTime -> m ()
onJobHeartbeat (WorkerConfig m payload -> ObservabilityHooks m payload
forall (m :: * -> *) payload.
WorkerConfig m payload -> ObservabilityHooks m payload
observabilityHooks WorkerConfig m payload
config) JobRead payload
job UTCTime
now UTCTime
start))
      }

-- | Run the guard loop. IO wrapper over 'Arbiter.Worker.Heartbeat.Guard.runHeartbeatGuard'.
runHeartbeatGuard :: (MonadIO m) => PoolGuard payload -> m Void
runHeartbeatGuard :: forall (m :: * -> *) payload.
MonadIO m =>
PoolGuard payload -> m Void
runHeartbeatGuard = IO Void -> m Void
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Void -> m Void)
-> (PoolGuard payload -> IO Void) -> PoolGuard payload -> m Void
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PoolGuard payload -> IO Void
forall (n :: * -> *) job.
(MonadFork n, MonadMask n, MonadTime n, MonadTimer n) =>
HeartbeatGuard n job -> n Void
Guard.runHeartbeatGuard

-- | Extend the batch that holds the job now.
recheckJob :: (MonadIO m) => PoolGuard payload -> JobId -> m ()
recheckJob :: forall (m :: * -> *) payload.
MonadIO m =>
PoolGuard payload -> JobId -> m ()
recheckJob PoolGuard payload
guard = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (JobId -> IO ()) -> JobId -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PoolGuard payload -> JobId -> IO ()
forall (n :: * -> *) job.
(MonadMonotonicTime n, MonadSTM n) =>
HeartbeatGuard n job -> JobId -> n ()
Guard.recheck PoolGuard payload
guard