module Stamina.HTTP (retry, handler) where

import Control.Applicative ((<|>))
import Control.Exception (SomeException, fromException)
import Control.Monad.Catch (MonadCatch)
import Control.Monad.IO.Class (MonadIO)
import Data.Time (UTCTime, defaultTimeLocale, readPTime, rfc822DateFormat, secondsToNominalDiffTime)
import Network.HTTP.Client qualified as HTTP
import Network.HTTP.Types.Header (hRetryAfter)
import Network.HTTP.Types.Status (statusIsServerError, tooManyRequests429)
import Stamina qualified
import Text.Read (Read (readPrec), readMaybe)
import Text.Read qualified as ReadPrec

-- | Retry handler for HTTP requests.
--
-- Retries a subset of HTTP exceptions and overrides the delay with the Retry-After header if present.
retry :: (MonadIO m, MonadCatch m) => Stamina.RetrySettings -> (Stamina.RetryStatus -> m a) -> m a
retry :: forall (m :: * -> *) a.
(MonadIO m, MonadCatch m) =>
RetrySettings -> (RetryStatus -> m a) -> m a
retry RetrySettings
settings = RetrySettings
-> (SomeException -> m RetryAction) -> (RetryStatus -> m a) -> m a
forall (m :: * -> *) exc a.
(Exception exc, MonadIO m, MonadCatch m) =>
RetrySettings
-> (exc -> m RetryAction) -> (RetryStatus -> m a) -> m a
Stamina.retryFor RetrySettings
settings SomeException -> m RetryAction
forall (m :: * -> *). MonadIO m => SomeException -> m RetryAction
handler

handler :: (MonadIO m) => SomeException -> m Stamina.RetryAction
handler :: forall (m :: * -> *). MonadIO m => SomeException -> m RetryAction
handler =
  Maybe HttpException -> m RetryAction
forall {m :: * -> *}.
Monad m =>
Maybe HttpException -> m RetryAction
httpExceptionToRetryAction (Maybe HttpException -> m RetryAction)
-> (SomeException -> Maybe HttpException)
-> SomeException
-> m RetryAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> Maybe HttpException
forall e. Exception e => SomeException -> Maybe e
fromException
  where
    -- httpExceptionToRetryAction :: Maybe HTTP.HttpException -> m Stamina.RetryAction
    httpExceptionToRetryAction :: Maybe HttpException -> m RetryAction
httpExceptionToRetryAction (Just exc :: HttpException
exc@(HTTP.HttpExceptionRequest Request
_ (HTTP.StatusCodeException Response ()
response ByteString
_))) = do
      case Response () -> Maybe RetryAfterHeader
forall body. Response body -> Maybe RetryAfterHeader
lookupRetryAfter Response ()
response of
        Just (RetryAfterSeconds Int
seconds) -> RetryAction -> m RetryAction
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (RetryAction -> m RetryAction) -> RetryAction -> m RetryAction
forall a b. (a -> b) -> a -> b
$ NominalDiffTime -> RetryAction
Stamina.RetryDelay (NominalDiffTime -> RetryAction) -> NominalDiffTime -> RetryAction
forall a b. (a -> b) -> a -> b
$ Pico -> NominalDiffTime
secondsToNominalDiffTime (Pico -> NominalDiffTime) -> Pico -> NominalDiffTime
forall a b. (a -> b) -> a -> b
$ Int -> Pico
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
seconds
        Just (RetryAfterDate UTCTime
date) -> RetryAction -> m RetryAction
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (RetryAction -> m RetryAction) -> RetryAction -> m RetryAction
forall a b. (a -> b) -> a -> b
$ UTCTime -> RetryAction
Stamina.RetryTime UTCTime
date
        Maybe RetryAfterHeader
Nothing ->
          if HttpException -> Bool
shouldRetryHttpException HttpException
exc
            then RetryAction -> m RetryAction
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return RetryAction
Stamina.Retry
            else RetryAction -> m RetryAction
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return RetryAction
Stamina.RaiseException
    httpExceptionToRetryAction (Just HttpException
exc) | HttpException -> Bool
shouldRetryHttpException HttpException
exc = RetryAction -> m RetryAction
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return RetryAction
Stamina.Retry
    httpExceptionToRetryAction Maybe HttpException
_ = RetryAction -> m RetryAction
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return RetryAction
Stamina.RaiseException

    -- https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
    lookupRetryAfter :: HTTP.Response body -> Maybe RetryAfterHeader
    lookupRetryAfter :: forall body. Response body -> Maybe RetryAfterHeader
lookupRetryAfter Response body
body = HeaderName -> [(HeaderName, ByteString)] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
hRetryAfter (Response body -> [(HeaderName, ByteString)]
forall body. Response body -> [(HeaderName, ByteString)]
HTTP.responseHeaders Response body
body) Maybe ByteString
-> (ByteString -> Maybe RetryAfterHeader) -> Maybe RetryAfterHeader
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> Maybe RetryAfterHeader
forall a. Read a => String -> Maybe a
readMaybe (String -> Maybe RetryAfterHeader)
-> (ByteString -> String) -> ByteString -> Maybe RetryAfterHeader
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
forall a. Show a => a -> String
show

data RetryAfterHeader
  = RetryAfterDate UTCTime
  | RetryAfterSeconds Int
  deriving (RetryAfterHeader -> RetryAfterHeader -> Bool
(RetryAfterHeader -> RetryAfterHeader -> Bool)
-> (RetryAfterHeader -> RetryAfterHeader -> Bool)
-> Eq RetryAfterHeader
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RetryAfterHeader -> RetryAfterHeader -> Bool
== :: RetryAfterHeader -> RetryAfterHeader -> Bool
$c/= :: RetryAfterHeader -> RetryAfterHeader -> Bool
/= :: RetryAfterHeader -> RetryAfterHeader -> Bool
Eq, Int -> RetryAfterHeader -> ShowS
[RetryAfterHeader] -> ShowS
RetryAfterHeader -> String
(Int -> RetryAfterHeader -> ShowS)
-> (RetryAfterHeader -> String)
-> ([RetryAfterHeader] -> ShowS)
-> Show RetryAfterHeader
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RetryAfterHeader -> ShowS
showsPrec :: Int -> RetryAfterHeader -> ShowS
$cshow :: RetryAfterHeader -> String
show :: RetryAfterHeader -> String
$cshowList :: [RetryAfterHeader] -> ShowS
showList :: [RetryAfterHeader] -> ShowS
Show)

instance Read RetryAfterHeader where
  readPrec :: ReadPrec RetryAfterHeader
readPrec = ReadPrec RetryAfterHeader
parseSeconds ReadPrec RetryAfterHeader
-> ReadPrec RetryAfterHeader -> ReadPrec RetryAfterHeader
forall a. ReadPrec a -> ReadPrec a -> ReadPrec a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ReadPrec RetryAfterHeader
parseWebDate
    where
      parseSeconds :: ReadPrec RetryAfterHeader
parseSeconds = Int -> RetryAfterHeader
RetryAfterSeconds (Int -> RetryAfterHeader)
-> ReadPrec Int -> ReadPrec RetryAfterHeader
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReadPrec Int
forall a. Read a => ReadPrec a
readPrec
      parseWebDate :: ReadPrec RetryAfterHeader
parseWebDate = ReadP RetryAfterHeader -> ReadPrec RetryAfterHeader
forall a. ReadP a -> ReadPrec a
ReadPrec.lift (ReadP RetryAfterHeader -> ReadPrec RetryAfterHeader)
-> ReadP RetryAfterHeader -> ReadPrec RetryAfterHeader
forall a b. (a -> b) -> a -> b
$ UTCTime -> RetryAfterHeader
RetryAfterDate (UTCTime -> RetryAfterHeader)
-> ReadP UTCTime -> ReadP RetryAfterHeader
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> TimeLocale -> String -> ReadP UTCTime
forall t. ParseTime t => Bool -> TimeLocale -> String -> ReadP t
readPTime Bool
True TimeLocale
defaultTimeLocale String
rfc822DateFormat

shouldRetryHttpException :: HTTP.HttpException -> Bool
shouldRetryHttpException :: HttpException -> Bool
shouldRetryHttpException (HTTP.InvalidUrlException String
_ String
_) = Bool
False
shouldRetryHttpException (HTTP.HttpExceptionRequest Request
_ HttpExceptionContent
reason) =
  case HttpExceptionContent
reason of
    HttpExceptionContent
HTTP.ConnectionClosed -> Bool
True
    HTTP.ConnectionFailure SomeException
_ -> Bool
True
    HttpExceptionContent
HTTP.ConnectionTimeout -> Bool
True
    HttpExceptionContent
HTTP.IncompleteHeaders -> Bool
True
    HTTP.InternalException SomeException
_ -> Bool
True
    HttpExceptionContent
HTTP.InvalidChunkHeaders -> Bool
True
    HTTP.InvalidProxyEnvironmentVariable Text
_ Text
_ -> Bool
True
    HTTP.InvalidStatusLine ByteString
_ -> Bool
True
    HttpExceptionContent
HTTP.NoResponseDataReceived -> Bool
True
    HTTP.ProxyConnectException ByteString
_ Int
_ Status
status
      | Status -> Bool
statusIsServerError Status
status -> Bool
True
    HTTP.ResponseBodyTooShort Word64
_ Word64
_ -> Bool
True
    HttpExceptionContent
HTTP.ResponseTimeout -> Bool
True
    HTTP.StatusCodeException Response ()
response ByteString
_
      | Response () -> Status
forall body. Response body -> Status
HTTP.responseStatus Response ()
response Status -> Status -> Bool
forall a. Eq a => a -> a -> Bool
== Status
tooManyRequests429 -> Bool
True
    HTTP.StatusCodeException Response ()
response ByteString
_
      | Status -> Bool
statusIsServerError (Response () -> Status
forall body. Response body -> Status
HTTP.responseStatus Response ()
response) -> Bool
True
    HTTP.HttpZlibException ZlibException
_ -> Bool
True
    HttpExceptionContent
_ -> Bool
False