{-# LANGUAGE OverloadedStrings #-}

-- | Spawn pool threads that survive transient failures.
module Arbiter.Worker.Retry
  ( spawnRetried
  ) where

import Arbiter.Core.Exceptions (displayEx)
import Arbiter.Core.Threads (labelArbiterThread, micros)
import Control.Monad (unless)
import Control.Monad.Trans.Cont (ContT (..))
import Data.Text (Text)
import Data.Time (NominalDiffTime)
import UnliftIO (MonadUnliftIO, liftIO, tryAny)
import UnliftIO.Async (Async, race, withAsync)
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.STM (TVar, atomically, checkSTM)

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

-- | Wait between attempts of a retried infrastructure thread.
retryBackoff :: NominalDiffTime
retryBackoff :: NominalDiffTime
retryBackoff = NominalDiffTime
5

-- | Spawn a thread that logs and retries transient failures until the pool shuts down.
spawnRetried
  :: (MonadUnliftIO m)
  => TVar WorkerState
  -> LogConfig
  -> Text
  -- ^ The queue this thread serves, for its RTS label.
  -> 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)
    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 <- STM Bool -> m Bool
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar WorkerState -> STM Bool
isShuttingDown TVar WorkerState
stateVar)
      unless stopping $ do
        tryLog logCfg Error $ label <> " error (retrying): " <> displayEx exception
        -- Shutdown wins the race.
        race awaitShutdown (liftIO (threadDelay (micros retryBackoff)))
          >>= either pure (const loop)
    awaitShutdown :: m ()
awaitShutdown = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar WorkerState -> STM Bool
isShuttingDown TVar WorkerState
stateVar STM Bool -> (Bool -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM)