{-# LANGUAGE OverloadedStrings #-}

-- | Guard loop for lease fencing, deadlines, and heartbeat extension.
-- Extensions run on separate threads to keep fencing responsive to stalled
-- statements. Registration wakes the loop for an earlier deadline.
--
-- The lease fence waits for an extend that carries the batch. It waits until
-- that extend gives up. If the extend lands, it waits a settle grace past
-- that. The fence and the settle each take the batch's status in one
-- transaction. Whichever commits first wins and the other sees it.
module Arbiter.Worker.Heartbeat.Guard.Loop
  ( runHeartbeatGuard
  , trySync
  ) where

import Arbiter.Core.Enum (showText)
import Arbiter.Core.Exceptions (JobDeadlineExceeded (..), JobForceCancelled (..), JobGoneException (..), displayEx)
import Arbiter.Core.HighLevel (SetVisibilityResult (..))
import Arbiter.Core.Job.Types (JobId)
import Control.Concurrent.Class.MonadSTM
  ( MonadSTM
  , STM
  , atomically
  , check
  , modifyTVar'
  , orElse
  , readTVar
  , retry
  , stateTVar
  , writeTVar
  )
import Control.Monad (filterM, forever, unless, void, when, (>=>))
import Control.Monad.Class.MonadFork (MonadFork (..))
import Control.Monad.Class.MonadThrow (Exception (..), MonadCatch (..), MonadMask (..), MonadThrow (..), SomeException)
import Control.Monad.Class.MonadTime.SI
  ( DiffTime
  , MonadMonotonicTime (..)
  , MonadTime (..)
  , Time
  , UTCTime
  , addTime
  , diffTime
  )
import Control.Monad.Class.MonadTimer.SI (MonadTimer (..))
import Data.Foldable (for_, toList, traverse_)
import Data.List.NonEmpty (nonEmpty)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust, isNothing, mapMaybe)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Void (Void)
import UnliftIO.Exception (isSyncException)

import Arbiter.Worker.Heartbeat.Guard.Signal (signal)
import Arbiter.Worker.Heartbeat.Guard.State
import Arbiter.Worker.Logger (LogLevel (..))

-- | The guard loop. Fences what is due, then extends what is due.
runHeartbeatGuard
  :: (MonadFork n, MonadMask n, MonadTime n, MonadTimer n)
  => HeartbeatGuard n job
  -> n Void
{-# SPECIALIZE runHeartbeatGuard :: HeartbeatGuard IO job -> IO Void #-}
runHeartbeatGuard :: forall (n :: * -> *) job.
(MonadFork n, MonadMask n, MonadTime n, MonadTimer n) =>
HeartbeatGuard n job -> n Void
runHeartbeatGuard HeartbeatGuard n job
guard = n () -> n Void
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (n () -> n Void) -> n () -> n Void
forall a b. (a -> b) -> a -> b
$ do
  (target, count) <- STM n (Maybe Time, Int) -> n (Maybe Time, Int)
forall a. HasCallStack => STM n a -> n a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (HeartbeatGuard n job -> STM n (Maybe Time, Int)
forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> STM n (Maybe Time, Int)
plan HeartbeatGuard n job
guard)
  sleepUntil guard count target
  woke <- getMonotonicTime
  (inFlight, current) <- atomically $ do
    abandon guard woke
    (,) <$> readTVar (guardInFlight guard) <*> snapshot guard
  -- A batch the fence stops gets no beat.
  leased <- filterM (fence guard woke) current
  let due = [Guarded n job
entry | (Guarded n job
entry, Status n
status) <- [(Guarded n job, Status n)]
leased, Status n -> Time
forall (n :: * -> *). Status n -> Time
beatAt Status n
status Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
<= Time
woke]
  unless (null due || isJust inFlight) (issue guard woke leased due)

-- ---------------------------------------------------------------------------
-- Sleeping
-- ---------------------------------------------------------------------------

-- | Publish the loop's target, the earliest due time.
plan :: (MonadSTM n) => HeartbeatGuard n job -> STM n (Maybe Time, Int)
plan :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> STM n (Maybe Time, Int)
plan HeartbeatGuard n job
guard = do
  inFlight <- TVar n (Maybe InFlight) -> STM n (Maybe InFlight)
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 (Maybe InFlight)
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe InFlight)
guardInFlight HeartbeatGuard n job
guard)
  abandoned <- readTVar (guardAbandoned guard)
  entries <- snapshot guard
  let times = Maybe Time -> [Time]
