{-# LANGUAGE OverloadedStrings #-}

-- | Retry combinators for worker infrastructure threads (notification listener,
-- cron scheduler, etc.) that should survive transient database failures.
module Arbiter.Worker.Retry
  ( isJobSignal
  , retryOnException
  , spawnRetried
  ) where

import Arbiter.Core.Exceptions (JobException, JobGoneException, displayEx)
import Arbiter.Core.Threads (labelArbiterThread)
import Control.Monad (unless)
import Control.Monad.Trans.Cont (ContT (..))
import Data.Maybe (isJust)
import Data.Text qualified as T
import UnliftIO (MonadUnliftIO, SomeException, fromException, liftIO)
import UnliftIO.Async (Async, race, withAsync)
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.Exception (tryAny)
import UnliftIO.STM (TVar, atomically, readTVar, readTVarIO, retrySTM)

import Arbiter.Worker.Logger (LogConfig, LogLevel (..), tryLog)
import Arbiter.Worker.WorkerState (WorkerState (..))

-- | Run an action in a retry loop that survives transient failures, logging each one
-- and waiting five seconds before the next attempt. Exits cleanly once the pool is
-- shutting down.
retryOnException
  :: (MonadUnliftIO m)
  => TVar WorkerState
  -> LogConfig
  -> T.Text
  -- ^ Label for log messages (e.g. "Notification listener")
  -> m ()
  -- ^ Action to run
  -> m ()
retryOnException :: forall (m :: * -> *).
MonadUnliftIO m =>
TVar WorkerState -> LogConfig -> Text -> m () -> m ()
retryOnException TVar WorkerState
stateVar LogConfig
logCfg Text
label m ()
action = m ()
loop
  where
    loop :: m ()
loop = m () -> m (Either SomeException ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny m ()
action m (Either SomeException ())
-> (Either SomeException () -> 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
>>= (SomeException -> m ())
-> (() -> m ()) -> Either SomeException () -> m ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> m ()
onFailure () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    onFailure :: SomeException -> m ()
onFailure SomeException
exception = do
      stopping <- (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
ShuttingDown) (WorkerState -> Bool) -> m WorkerState -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar WorkerState -> m WorkerState
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO TVar WorkerState
stateVar
      unless stopping $ do
        tryLog logCfg Error $ label <> " error (retrying): " <> displayEx exception
        -- Shutdown wins the race.
        race awaitShutdown (liftIO (threadDelay retryBackoffMicros))
          >>= either pure (const loop)
    awaitShutdown :: m ()
awaitShutdown = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (STM () -> IO ()) -> STM () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      state <- TVar WorkerState -> STM WorkerState
forall a. TVar a -> STM a
readTVar TVar WorkerState
stateVar
      unless (state == ShuttingDown) retrySTM

-- | Wait between attempts of a retried infrastructure thread.
retryBackoffMicros :: Int
retryBackoffMicros :: Int
retryBackoffMicros = Int
5_000_000

-- | A signal that only the worker layer can act on. A retry loop rethrows it.
isJobSignal :: SomeException -> Bool
isJobSignal :: SomeException -> Bool
isJobSignal SomeException
exception =
  Maybe JobException -> Bool
forall a. Maybe a -> Bool
isJust (SomeException -> Maybe JobException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
exception :: Maybe JobException)
    Bool -> Bool -> Bool
|| Maybe JobGoneException -> Bool
forall a. Maybe a -> Bool
isJust (SomeException -> Maybe JobGoneException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
exception :: Maybe JobGoneException)

-- | Spawn a thread under 'withAsync' and 'retryOnException'. A transient failure
-- restarts the action.
spawnRetried
  :: (MonadUnliftIO m)
  => TVar WorkerState
  -> LogConfig
  -> T.Text
  -- ^ The queue this thread serves, for its RTS label.
  -> T.Text
  -- ^ Label for log messages, and the role in its RTS label.
  -> m ()
  -- ^ Action to run.
  -> ContT r m (Async ())
spawnRetried :: forall (m :: * -> *) r.
MonadUnliftIO m =>
TVar WorkerState
-> LogConfig -> Text -> Text -> m () -> ContT r m (Async ())
spawnRetried TVar WorkerState
stateVar LogConfig
logCfg Text
queue Text
label m ()
action =
  ((Async () -> m r) -> m r) -> ContT r m (Async ())
forall {k} (r :: k) (m :: k -> *) a.
((a -> m r) -> m r) -> ContT r m a
ContT (((Async () -> m r) -> m r) -> ContT r m (Async ()))
-> (m () -> (Async () -> m r) -> m r)
-> m ()
-> ContT r m (Async ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m () -> (Async () -> m r) -> m r
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync (m () -> ContT r m (Async ())) -> m () -> ContT r m (Async ())
forall a b. (a -> b) -> a -> b
$ do
    Text -> Maybe Text -> m ()
forall (m :: * -> *). MonadIO m => Text -> Maybe Text -> m ()
labelArbiterThread Text
label (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
queue)
    TVar WorkerState -> LogConfig -> Text -> m () -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
TVar WorkerState -> LogConfig -> Text -> m () -> m ()
retryOnException TVar WorkerState
stateVar LogConfig
logCfg Text
label m ()
action