{-# LANGUAGE OverloadedStrings #-}

-- | The guard's shared state. The handler, the guard loop, and the extend
-- thread the loop forks write it.
module Arbiter.Worker.Heartbeat.Guard.State
  ( GuardConfig (..)
  , Batch (..)
  , Guarded (..)
  , Status (..)
  , InFlight (..)
  , HeartbeatGuard (..)
  , Wake (..)
  , newHeartbeatGuard
  , guardKey
  , snapshot
  , adjust
  , pendingOf
  , wakeFor
  , wake
  , heartbeatWait
  , rebeat
  , toDiffTime
  , minRetryPause
  , settleGrace
  , leaseExpiredReason
  , reclaimedReason
  , deletedReason
  ) where

import Arbiter.Core.HighLevel (SetVisibilityResult (..))
import Arbiter.Core.Job.Types (JobId)
import Control.Concurrent.Class.MonadSTM (MonadSTM, STM, TVar, atomically, modifyTVar', newTVarIO, readTVar, readTVarIO)
import Control.Monad (when)
import Control.Monad.Class.MonadFork (MonadThread (..))
import Control.Monad.Class.MonadTime.SI (DiffTime, Time, UTCTime)
import Data.Fixed (Fixed (..))
import Data.List.NonEmpty (NonEmpty)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust)
import Data.Set (Set)
import Data.Text (Text)
import Data.Time.Clock (NominalDiffTime, nominalDiffTimeToSeconds, picosecondsToDiffTime)

import Arbiter.Worker.Logger (LogLevel (..))

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

-- | Settle time a landed extend gets past its give-up.
settleGrace :: DiffTime
settleGrace :: DiffTime
settleGrace = DiffTime
0.25

-- | The reason a batch is stopped at its lease.
leaseExpiredReason :: Text
leaseExpiredReason :: Text
leaseExpiredReason = Text
"lease expired without renewal"

-- | The reason a batch another worker reclaimed is stopped.
reclaimedReason :: Text
reclaimedReason :: Text
reclaimedReason = Text
"reclaimed by another worker"

-- | The reason a batch whose row was deleted under it is stopped.
deletedReason :: Text
deletedReason :: Text
deletedReason = Text
"deleted while held by this worker"

-- | A wall-clock span as a monotonic one.
toDiffTime :: NominalDiffTime -> DiffTime
toDiffTime :: NominalDiffTime -> DiffTime
toDiffTime NominalDiffTime
elapsed = Integer -> DiffTime
picosecondsToDiffTime Integer
picos
  where
    MkFixed Integer
picos = NominalDiffTime -> Fixed E12
nominalDiffTimeToSeconds NominalDiffTime
elapsed

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

-- | Set the next extend. An earlier one a recheck asked for stands.
rebeat :: Time -> Status n -> Status n
rebeat :: forall (n :: * -> *). Time -> Status n -> Status n
rebeat Time
at Status n
status = Status n
status {beatAt = if recheckAsked status then min at (beatAt status) else at}

-- | What the guard needs from the pool.
data GuardConfig n job = GuardConfig
  { forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configInterval :: DiffTime
  -- ^ Gap between extends of a batch.
  , forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configTimeout :: DiffTime
  -- ^ Lease length that each extend sets.
  , forall (n :: * -> *) job. GuardConfig n job -> Maybe DiffTime
configMaxDuration :: Maybe DiffTime
  -- ^ Interrupt a batch that runs longer than this. 'Nothing' sets no bound.
  , forall (n :: * -> *) job. GuardConfig n job -> job -> JobId
configKey :: job -> JobId
  -- ^ The job's primary key.
  , forall (n :: * -> *) job. GuardConfig n job -> job -> Maybe UTCTime
configLease :: job -> Maybe UTCTime
  -- ^ The row's lease deadline, as the claim read it back.
  , forall (n :: * -> *) job.
GuardConfig n job -> [job] -> n [SetVisibilityResult]
configExtend :: [job] -> n [SetVisibilityResult]
  -- ^ Extend the jobs' leases by 'configTimeout', reporting each.
  , forall (n :: * -> *) job. GuardConfig n job -> n ()
configExtended :: n ()
  -- ^ Runs after each extend that lands before the guard abandons it.
  , forall (n :: * -> *) job.
GuardConfig n job -> LogLevel -> [job] -> Text -> n ()
configLog :: LogLevel -> [job] -> Text -> n ()
  -- ^ The pool log, with the jobs in context.
  , forall (n :: * -> *) job.
GuardConfig n job -> job -> UTCTime -> UTCTime -> n ()
configHeartbeat :: job -> UTCTime -> UTCTime -> n ()
  -- ^ The heartbeat hook: the job, now, and the batch start.
  }

-- | A batch under guard.
data Batch n job = Batch
  { forall (n :: * -> *) job. Batch n job -> NonEmpty job
batchJobs :: NonEmpty job
  -- ^ The jobs the batch claimed.
  , forall (n :: * -> *) job. Batch n job -> n [job]
batchPending :: n [job]
  -- ^ Read on the guard's threads, for the jobs still awaiting an outcome.
  , forall (n :: * -> *) job. Batch n job -> UTCTime
batchStart :: UTCTime
  -- ^ When the batch started.
  , forall (n :: * -> *) job. Batch n job -> n () -> n ()
batchInherit :: n () -> n ()
  -- ^ Runs the heartbeat hooks under the batch's context.
  }

-- | A registered batch. Fixed at register.
data Guarded n job = Guarded
  { forall (n :: * -> *) job. Guarded n job -> Int
guardedToken :: Int
  , forall (n :: * -> *) job. Guarded n job -> Batch n job
guardedBatch :: Batch n job
  , forall (n :: * -> *) job. Guarded n job -> ThreadId n
guardedHandler :: ThreadId n
  , forall (n :: * -> *) job. Guarded n job -> Maybe Time
guardedDeadline :: Maybe Time
  -- ^ When the duration fence fires.
  , forall (n :: * -> *) job. Guarded n job -> TVar n (Status n)
guardedStatus :: TVar n (Status n)
  }

-- | A registered batch's timers and flags.
data Status n = Status
  { forall (n :: * -> *). Status n -> Time
leaseAt :: !Time
  -- ^ When the lease runs out. Register and settle write it.
  , forall (n :: * -> *). Status n -> Time
beatAt :: !Time
  -- ^ When the next extend is due. Register, settle, a failed extend and a recheck write it.
  , forall (n :: * -> *). Status n -> Bool
recheckAsked :: !Bool
  -- ^ A recheck asked for an extend that no issued extend covers yet. A recheck sets it, issue clears it.
  , forall (n :: * -> *). Status n -> Bool
leaseLapsed :: !Bool
  -- ^ The lease lapsed and the batch gets no further extend. The fence writes it.
  , forall (n :: * -> *). Status n -> Bool
deadlineSent :: !Bool
  -- ^ The deadline signal went out. The fence writes it.
  , forall (n :: * -> *). Status n -> Maybe Time
signalledAt :: !(Maybe Time)
  -- ^ When the last signal went out. Signal writes it.
  , forall (n :: * -> *). Status n -> Set JobId
goneSeen :: !(Set JobId)
  -- ^ Pending jobs the last settled extend found gone. Settle writes it.
  , forall (n :: * -> *). Status n -> Maybe [ThreadId n]
couriers :: !(Maybe [ThreadId n])
  -- ^ Threads carrying a signal to the handler. Nothing once unregistered. Couriers and unregister write it.
  }

-- | The extend statement in flight. The loop issues it, the extend thread lands and clears it.
-- The loop abandons it past its give-up and settle grace.
data InFlight = InFlight
  { InFlight -> Time
issuedAt :: !Time
  , InFlight -> Time
givesUp :: !Time
  -- ^ When its timeout fires.
  , InFlight -> Set Int
carries :: !(Set Int)
  -- ^ The batches it covers.
  , InFlight -> Bool
landed :: !Bool
  -- ^ It returned. Its results are being settled.
  }

-- | The pool's heartbeat guard.
data HeartbeatGuard n job = HeartbeatGuard
  { forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig :: GuardConfig n job
  , forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Map Int (Guarded n job))
guardEntries :: TVar n (Map Int (Guarded n job))
  , forall (n :: * -> *) job. HeartbeatGuard n job -> TVar n Int
guardNextToken :: TVar n Int
  , forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe InFlight)