forall a. Maybe a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Maybe Time -> Maybe InFlight -> Maybe Time
abandonAt Maybe Time
abandoned Maybe InFlight
inFlight) [Time] -> [Time] -> [Time]
forall a. Semigroup a => a -> a -> a
<> ((Guarded n job, Status n) -> [Time])
-> [(Guarded n job, Status n)] -> [Time]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Maybe InFlight -> (Guarded n job, Status n) -> [Time]
forall (n :: * -> *) job.
Maybe InFlight -> (Guarded n job, Status n) -> [Time]
dueTimes Maybe InFlight
inFlight) [(Guarded n job, Status n)]
entries
  let target = NonEmpty Time -> Time
forall a. Ord a => NonEmpty a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (NonEmpty Time -> Time) -> Maybe (NonEmpty Time) -> Maybe Time
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Time] -> Maybe (NonEmpty Time)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [Time]
times
  writeTVar (wakeTarget (guardWake guard)) target
  (,) target <$> readTVar (wakeCount (guardWake guard))

-- | When the guard next acts on a batch. Beats wait for the extend in flight.
dueTimes :: Maybe InFlight -> (Guarded n job, Status n) -> [Time]
dueTimes :: forall (n :: * -> *) job.
Maybe InFlight -> (Guarded n job, Status n) -> [Time]
dueTimes Maybe InFlight
inFlight (Guarded n job
entry, Status n
status) =
  [Maybe InFlight -> Guarded n job -> Status n -> Time
forall (n :: * -> *) job.
Maybe InFlight -> Guarded n job -> Status n -> Time
leaseFenceAt Maybe InFlight
inFlight Guarded n job
entry Status n
status | Bool -> Bool
not (Status n -> Bool
forall (n :: * -> *). Status n -> Bool
leaseLapsed Status n
status)]
    [Time] -> [Time] -> [Time]
forall a. Semigroup a => a -> a -> a
<> [Status n -> Time
forall (n :: * -> *). Status n -> Time
beatAt Status n
status | Bool -> Bool
not (Status n -> Bool
forall (n :: * -> *). Status n -> Bool
leaseLapsed Status n
status), Maybe InFlight -> Bool
forall a. Maybe a -> Bool
isNothing Maybe InFlight
inFlight]
    [Time] -> [Time] -> [Time]
forall a. Semigroup a => a -> a -> a
<> [Time
deadline | Bool -> Bool
not (Status n -> Bool
forall (n :: * -> *). Status n -> Bool
deadlineSent Status n
status), Just Time
deadline <- [Guarded n job -> Maybe Time
forall (n :: * -> *) job. Guarded n job -> Maybe Time
guardedDeadline Guarded n job
entry]]

-- | When the lease fence fires. An extend carrying the batch holds it until the extend
-- gives up, or, once landed, a settle grace past that.
leaseFenceAt :: Maybe InFlight -> Guarded n job -> Status n -> Time
leaseFenceAt :: forall (n :: * -> *) job.
Maybe InFlight -> Guarded n job -> Status n -> Time
leaseFenceAt Maybe InFlight
inFlight Guarded n job
entry Status n
status = case Maybe InFlight
inFlight of
  Just InFlight
running
    | InFlight -> Bool
carried InFlight
running, InFlight -> Bool
landed InFlight
running -> Time -> Time -> Time
forall a. Ord a => a -> a -> a
max Time
lease (DiffTime -> Time -> Time
addTime DiffTime
settleGrace (InFlight -> Time
givesUp InFlight
running))
    | InFlight -> Bool
carried InFlight
running -> Time -> Time -> Time
forall a. Ord a => a -> a -> a
max Time
lease (InFlight -> Time
givesUp InFlight
running)
  Maybe InFlight
_ -> Time
lease
  where
    lease :: Time
lease = Status n -> Time
forall (n :: * -> *). Status n -> Time
leaseAt Status n
status
    carried :: InFlight -> Bool
carried InFlight
running = Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Guarded n job -> Int
forall (n :: * -> *) job. Guarded n job -> Int
guardedToken Guarded n job
entry) (InFlight -> Set Int
carries InFlight
running) Bool -> Bool -> Bool
&& InFlight -> Time
issuedAt InFlight
running Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
< Time
lease

-- | Sleep until @at@ or the next wake.
sleepUntil :: (MonadTimer n) => HeartbeatGuard n job -> Int -> Maybe Time -> n ()
sleepUntil :: forall (n :: * -> *) job.
MonadTimer n =>
HeartbeatGuard n job -> Int -> Maybe Time -> n ()
sleepUntil HeartbeatGuard n job
guard Int
count Maybe Time
at = do
  alarm <- (Time -> n (TVar n Bool)) -> Maybe Time -> n (Maybe (TVar n Bool))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Time -> n (TVar n Bool)
forall {m :: * -> *}. MonadTimer m => Time -> m (TVar m Bool)
arm Maybe Time
at
  atomically $
    (readTVar (wakeCount (guardWake guard)) >>= check . (/= count))
      `orElse` maybe retry (readTVar >=> check) alarm
  where
    arm :: Time -> m (TVar m Bool)
arm Time
target = do
      now <- m Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
      registerDelay (max 0 (target `diffTime` now))

-- ---------------------------------------------------------------------------
-- Fencing
-- ---------------------------------------------------------------------------

-- | Signal a handler past its lease or its deadline. Whether its lease still stands.
fence :: (MonadFork n, MonadSTM n) => HeartbeatGuard n job -> Time -> (Guarded n job, Status n) -> n Bool
fence :: forall (n :: * -> *) job.
(MonadFork n, MonadSTM n) =>
HeartbeatGuard n job -> Time -> (Guarded n job, Status n) -> n Bool
fence HeartbeatGuard n job
guard Time
woke (Guarded n job
entry, Status n
status) = do
  (lapsed, standing) <- STM n (Bool, Bool) -> n (Bool, Bool)
forall a. HasCallStack => STM n a -> n a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (HeartbeatGuard n job -> Time -> Guarded n job -> STM n (Bool, Bool)
forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> Guarded n job -> STM n (Bool, Bool)
lapse HeartbeatGuard n job
guard Time
woke Guarded n job
entry)
  when lapsed $ do
    live <- pendingOf entry
    unless (null live) $
      signal guard woke entry (toException (JobGoneException leaseExpiredReason (map (guardKey guard) live)))
  for_ (guardedDeadline entry) $ \Time
deadline ->
    Bool -> n () -> n ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (Status n -> Bool
forall (n :: * -> *). Status n -> Bool
deadlineSent Status n
status) Bool -> Bool -> Bool
&& Time
woke Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
>= Time
deadline) (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$ do
      Guarded n job -> (Status n -> Status n) -> n ()
forall (n :: * -> *) job.
MonadSTM n =>
Guarded n job -> (Status n -> Status n) -> n ()
adjust Guarded n job
entry (\Status n
current -> Status n
current {deadlineSent = True})
      HeartbeatGuard n job
-> Time -> Guarded n job -> SomeException -> n ()
forall (n :: * -> *) job.
(MonadFork n, MonadSTM n) =>
HeartbeatGuard n job
-> Time -> Guarded n job -> SomeException -> n ()
signal HeartbeatGuard n job
guard Time
woke Guarded n job
entry (JobDeadlineExceeded -> SomeException
forall e. Exception e => e -> SomeException
toException (Text -> JobDeadlineExceeded
JobDeadlineExceeded (HeartbeatGuard n job -> Text
forall (n :: * -> *) job. HeartbeatGuard n job -> Text
durationMessage HeartbeatGuard n job
guard)))
  pure standing

-- | Mark the lease lapsed if it is, against the extend in flight. Whether it lapsed now, and whether it stands.
lapse :: (MonadSTM n) => HeartbeatGuard n job -> Time -> Guarded n job -> STM n (Bool, Bool)
lapse :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> Guarded n job -> STM n (Bool, Bool)
lapse HeartbeatGuard n job
guard Time
woke Guarded n job
entry = do
  inFlight <- TVar n (Maybe InFlight) -> STM n (Maybe InFlight)
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 (Maybe InFlight)
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe InFlight)
guardInFlight HeartbeatGuard n job
guard)
  status <- readTVar (guardedStatus entry)
  let lapsed = Bool -> Bool
not (Status n -> Bool
forall (n :: * -> *). Status n -> Bool
leaseLapsed Status n
status) Bool -> Bool -> Bool
&& Time
woke Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
>= Maybe InFlight -> Guarded n job -> Status n -> Time
forall (n :: * -> *) job.
Maybe InFlight -> Guarded n job -> Status n -> Time
leaseFenceAt Maybe InFlight
inFlight Guarded n job
entry Status n
status
  when lapsed (writeTVar (guardedStatus entry) status {leaseLapsed = True})
  pure (lapsed, not (leaseLapsed status || lapsed))

durationMessage :: HeartbeatGuard n job -> Text
durationMessage :: forall (n :: * -> *) job. HeartbeatGuard n job -> Text
durationMessage HeartbeatGuard n job
guard =
  Text
"claim hooks and handler ran past the maximum job duration"
    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> (DiffTime -> Text) -> Maybe DiffTime -> 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) -> (DiffTime -> Text) -> DiffTime -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DiffTime -> Text
