-- | The handler side of the guard: registering a batch, asking for its extend,
-- and asking it to stop.
--
-- A courier thread throws each signal. A masked handler does not stall the
-- guard. Unregister kills the couriers, which revokes a signal still in
-- flight, and drains one that already landed. A second signal within one beat
-- is dropped.
module Arbiter.Worker.Heartbeat.Guard.Signal
  ( guardBatch
  , recheck
  , signal
  ) where

import Arbiter.Core.Exceptions (JobGoneException (..))
import Arbiter.Core.Job.Types (JobId)
import Control.Concurrent.Class.MonadSTM (MonadSTM, atomically, modifyTVar', newTVarIO, readTVar, stateTVar)
import Control.Exception (asyncExceptionFromException, asyncExceptionToException)
import Control.Monad (unless, void, when)
import Control.Monad.Class.MonadFork (MonadFork (..), MonadThread (..))
import Control.Monad.Class.MonadThrow (Exception (..), MonadCatch (..), MonadMask (..), MonadThrow (..), SomeException)
import Control.Monad.Class.MonadTime.SI (MonadMonotonicTime (..), MonadTime (..), Time, addTime, diffTime, diffUTCTime)
import Data.Foldable (fold, toList, traverse_)
import Data.Map.Strict qualified as Map
import Data.Maybe (maybeToList)
import Data.Set qualified as Set

import Arbiter.Worker.Heartbeat.Guard.State

-- | A guard signal, rethrown as sync at the handler boundary.
newtype GuardSignal = GuardSignal SomeException
  deriving stock (Int -> GuardSignal -> ShowS
[GuardSignal] -> ShowS
GuardSignal -> String
(Int -> GuardSignal -> ShowS)
-> (GuardSignal -> String)
-> ([GuardSignal] -> ShowS)
-> Show GuardSignal
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> GuardSignal -> ShowS
showsPrec :: Int -> GuardSignal -> ShowS
$cshow :: GuardSignal -> String
show :: GuardSignal -> String
$cshowList :: [GuardSignal] -> ShowS
showList :: [GuardSignal] -> ShowS
Show)

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

-- | Run @action@ on the calling thread with the batch registered with the guard.
guardBatch
  :: (MonadFork n, MonadMask n, MonadMonotonicTime n, MonadSTM n, MonadTime n)
  => HeartbeatGuard n job
  -> Batch n job
  -> n a
  -> n a
{-# SPECIALIZE guardBatch :: HeartbeatGuard IO job -> Batch IO job -> IO a -> IO a #-}
guardBatch :: forall (n :: * -> *) job a.
(MonadFork n, MonadMask n, MonadMonotonicTime n, MonadSTM n,
 MonadTime n) =>
HeartbeatGuard n job -> Batch n job -> n a -> n a
guardBatch HeartbeatGuard n job
guard Batch n job
batch n a
action =
  n a -> n a
forall {m :: * -> *} {a}. MonadCatch m => m a -> m a
asSync (n (Guarded n job)
-> (Guarded n job -> n ()) -> (Guarded n job -> n a) -> n a
forall a b c. n a -> (a -> n b) -> (a -> n c) -> n c
forall (m :: * -> *) a b c.
MonadThrow m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket n (Guarded n job)
register Guarded n job -> n ()
unregister (n a -> Guarded n job -> n a
forall a b. a -> b -> a
const n a
action))
  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
    register :: n (Guarded n job)
register = do
      wallNow <- n UTCTime
forall (m :: * -> *). MonadTime m => m UTCTime
getCurrentTime
      monoNow <- getMonotonicTime
      handler <- myThreadId
      let elapsed = NominalDiffTime -> DiffTime
toDiffTime (UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
wallNow (Batch n job -> UTCTime
forall (n :: * -> *) job. Batch n job -> UTCTime
batchStart Batch n job
batch))
          -- The rows' own deadlines, which the claim set before the batch started.
          rowLeases =
            [ (job
job, DiffTime -> Time -> Time
addTime (NominalDiffTime -> DiffTime
toDiffTime (UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
at UTCTime
wallNow)) Time
monoNow)
            | job
job <- 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 Batch n job
batch)
            , Just UTCTime
at <- [GuardConfig n job -> job -> Maybe UTCTime
forall (n :: * -> *) job. GuardConfig n job -> job -> Maybe UTCTime
configLease GuardConfig n job
config job
job]
            ]
          leaseUntil = [Time] -> Time
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (DiffTime -> Time -> Time
addTime (GuardConfig n job -> DiffTime
forall (n :: * -> *) job. GuardConfig n job -> DiffTime
configTimeout GuardConfig n job
config DiffTime -> DiffTime -> DiffTime
forall a. Num a => a -> a -> a
- DiffTime
elapsed) Time
monoNow Time -> [Time] -> [Time]
forall a. a -> [a] -> [a]
: ((job, Time) -> Time) -> [(job, Time)] -> [Time]
forall a b. (a -> b) -> [a] -> [b]
map (job, Time) -> Time
forall a b. (a, b) -> b
snd [(job, Time)]
rowLeases)
          firstBeat = 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
leaseUntil Time -> Time -> DiffTime
`diffTime` Time
monoNow)) Time
monoNow
          deadline = (DiffTime -> Time -> Time
`addTime` Time
monoNow) (DiffTime -> Time) -> Maybe DiffTime -> Maybe Time
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GuardConfig n job -> Maybe DiffTime
forall (n :: * -> *) job. GuardConfig n job -> Maybe DiffTime
configMaxDuration GuardConfig n job
config
      -- A delayed claim response or queueing pause can outlive the row's lease.
      -- Refuse synchronously: an asynchronous fence could lose to a fast handler.
      -- Clock skew alone is not proof that this worker still owns the claim.
      when (leaseUntil <= monoNow) $
        throwIO (JobGoneException leaseExpiredReason (map (configKey config) (toList (batchJobs batch))))
      status <-
        newTVarIO
          Status
            { leaseAt = leaseUntil
            , beatAt = firstBeat
            , recheckAsked = False
            , leaseLapsed = False
            , deadlineSent = False
            , signalledAt = Nothing
            , goneSeen = Set.empty
            , couriers = Just []
            }
      atomically $ do
        token <- stateTVar (guardNextToken guard) (\Int
next -> (Int
next, Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        let entry =
              Guarded
                { guardedToken :: Int
guardedToken = Int
token
                , guardedBatch :: Batch n job
guardedBatch = Batch n job
batch
                , guardedHandler :: ThreadId n
guardedHandler = ThreadId n
handler
                , guardedDeadline :: Maybe Time
guardedDeadline = Maybe Time
deadline
                , guardedStatus :: TVar n (Status n)
guardedStatus = TVar n (Status n)
status
                }
        modifyTVar' (guardEntries guard) (Map.insert token entry)
        wakeFor guard (minimum (leaseUntil : firstBeat : maybeToList deadline))
        pure entry

    asSync :: m a -> m a
asSync m a
act = m a
act m a -> (GuardSignal -> m a) -> m a
forall e a. Exception e => m a -> (e -> m a) -> m a
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` (\(GuardSignal SomeException
exc) -> SomeException -> m a
forall e a. Exception e => e -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwIO SomeException
exc)

    unregister :: Guarded n job -> n ()
unregister Guarded n job
entry = do
      carrying <- n [ThreadId n] -> n [ThreadId n]
forall a. n a -> n a
forall (m :: * -> *) a. MonadMask m => m a -> m a
uninterruptibleMask_ (n [ThreadId n] -> n [ThreadId n])
-> n [ThreadId n] -> n [ThreadId n]
forall a b. (a -> b) -> a -> b
$ do
        carrying <- STM n [ThreadId n] -> n [ThreadId n]
forall a. HasCallStack => STM n a -> n a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM n [ThreadId n] -> n [ThreadId n])
-> STM n [ThreadId n] -> n [ThreadId n]
forall a b. (a -> b) -> a -> b
$ do
          TVar n (Map Int (Guarded n job))
-> (Map Int (Guarded n job) -> Map Int (Guarded n job)) -> STM n ()
forall a. TVar n a -> (a -> a) -> STM n ()
forall (m :: * -> *) a.
MonadSTM m =>
TVar m a -> (a -> a) -> STM m ()
modifyTVar' (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) (Int -> Map Int (Guarded n job) -> Map Int (Guarded n job)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (Guarded n job -> Int
forall (n :: * -> *) job. Guarded n job -> Int
guardedToken Guarded n job
entry))
          TVar n (Status n)
-> (Status n -> ([ThreadId n], Status n)) -> STM n [ThreadId n]
forall s a. TVar n s -> (s -> (a, s)) -> STM n a
forall (m :: * -> *) s a.
MonadSTM m =>
TVar m s -> (s -> (a, s)) -> STM m a
stateTVar (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 -> (Maybe [ThreadId n] -> [ThreadId n]
forall m. Monoid m => Maybe m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold (Status n -> Maybe [ThreadId n]
forall (n :: * -> *). Status n -> Maybe [ThreadId n]
couriers Status n
status), Status n
status {couriers = Nothing}))
        carrying <$ traverse_ killThread carrying
      -- A signal a courier had already sent lands here, inside the boundary, and is dropped.
      unless (null carrying) $ interruptible (pure ()) `catch` \GuardSignal {} -> () -> n ()
forall a. a -> n a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Extend the batch that holds the job now. The extend's verdict signals its handler.
recheck :: (MonadMonotonicTime n, MonadSTM n) => HeartbeatGuard n job -> JobId -> n ()
recheck :: forall (n :: * -> *) job.
(MonadMonotonicTime n, MonadSTM n) =>
HeartbeatGuard n job -> Int64 -> n ()
recheck HeartbeatGuard n job
guard Int64
jobId = do
  now <- n Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
  atomically $ do
    entries <- readTVar (guardEntries guard)
    let holding = (Guarded n job -> Bool) -> [Guarded n job] -> [Guarded n job]
forall a. (a -> Bool) -> [a] -> [a]
filter ((job -> Bool) -> NonEmpty job -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
jobId) (Int64 -> Bool) -> (job -> Int64) -> job -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HeartbeatGuard n job -> job -> Int64
forall (n :: * -> *) job. HeartbeatGuard n job -> job -> Int64
guardKey HeartbeatGuard n job
guard) (NonEmpty job -> Bool)
-> (Guarded n job -> NonEmpty job) -> Guarded n job -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Batch n job -> NonEmpty job
forall (n :: * -> *) job. Batch n job -> NonEmpty job
batchJobs (Batch n job -> NonEmpty job)
-> (Guarded n job -> Batch n job) -> Guarded n job -> NonEmpty job
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Guarded n job -> Batch n job
forall (n :: * -> *) job. Guarded n job -> Batch n job
guardedBatch) (Map Int (Guarded n job) -> [Guarded n job]
forall k a. Map k a -> [a]
Map.elems Map Int (Guarded n job)
entries)
    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 {beatAt = min now (beatAt status), recheckAsked = True}))
      holding
    unless (null holding) (wakeFor guard now)

-- | Ask the handler to stop, from a courier thread. A second ask within one beat is dropped.
signal :: (MonadFork n, MonadSTM n) => HeartbeatGuard n job -> Time -> Guarded n job -> SomeException -> n ()
signal :: forall (n :: * -> *) job.
(MonadFork n, MonadSTM n) =>
HeartbeatGuard n job
-> Time -> Guarded n job -> SomeException -> n ()
signal HeartbeatGuard n job
guard Time
now Guarded n job
entry SomeException
exc = do
  fresh <- 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
$ TVar n (Status n) -> (Status n -> (Bool, Status n)) -> STM n Bool
forall s a. TVar n s -> (s -> (a, s)) -> STM n a
forall (m :: * -> *) s a.
MonadSTM m =>
TVar m s -> (s -> (a, s)) -> STM m a
stateTVar (Guarded n job -> TVar n (Status n)
forall (n :: * -> *) job. Guarded n job -> TVar n (Status n)
guardedStatus Guarded n job
entry) ((Status n -> (Bool, Status n)) -> STM n Bool)
-> (Status n -> (Bool, Status n)) -> STM n Bool
forall a b. (a -> b) -> a -> b
$ \Status n
status ->
    let due :: Bool
due = Bool -> (Time -> Bool) -> Maybe Time -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (\Time
at -> Time
now Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
>= DiffTime -> Time -> Time
addTime (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)) Time
at) (Status n -> Maybe Time
forall (n :: * -> *). Status n -> Maybe Time
signalledAt Status n
status)
     in (Bool
due, if Bool
due then Status n
status {signalledAt = Just now} else Status n
status)
  when fresh . void . forkIO $ do
    courier <- myThreadId
    claimed <- atomically $ stateTVar (guardedStatus entry) $ \Status n
status -> case Status n -> Maybe [ThreadId n]
forall (n :: * -> *). Status n -> Maybe [ThreadId n]
couriers Status n
status of
      Maybe [ThreadId n]
Nothing -> (Bool
False, Status n
status)
      Just [ThreadId n]
carrying -> (Bool
True, Status n
status {couriers = Just (courier : carrying)})
    when claimed (throwTo (guardedHandler entry) (GuardSignal exc))