guardInFlight :: TVar n (Maybe InFlight)
  , forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe Time)
guardAbandoned :: TVar n (Maybe Time)
  -- ^ The issue time of the abandoned extend still running. The loop sets it,
  -- that extend's thread clears it.
  , forall (n :: * -> *) job. HeartbeatGuard n job -> Wake n
guardWake :: Wake n
  }

-- | How the guard loop is woken before its target.
data Wake n = Wake
  { forall (n :: * -> *). Wake n -> TVar n Int
wakeCount :: TVar n Int
  -- ^ Bumped by each wake.
  , forall (n :: * -> *). Wake n -> TVar n (Maybe Time)
wakeTarget :: TVar n (Maybe Time)
  -- ^ The time the loop sleeps until.
  }

-- | A guard with no batches registered.
newHeartbeatGuard :: (MonadSTM n) => GuardConfig n job -> n (HeartbeatGuard n job)
newHeartbeatGuard :: forall (n :: * -> *) job.
MonadSTM n =>
GuardConfig n job -> n (HeartbeatGuard n job)
newHeartbeatGuard GuardConfig n job
config =
  GuardConfig n job
-> TVar n (Map Int (Guarded n job))
-> TVar n Int
-> TVar n (Maybe InFlight)
-> TVar n (Maybe Time)
-> Wake n
-> HeartbeatGuard n job
forall (n :: * -> *) job.
GuardConfig n job
-> TVar n (Map Int (Guarded n job))
-> TVar n Int
-> TVar n (Maybe InFlight)
-> TVar n (Maybe Time)
-> Wake n
-> HeartbeatGuard n job
HeartbeatGuard GuardConfig n job
config
    (TVar n (Map Int (Guarded n job))
 -> TVar n Int
 -> TVar n (Maybe InFlight)
 -> TVar n (Maybe Time)
 -> Wake n
 -> HeartbeatGuard n job)
-> n (TVar n (Map Int (Guarded n job)))
-> n (TVar n Int
      -> TVar n (Maybe InFlight)
      -> TVar n (Maybe Time)
      -> Wake n
      -> HeartbeatGuard n job)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Int (Guarded n job) -> n (TVar n (Map Int (Guarded n job)))
forall a. a -> n (TVar n a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Map Int (Guarded n job)
forall k a. Map k a
Map.empty
    n (TVar n Int
   -> TVar n (Maybe InFlight)
   -> TVar n (Maybe Time)
   -> Wake n
   -> HeartbeatGuard n job)
-> n (TVar n Int)
-> n (TVar n (Maybe InFlight)
      -> TVar n (Maybe Time) -> Wake n -> HeartbeatGuard n job)
forall a b. n (a -> b) -> n a -> n b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> n (TVar n Int)
forall a. a -> n (TVar n a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Int
0
    n (TVar n (Maybe InFlight)
   -> TVar n (Maybe Time) -> Wake n -> HeartbeatGuard n job)
-> n (TVar n (Maybe InFlight))
-> n (TVar n (Maybe Time) -> Wake n -> HeartbeatGuard n job)
forall a b. n (a -> b) -> n a -> n b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe InFlight -> n (TVar n (Maybe InFlight))
forall a. a -> n (TVar n a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Maybe InFlight
forall a. Maybe a
Nothing
    n (TVar n (Maybe Time) -> Wake n -> HeartbeatGuard n job)
-> n (TVar n (Maybe Time)) -> n (Wake n -> HeartbeatGuard n job)
forall a b. n (a -> b) -> n a -> n b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Time -> n (TVar n (Maybe Time))
forall a. a -> n (TVar n a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Maybe Time
forall a. Maybe a
Nothing
    n (Wake n -> HeartbeatGuard n job)
-> n (Wake n) -> n (HeartbeatGuard n job)
forall a b. n (a -> b) -> n a -> n b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TVar n Int -> TVar n (Maybe Time) -> Wake n
forall (n :: * -> *). TVar n Int -> TVar n (Maybe Time) -> Wake n
Wake (TVar n Int -> TVar n (Maybe Time) -> Wake n)
-> n (TVar n Int) -> n (TVar n (Maybe Time) -> Wake n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> n (TVar n Int)
forall a. a -> n (TVar n a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Int
0 n (TVar n (Maybe Time) -> Wake n)
-> n (TVar n (Maybe Time)) -> n (Wake n)
forall a b. n (a -> b) -> n a -> n b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Time -> n (TVar n (Maybe Time))
forall a. a -> n (TVar n a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Maybe Time
forall a. Maybe a
Nothing)

-- | The key the guard tracks jobs by.
guardKey :: HeartbeatGuard n job -> job -> JobId
guardKey :: forall (n :: * -> *) job. HeartbeatGuard n job -> job -> JobId
guardKey = GuardConfig n job -> job -> JobId
forall (n :: * -> *) job. GuardConfig n job -> job -> JobId
configKey (GuardConfig n job -> job -> JobId)
-> (HeartbeatGuard n job -> GuardConfig n job)
-> HeartbeatGuard n job
-> job
-> JobId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig

-- | Every registered batch with its status.
snapshot :: (MonadSTM n) => HeartbeatGuard n job -> STM n [(Guarded n job, Status n)]
snapshot :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> STM n [(Guarded n job, Status n)]
snapshot HeartbeatGuard n job
guard = do
  entries <- TVar n (Map Int (Guarded n job)) -> STM n (Map Int (Guarded n job))
forall a. TVar n a -> STM n a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar (HeartbeatGuard n job -> TVar n (Map Int (Guarded n job))
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Map Int (Guarded n job))
guardEntries HeartbeatGuard n job
guard)
  traverse (\Guarded n job
entry -> (,) Guarded n job
entry (Status n -> (Guarded n job, Status n))
-> STM n (Status n) -> STM n (Guarded n job, Status n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar n (Status n) -> STM n (Status n)
forall a. TVar n a -> STM n a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar (Guarded n job -> TVar n (Status n)
forall (n :: * -> *) job. Guarded n job -> TVar n (Status n)
guardedStatus Guarded n job
entry)) (Map.elems entries)

-- | Update a batch's status in one transaction.
adjust :: (MonadSTM n) => Guarded n job -> (Status n -> Status n) -> n ()
adjust :: forall (n :: * -> *) job.
MonadSTM n =>
Guarded n job -> (Status n -> Status n) -> n ()
adjust Guarded n job
entry = STM n () -> n ()
forall a. HasCallStack => STM n a -> n a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM n () -> n ())
-> ((Status n -> Status n) -> STM n ())
-> (Status n -> Status n)
-> n ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar n (Status n) -> (Status n -> Status n) -> STM n ()
forall a. TVar n a -> (a -> a) -> STM n ()
forall (m :: * -> *) a.
MonadSTM m =>
TVar m a -> (a -> a) -> STM m ()
modifyTVar' (Guarded n job -> TVar n (Status n)
forall (n :: * -> *) job. Guarded n job -> TVar n (Status n)
guardedStatus Guarded n job
entry)

-- | The batch's jobs still awaiting an outcome. None once it unregistered.
pendingOf :: (MonadSTM n) => Guarded n job -> n [job]
pendingOf :: forall (n :: * -> *) job. MonadSTM n => Guarded n job -> n [job]
pendingOf Guarded n job
entry = do
  status <- TVar n (Status n) -> n (Status n)
forall a. TVar n a -> n a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> m a
readTVarIO (Guarded n job -> TVar n (Status n)
forall (n :: * -> *) job. Guarded n job -> TVar n (Status n)
guardedStatus Guarded n job
entry)
  if isJust (couriers status) then batchPending (guardedBatch entry) else pure []

-- | Wake the loop when @at@ precedes its target.
wakeFor :: (MonadSTM n) => HeartbeatGuard n job -> Time -> STM n ()
wakeFor :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> STM n ()
wakeFor HeartbeatGuard n job
guard Time
at = do
  target <- TVar n (Maybe Time) -> STM n (Maybe Time)
forall a. TVar n a -> STM n a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar (Wake n -> TVar n (Maybe Time)
forall (n :: * -> *). Wake n -> TVar n (Maybe Time)
wakeTarget (HeartbeatGuard n job -> Wake n
forall (n :: * -> *) job. HeartbeatGuard n job -> Wake n
guardWake HeartbeatGuard n job
guard))
  when (maybe True (at <) target) (wake guard)

-- | Wake the loop now.
wake :: (MonadSTM n) => HeartbeatGuard n job -> STM n ()
wake :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> STM n ()
wake HeartbeatGuard n job
guard = TVar n Int -> (Int -> Int) -> STM n ()
forall a. TVar n a -> (a -> a) -> STM n ()
forall (m :: * -> *) a.
MonadSTM m =>
TVar m a -> (a -> a) -> STM m ()
modifyTVar' (Wake n -> TVar n Int
forall (n :: * -> *). Wake n -> TVar n Int
wakeCount (HeartbeatGuard n job -> Wake n
forall (n :: * -> *) job. HeartbeatGuard n job -> Wake n
guardWake HeartbeatGuard n job
guard)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)