forall a. Show a => a -> Text
showText) (GuardConfig n job -> Maybe DiffTime
forall (n :: * -> *) job. GuardConfig n job -> Maybe DiffTime
configMaxDuration (HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig HeartbeatGuard n job
guard))

-- ---------------------------------------------------------------------------
-- The extend in flight
-- ---------------------------------------------------------------------------

-- | Issue one extend over @due@, bounded by the earliest lease among @leased@.
issue
  :: (MonadFork n, MonadMask n, MonadTime n, MonadTimer n)
  => HeartbeatGuard n job
  -> Time
  -> [(Guarded n job, Status n)]
  -> [Guarded n job]
  -> n ()
issue :: forall (n :: * -> *) job.
(MonadFork n, MonadMask n, MonadTime n, MonadTimer n) =>
HeartbeatGuard n job
-> Time -> [(Guarded n job, Status n)] -> [Guarded n job] -> n ()
issue HeartbeatGuard n job
guard Time
woke [(Guarded n job, Status n)]
leased [Guarded n job]
due = do
  issued <- n Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
  atomically $ do
    writeTVar
      (guardInFlight guard)
      ( Just
          InFlight
            { issuedAt = issued
            , givesUp = addTime bound issued
            , carries = Set.fromList (map guardedToken due)
            , landed = False
            }
      )
    traverse_ (\Guarded n job
entry -> 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) (\Status n
status -> Status n
status {recheckAsked = False})) due
  void . forkIO $ (extend guard issued bound due `finally` finish guard issued) >>= traverse_ (report guard due)
  where
    bound :: DiffTime
bound =
      DiffTime -> DiffTime -> DiffTime
forall a. Ord a => a -> a -> a
max
        DiffTime
minRetryPause
        ([DiffTime] -> DiffTime
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (GuardConfig n job -> DiffTime
forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configTimeout (HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig HeartbeatGuard n job
guard) DiffTime -> [DiffTime] -> [DiffTime]
forall a. a -> [a] -> [a]
: [Status n -> Time
forall (n :: * -> *). Status n -> Time
leaseAt Status n
status Time -> Time -> DiffTime
`diffTime` Time
woke | (Guarded n job
_, Status n
status) <- [(Guarded n job, Status n)]
leased]))

-- | Abandon the extend in flight past its give-up and settle grace.
abandon :: (MonadSTM n) => HeartbeatGuard n job -> Time -> STM n ()
abandon :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> STM n ()
abandon HeartbeatGuard n job
guard Time
woke = do
  inFlight <- TVar n (Maybe InFlight) -> STM n (Maybe InFlight)
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 (Maybe InFlight)
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe InFlight)
guardInFlight HeartbeatGuard n job
guard)
  abandoned <- readTVar (guardAbandoned guard)
  for_ ((,) <$> inFlight <*> abandonAt abandoned inFlight) $ \(InFlight
running, Time
at) ->
    Bool -> STM n () -> STM n ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Time
woke Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
>= Time
at) (STM n () -> STM n ()) -> STM n () -> STM n ()
forall a b. (a -> b) -> a -> b
$ do
      TVar n (Maybe InFlight) -> Maybe InFlight -> STM n ()
forall a. TVar n a -> a -> STM n ()
forall (m :: * -> *) a. MonadSTM m => TVar m a -> a -> STM m ()
writeTVar (HeartbeatGuard n job -> TVar n (Maybe InFlight)
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe InFlight)
guardInFlight HeartbeatGuard n job
guard) Maybe InFlight
forall a. Maybe a
Nothing
      TVar n (Maybe Time) -> Maybe Time -> STM n ()
forall a. TVar n a -> a -> STM n ()
forall (m :: * -> *) a. MonadSTM m => TVar m a -> a -> STM m ()
writeTVar (HeartbeatGuard n job -> TVar n (Maybe Time)
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe Time)
guardAbandoned HeartbeatGuard n job
guard) (Time -> Maybe Time
forall a. a -> Maybe a
Just (InFlight -> Time
issuedAt InFlight
running))

-- | When the extend in flight is abandoned. Nothing while another abandoned extend runs.
abandonAt :: Maybe Time -> Maybe InFlight -> Maybe Time
abandonAt :: Maybe Time -> Maybe InFlight -> Maybe Time
abandonAt Maybe Time
abandoned Maybe InFlight
inFlight = case (Maybe Time
abandoned, Maybe InFlight
inFlight) of
  (Maybe Time
Nothing, Just InFlight
running) -> Time -> Maybe Time
forall a. a -> Maybe a
Just (DiffTime -> Time -> Time
addTime DiffTime
settleGrace (InFlight -> Time
givesUp InFlight
running))
  (Maybe Time, Maybe InFlight)
_ -> Maybe Time
forall a. Maybe a
Nothing

-- | The statement returned. The fence holds its batches until it is over.
land :: (MonadSTM n) => HeartbeatGuard n job -> Time -> n Bool
land :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> n Bool
land HeartbeatGuard n job
guard Time
issued = STM n Bool -> n Bool
forall a. HasCallStack => STM n a -> n a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM n Bool -> n Bool) -> STM n Bool -> n Bool
forall a b. (a -> b) -> a -> b
$ do
  current <- HeartbeatGuard n job -> Time -> STM n Bool
forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> STM n Bool
ownsAttempt HeartbeatGuard n job
guard Time
issued
  when current $ modifyTVar' (guardInFlight guard) (fmap (\InFlight
running -> InFlight
running {landed = True}))
  pure current

ownsAttempt :: (MonadSTM n) => HeartbeatGuard n job -> Time -> STM n Bool
ownsAttempt :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> STM n Bool
ownsAttempt HeartbeatGuard n job
guard Time
issued = Bool -> (InFlight -> Bool) -> Maybe InFlight -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False ((Time -> Time -> Bool
forall a. Eq a => a -> a -> Bool
== Time
issued) (Time -> Bool) -> (InFlight -> Time) -> InFlight -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InFlight -> Time
issuedAt) (Maybe InFlight -> Bool) -> STM n (Maybe InFlight) -> STM n Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar n (Maybe InFlight) -> STM n (Maybe InFlight)
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 (Maybe InFlight)
forall (n :: * -> *) job.
HeartbeatGuard n job -> TVar n (Maybe InFlight)
guardInFlight HeartbeatGuard n job
guard)

-- | The extend is over. The fence is due at the leases again.
finish :: (MonadSTM n) => HeartbeatGuard n job -> Time -> n ()
finish :: forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> n ()
finish HeartbeatGuard n job
guard Time
issued = 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 ()) -> STM n () -> n ()
forall a b. (a -> b) -> a -> b
$ do
  current <- HeartbeatGuard n job -> Time -> STM n Bool
forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> STM n Bool
ownsAttempt HeartbeatGuard n job
guard Time
issued
  when current (writeTVar (guardInFlight guard) Nothing *> wake guard)
  abandoned <- readTVar (guardAbandoned guard)
  when (abandoned == Just issued) (writeTVar (guardAbandoned guard) Nothing *> wake guard)

-- | One extend statement over every due batch, bounded by @bound@.
extend
  :: (MonadCatch n, MonadFork n, MonadTime n, MonadTimer n)
  => HeartbeatGuard n job
  -> Time
  -> DiffTime
  -> [Guarded n job]
  -> n (Maybe SomeException)
extend :: forall (n :: * -> *) job.
(MonadCatch n, MonadFork n, MonadTime n, MonadTimer n) =>
HeartbeatGuard n job
-> Time -> DiffTime -> [Guarded n job] -> n (Maybe SomeException)
extend HeartbeatGuard n job
guard Time
issued DiffTime
bound [Guarded n job]
due = do
  lives <- (Guarded n job -> n (Guarded n job, [job]))
-> [Guarded n job] -> n [(Guarded n job, [job])]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\Guarded n job
entry -> (,) Guarded n job
entry ([job] -> (Guarded n job, [job]))
-> n [job] -> n (Guarded n job, [job])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Guarded n job -> n [job]
forall (n :: * -> *) job. MonadSTM n => Guarded n job -> n [job]
pendingOf Guarded n job
entry) [Guarded n job]
due
  outcome <- timeout bound (trySync (configExtend config (concatMap snd lives)))
  case outcome of
    Maybe (Either SomeException [SetVisibilityResult])
Nothing -> Maybe SomeException
forall a. Maybe a
Nothing Maybe SomeException -> n () -> n (Maybe SomeException)
forall a b. a -> n b -> n a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (Guarded n job -> n ()) -> [Guarded n job] -> n ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (HeartbeatGuard n job -> Time -> Guarded n job -> n ()
forall (n :: * -> *) job.
(MonadMonotonicTime n, MonadSTM n) =>
HeartbeatGuard n job -> Time -> Guarded n job -> n ()
retryLater HeartbeatGuard n job
guard Time
issued) [Guarded n job]
due
    Just (Left SomeException
exception) -> SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
exception Maybe SomeException -> n () -> n (Maybe SomeException)
forall a b. a -> n b -> n a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (Guarded n job -> n ()) -> [Guarded n job] -> n ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (HeartbeatGuard n job -> Time -> Guarded n job -> n ()
forall (n :: * -> *) job.
(MonadMonotonicTime n, MonadSTM n) =>
HeartbeatGuard n job -> Time -> Guarded n job -> n ()
retryLater HeartbeatGuard n job
guard Time
issued) [Guarded n job]
due
    Just (Right [SetVisibilityResult]
results) -> do
      current <- HeartbeatGuard n job -> Time -> n Bool
