content: module Cachix.Client.Retry
    ( retryAll,
      retryAllWithPolicy,
      retryHttp,
      retryHttpWith,
      endlessRetryPolicy,
      endlessConstantRetryPolicy,
    )
  where
  
  import Cachix.Client.Exception (CachixException (..))
  import Control.Concurrent.Async qualified as Async
  import Control.Exception.Safe (Handler (..))
  import Control.Monad.Catch (MonadCatch, MonadMask, handleJust, throwM)
  import Control.Retry
  import Data.List (lookup)
  import Data.Time
    ( UTCTime,
      defaultTimeLocale,
      diffUTCTime,
      getCurrentTime,
      nominalDiffTimeToSeconds,
      readPTime,
      rfc822DateFormat,
    )
  import GHC.Read (Read (readPrec))
  import Network.HTTP.Client qualified as HTTP
  import Network.HTTP.Types.Header (hRetryAfter)
  import Network.HTTP.Types.Status qualified as HTTP
  import Protolude hiding (Handler (..), handleJust)
  import Servant.Client (ClientError (..))
  import Servant.Client qualified as Servant
  import Text.ParserCombinators.ReadPrec qualified as ReadPrec (lift)
  
  defaultRetryPolicy :: RetryPolicy
  defaultRetryPolicy =
    exponentialBackoff (1000 * 1000) <> limitRetries 5
  
  endlessConstantRetryPolicy :: RetryPolicy
  endlessConstantRetryPolicy =
    constantDelay (1000 * 1000)
  
  endlessRetryPolicy :: RetryPolicy
  endlessRetryPolicy =
    exponentialBackoff (1000 * 1000)
  
  retryAll :: (MonadIO m, MonadMask m) => (RetryStatus -> m a) -> m a
  retryAll = retryAllWithPolicy defaultRetryPolicy
  
  retryAllWithPolicy :: (MonadIO m, MonadMask m) => RetryPolicyM m -> (RetryStatus -> m a) -> m a
  retryAllWithPolicy policy f =
    recovering policy handlers $
      rethrowLinkedThreadExceptions . f
    where
      handlers = skipAsyncExceptions ++ [exitCodeHandler, cachixExceptionsHandler, allHandler]
  
      -- Skip over exitSuccess/exitFailure
      exitCodeHandler _ = Handler $ \(_ :: ExitCode) -> return False
  
      -- Skip over fatal Cachix exceptions
      cachixExceptionsHandler _ = Handler $ \(_ :: CachixException) -> return False
  
      -- Retry everything else
      allHandler _ = Handler $ \(_ :: SomeException) -> return True
  
  -- | Unwrap 'Async.ExceptionInLinkedThread' exceptions and rethrow the inner exception.
  rethrowLinkedThreadExceptions :: (MonadCatch m) => m a -> m a
  rethrowLinkedThreadExceptions =
    handleJust unwrapLinkedThreadException throwM
  
  unwrapLinkedThreadException :: SomeException -> Maybe SomeException
  unwrapLinkedThreadException e
    | Just (Async.ExceptionInLinkedThread _ e') <- fromException e = Just e'
    | otherwise = Nothing
  
  retryHttp :: (MonadIO m, MonadMask m) => m a -> m a
  retryHttp = retryHttpWith defaultRetryPolicy
  
  -- | Retry policy for HTTP requests.
  --
  -- Retries a subset of HTTP exceptions and overrides the delay with the Retry-After header if present.
  retryHttpWith :: forall m a. (MonadIO m, MonadMask m) => RetryPolicyM m -> m a -> m a
  retryHttpWith policy = recoveringDynamic policy handlers . const
    where
      handlers :: [RetryStatus -> Handler m RetryAction]
      handlers =
        skipAsyncExceptions' ++ [retryHttpExceptions, retryClientExceptions]
  
      skipAsyncExceptions' = map (fmap toRetryAction .) skipAsyncExceptions
  
      retryHttpExceptions _ = Handler httpExceptionToRetryHandler
      retryClientExceptions _ = Handler clientExceptionToRetryHandler
  
      httpExceptionToRetryHandler :: HTTP.HttpException -> m RetryAction
      httpExceptionToRetryHandler (HTTP.HttpExceptionRequest _ (HTTP.StatusCodeException response _))
        | statusMayHaveRetryHeader (HTTP.responseStatus response) = overrideDelayWithRetryAfter response
      httpExceptionToRetryHandler ex = return . toRetryAction . shouldRetryHttpException $ ex
  
      clientExceptionToRetryHandler :: ClientError -> m RetryAction
      clientExceptionToRetryHandler (FailureResponse _req res)
        | shouldRetryHttpStatusCode (Servant.responseStatusCode res) =
            return ConsultPolicy
      clientExceptionToRetryHandler (ConnectionError ex) =
        case fromException ex of
          Just httpException -> httpExceptionToRetryHandler httpException
          Nothing -> return DontRetry
      clientExceptionToRetryHandler _ = return DontRetry
  
  data RetryAfter
    = RetryAfterDate UTCTime
    | RetryAfterSeconds Int
    deriving (Eq, Show)
  
  instance Read RetryAfter where
    readPrec = parseSeconds <|> parseWebDate
      where
        parseSeconds = RetryAfterSeconds <$> readPrec
        parseWebDate = ReadPrec.lift $ RetryAfterDate <$> readPTime True defaultTimeLocale rfc822DateFormat
  
  overrideDelayWithRetryAfter :: (MonadIO m) => HTTP.Response a -> m RetryAction
  overrideDelayWithRetryAfter response =
    case lookupRetryAfter response of
      Nothing ->
        return ConsultPolicy
      Just (RetryAfterSeconds seconds) ->
        return $ ConsultPolicyOverrideDelay (seconds * 1000 * 1000)
      Just (RetryAfterDate date) -> do
        seconds <- secondsFromNow date
        return $
          if seconds > 0
            then ConsultPolicyOverrideDelay (seconds * 1000 * 1000)
            else ConsultPolicy
    where
      secondsFromNow date = do
        now <- liftIO getCurrentTime
        return $ ceiling $ nominalDiffTimeToSeconds (date `diffUTCTime` now)
  
      lookupRetryAfter = readMaybe . decodeUtf8 <=< lookup hRetryAfter . HTTP.responseHeaders
  
  -- | Determine whether the HTTP exception is worth retrying.
  --
  -- Temporary connection or network transfer errors are good candidates.
  shouldRetryHttpException :: HTTP.HttpException -> Bool
  shouldRetryHttpException (HTTP.InvalidUrlException _ _) = False
  shouldRetryHttpException (HTTP.HttpExceptionRequest _ reason) =
    case reason of
      HTTP.ConnectionClosed -> True
      HTTP.ConnectionFailure _ -> True
      HTTP.ConnectionTimeout -> True
      HTTP.IncompleteHeaders -> True
      HTTP.InternalException _ -> True
      HTTP.InvalidChunkHeaders -> True
      HTTP.InvalidProxyEnvironmentVariable _ _ -> True
      HTTP.InvalidStatusLine _ -> True
      HTTP.NoResponseDataReceived -> True
      HTTP.ProxyConnectException _ _ status
        | HTTP.statusIsServerError status -> True
      HTTP.ResponseBodyTooShort _ _ -> True
      HTTP.ResponseTimeout -> True
      HTTP.StatusCodeException response _ ->
        shouldRetryHttpStatusCode (HTTP.responseStatus response)
      HTTP.HttpZlibException _ -> True
      _ -> False
  
  -- | Determine whether the HTTP status code is worth retrying.
  shouldRetryHttpStatusCode :: HTTP.Status -> Bool
  shouldRetryHttpStatusCode code | code == HTTP.tooManyRequests429 = True
  shouldRetryHttpStatusCode code | HTTP.statusIsServerError code = True
  shouldRetryHttpStatusCode _ = False
  
  statusMayHaveRetryHeader :: HTTP.Status -> Bool
  statusMayHaveRetryHeader = flip elem [HTTP.tooManyRequests429, HTTP.serviceUnavailable503]
  