forall (n :: * -> *) job.
MonadSTM n =>
HeartbeatGuard n job -> Time -> n Bool
land HeartbeatGuard n job
guard Time
issued
      when current $ do
        configExtended config
        currentTime <- getCurrentTime
        let byJob = [(Int64, SetVisibilityResult)] -> Map Int64 SetVisibilityResult
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(SetVisibilityResult -> Int64
resultId SetVisibilityResult
result, SetVisibilityResult
result) | SetVisibilityResult
result <- [SetVisibilityResult]
results]
        traverse_ (settle guard issued currentTime byJob) lives
      pure Nothing
  where
    config :: GuardConfig n job
config = HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig HeartbeatGuard n job
guard

-- | Log a refused extend.
report :: (Applicative n) => HeartbeatGuard n job -> [Guarded n job] -> SomeException -> n ()
report :: forall (n :: * -> *) job.
Applicative n =>
HeartbeatGuard n job -> [Guarded n job] -> SomeException -> n ()
report HeartbeatGuard n job
guard [Guarded n job]
due SomeException
exception =
  [Guarded n job] -> (Guarded n job -> n ()) -> n ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Guarded n job]
due ((Guarded n job -> n ()) -> n ())
-> (Guarded n job -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \Guarded n job
entry ->
    GuardConfig n job -> LogLevel -> [job] -> Text -> n ()
forall (n :: * -> *) job.
GuardConfig n job -> LogLevel -> [job] -> Text -> n ()
configLog
      (HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig HeartbeatGuard n job
guard)
      LogLevel
Error
      (NonEmpty job -> [job]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Batch n job -> NonEmpty job
forall (n :: * -> *) job. Batch n job -> NonEmpty job
batchJobs (Guarded n job -> Batch n job
forall (n :: * -> *) job. Guarded n job -> Batch n job
guardedBatch Guarded n job
entry)))
      (Text
"Heartbeat error (retrying): " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SomeException -> Text
displayEx SomeException
exception)

-- | Beat again after a failed extend.
retryLater :: (MonadMonotonicTime n, MonadSTM n) => HeartbeatGuard n job -> Time -> Guarded n job -> n ()
retryLater :: forall (n :: * -> *) job.
(MonadMonotonicTime n, MonadSTM n) =>
HeartbeatGuard n job -> Time -> Guarded n job -> n ()
retryLater HeartbeatGuard n job
guard Time
issued Guarded n job
entry = do
  now <- n Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
  atomically $ do
    current <- ownsAttempt guard issued
    when current $ modifyTVar' (guardedStatus entry) $ \Status n
status ->
      Time -> Status n -> Status n
forall (n :: * -> *). Time -> Status n -> Status n
rebeat (DiffTime -> Time -> Time
addTime (DiffTime -> Bool -> DiffTime -> DiffTime
heartbeatWait (GuardConfig n job -> DiffTime
forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configInterval (HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig HeartbeatGuard n job
guard)) Bool
False (Status n -> Time
forall (n :: * -> *). Status n -> Time
leaseAt Status n
status Time -> Time -> DiffTime
`diffTime` Time
now)) Time
now) Status n
status

-- | @try@ for synchronous exceptions only.
trySync :: (MonadCatch n) => n a -> n (Either SomeException a)
trySync :: forall (n :: * -> *) a.
MonadCatch n =>
n a -> n (Either SomeException a)
trySync = (SomeException -> Maybe SomeException)
-> n a -> n (Either SomeException a)
forall e b a.
Exception e =>
(e -> Maybe b) -> n a -> n (Either b a)
forall (m :: * -> *) e b a.
(MonadCatch m, Exception e) =>
(e -> Maybe b) -> m a -> m (Either b a)
tryJust (\SomeException
exc -> if SomeException -> Bool
forall e. Exception e => e -> Bool
isSyncException SomeException
exc then SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
exc else Maybe SomeException
forall a. Maybe a
Nothing)

-- ---------------------------------------------------------------------------
-- Settling
-- ---------------------------------------------------------------------------

-- | Act on one batch's verdicts from the extend. A stopped batch takes none.
settle
  :: (MonadFork n, MonadMonotonicTime n, MonadSTM n)
  => HeartbeatGuard n job
  -> Time
  -> UTCTime
  -> Map JobId SetVisibilityResult
  -> (Guarded n job, [job])
  -> n ()
settle :: forall (n :: * -> *) job.
(MonadFork n, MonadMonotonicTime n, MonadSTM n) =>
HeartbeatGuard n job
-> Time
-> UTCTime
-> Map Int64 SetVisibilityResult
-> (Guarded n job, [job])
-> n ()
settle HeartbeatGuard n job
guard Time
issued UTCTime
currentTime Map Int64 SetVisibilityResult
byJob (Guarded n job
entry, [job]
live) = do
  -- Rows this worker settled during the statement do not count.
  stillPending <- [Int64] -> Set Int64
forall a. Ord a => [a] -> Set a
Set.fromList ([Int64] -> Set Int64) -> ([job] -> [Int64]) -> [job] -> Set Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (job -> Int64) -> [job] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map job -> Int64
key ([job] -> Set Int64) -> n [job] -> n (Set Int64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Guarded n job -> n [job]
forall (n :: * -> *) job. MonadSTM n => Guarded n job -> n [job]
pendingOf Guarded n job
entry
  now <- getMonotonicTime
  let pendingLive = (job -> Bool) -> [job] -> [job]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int64 -> Set Int64 -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Int64
stillPending) (Int64 -> Bool) -> (job -> Int64) -> job -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. job -> Int64
key) [job]
live
      verdictOf job
job = Int64 -> Map Int64 SetVisibilityResult -> Maybe SetVisibilityResult
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (job -> Int64
key job
job) Map Int64 SetVisibilityResult
byJob
      verdicts = (job -> Maybe SetVisibilityResult)
-> [job] -> [SetVisibilityResult]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe job -> Maybe SetVisibilityResult
verdictOf [job]
pendingLive
      cancelledJobs = [Int64
jobId | JobCancelled Int64
jobId <- [SetVisibilityResult]
verdicts]
      stolenJobs = [Int64
jobId | JobReclaimed Int64
jobId Int64
_ Int64
_ <- [SetVisibilityResult]
verdicts]
      gone = [Int64] -> Set Int64
forall a. Ord a => [a] -> Set a
Set.fromList [Int64
jobId | JobGone Int64
jobId <- [SetVisibilityResult]
verdicts]
      isRenewed = (SetVisibilityResult -> Bool) -> Maybe SetVisibilityResult -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any SetVisibilityResult -> Bool
renewed (Maybe SetVisibilityResult -> Bool)
-> (job -> Maybe SetVisibilityResult) -> job -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. job -> Maybe SetVisibilityResult
verdictOf
      extended = (job -> Bool) -> [job] -> [job]
forall a. (a -> Bool) -> [a] -> [a]
filter job -> Bool
isRenewed [job]
pendingLive
      leaseFrom Status n
status
        | (job -> Bool) -> [job] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all job -> Bool
isRenewed [job]
pendingLive = DiffTime -> Time -> Time
addTime (GuardConfig n job -> DiffTime
forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configTimeout GuardConfig n job
config) Time
issued
        | Bool
otherwise = Status n -> Time
forall (n :: * -> *). Status n -> Time
leaseAt Status n
status
      cadence Time
lease = DiffTime -> Time -> Time
addTime (DiffTime -> Bool -> DiffTime -> DiffTime
heartbeatWait (GuardConfig n job -> DiffTime
forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configInterval GuardConfig n job
config) Bool
True (Time
lease Time -> Time -> DiffTime
`diffTime` Time
issued)) Time
issued
      -- A first sighting can be an ack not yet recorded. The next extend confirms it.
      beatFrom Status n
status
        | Set Int64
gone Set Int64 -> Set Int64 -> Bool
forall a. Ord a => Set a -> Set a -> Bool
`Set.isSubsetOf` Status n -> Set Int64
forall (n :: * -> *). Status n -> Set Int64
goneSeen Status n
status = Time -> Time
cadence (Status n -> Time
leaseFrom Status n
status)
        | Bool
otherwise = Time -> Time -> Time
forall a. Ord a => a -> a -> a
min (Time -> Time
cadence (Status n -> Time
leaseFrom Status n
status)) (DiffTime -> Time -> Time
addTime DiffTime
minRetryPause Time
now)
      renew Status n
status = Time -> Status n -> Status n
forall (n :: * -> *). Time -> Status n -> Status n
rebeat (Status n -> Time
beatFrom Status n
status) Status n
status {leaseAt = leaseFrom status, goneSeen = gone}
      stopFor Set Int64
seen
        | Bool -> Bool
not ([Int64] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int64]
cancelledJobs) = SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (JobForceCancelled -> SomeException
forall e. Exception e => e -> SomeException
toException ([Int64] -> [Int64] -> JobForceCancelled
JobForceCancelled [Int64]
cancelledJobs ([Int64]
stolenJobs [Int64] -> [Int64] -> [Int64]
forall a. Semigroup a => a -> a -> a
<> Set Int64 -> [Int64]
forall a. Set a -> [a]
Set.toList Set Int64
gone)))
        | Bool -> Bool
not ([Int64] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int64]
stolenJobs) = SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (JobGoneException -> SomeException
forall e. Exception e => e -> SomeException
toException (Text -> [Int64] -> JobGoneException
JobGoneException Text
reclaimedReason [Int64]
stolenJobs))
        | Bool -> Bool
not (Set Int64 -> Bool
forall a. Set a -> Bool
Set.null Set Int64
deleted) = SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (JobGoneException -> SomeException
forall e. Exception e => e -> SomeException
toException (Text -> [Int64] -> JobGoneException
JobGoneException Text
deletedReason (Set Int64 -> [Int64]
forall a. Set a -> [a]
Set.toList Set Int64
deleted)))
        | Bool
otherwise = Maybe SomeException
forall a. Maybe a
Nothing
        where
          deleted :: Set Int64
deleted = Set Int64 -> Set Int64 -> Set Int64
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection Set Int64
gone Set Int64
seen
  -- What the last extend found gone. Nothing for a batch the fence stopped.
  seenBefore <- atomically $ do
    current <- ownsAttempt guard issued
    stateTVar (guardedStatus entry) $ \Status n
status ->
      if Bool
current Bool -> Bool -> Bool
&& Bool -> Bool
not (Status n -> Bool
forall (n :: * -> *). Status n -> Bool
leaseLapsed Status n
status) then (Set Int64 -> Maybe (Set Int64)
forall a. a -> Maybe a
Just (Status n -> Set Int64
forall (n :: * -> *). Status n -> Set Int64
goneSeen Status n
status), Status n -> Status n
renew Status n
status) else (Maybe (Set Int64)
forall a. Maybe a
Nothing, Status n
status)
  for_ seenBefore (maybe (heartbeats extended) (signal guard now entry) . stopFor)
  where
    heartbeats :: [job] -> n ()
heartbeats [job]
extended =
      Bool -> n () -> n ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([job] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [job]
extended) (n () -> n ()) -> (n () -> n ()) -> n () -> n ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n (ThreadId n) -> n ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (n (ThreadId n) -> n ())
-> (n () -> n (ThreadId n)) -> n () -> n ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n () -> n (ThreadId n)
forall (m :: * -> *). MonadFork m => m () -> m (ThreadId m)
forkIO (n () -> n (ThreadId n))
-> (n () -> n ()) -> n () -> n (ThreadId n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Batch n job -> n () -> n ()
forall (n :: * -> *) job. Batch n job -> n () -> n ()
batchInherit Batch n job
batch (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
        [job] -> (job -> n ()) -> n ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [job]
extended ((job -> n ()) -> n ()) -> (job -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \job
job -> do
          pending <- Guarded n job -> n [job]
forall (n :: * -> *) job. MonadSTM n => Guarded n job -> n [job]
pendingOf Guarded n job
entry
          when (any ((== key job) . key) pending) $
            configHeartbeat config job currentTime (batchStart batch)
    config :: GuardConfig n job
config = HeartbeatGuard n job -> GuardConfig n job
forall (n :: * -> *) job. HeartbeatGuard n job -> GuardConfig n job
guardConfig HeartbeatGuard n job
guard
    key :: job -> Int64
key = GuardConfig n job -> job -> Int64
forall (n :: * -> *) job. GuardConfig n job -> job -> Int64
configKey GuardConfig n job
config
    batch :: Batch n job
batch = Guarded n job -> Batch n job
forall (n :: * -> *) job. Guarded n job -> Batch n job
guardedBatch Guarded n job
entry

-- | Only a successful extend renews the lease.
renewed :: SetVisibilityResult -> Bool
renewed :: SetVisibilityResult -> Bool
renewed VisibilityExtended {} = Bool
True
renewed SetVisibilityResult
_ = Bool
False

resultId :: SetVisibilityResult -> JobId
resultId :: SetVisibilityResult -> Int64
resultId SetVisibilityResult
result = case SetVisibilityResult
result of
  VisibilityExtended Int64
jobId -> Int64
jobId
  VisibilityUnchanged Int64
jobId -> Int64
jobId
  JobCancelled Int64
jobId -> Int64
jobId
  JobReclaimed Int64
jobId Int64
_ Int64
_ -> Int64
jobId
  JobGone Int64
jobId -> Int64
jobId
  JobSuspended Int64
jobId -> Int64
jobId