module Freckle.App.Http.Retry
  ( RetriesExhausted(..)
  , rateLimited
  , rateLimited'
  ) where

import Freckle.App.Prelude

import Control.Retry
import qualified Data.ByteString.Char8 as BS8
import Network.HTTP.Client (Request(..))
import Network.HTTP.Simple
import Network.HTTP.Types.Status (status429)
import Text.Read (readMaybe)
import UnliftIO.Exception (Exception(..), throwIO)

-- | Thrown if we exhaust our retries limit and still see a @429@
--
-- This typically means the API is not sending back accurate @Retry-In@ values
-- with 429 responses.
--
-- __Rationale__:
--
-- In order for 'rateLimited' to function in the case when the 'Request' is
-- using 'throwErrorStatusCodes' for 'checkResponse', we have to modify it to
-- not throw on 429s specifically. Otherwise, the first response would just
-- throw due to 4XX and never retry. However, in that case of someone expecting
-- invalid statuses to throw an exception, if we exhaust our retries and still
-- see a 429 at the end, an exception should be thrown.
--
-- Unfortunately, it's not possible to reuse the user-defined 'checkResponse' in
-- order to throw a uniform 'HttpException' in this case; so we throw this
-- ourselves instead.
--
data RetriesExhausted = RetriesExhausted
  { RetriesExhausted -> Int
reLimit :: Int
  , RetriesExhausted -> Response ()
reResponse :: Response ()
  }
  deriving stock Int -> RetriesExhausted -> ShowS
[RetriesExhausted] -> ShowS
RetriesExhausted -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RetriesExhausted] -> ShowS
$cshowList :: [RetriesExhausted] -> ShowS
show :: RetriesExhausted -> String
$cshow :: RetriesExhausted -> String
showsPrec :: Int -> RetriesExhausted -> ShowS
$cshowsPrec :: Int -> RetriesExhausted -> ShowS
Show

instance Exception RetriesExhausted where
  displayException :: RetriesExhausted -> String
displayException RetriesExhausted {Int
Response ()
reResponse :: Response ()
reLimit :: Int
reResponse :: RetriesExhausted -> Response ()
reLimit :: RetriesExhausted -> Int
..} =
    String
"Retries exhaused after "
      forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Int
reLimit
      forall a. Semigroup a => a -> a -> a
<> String
" attempts. Final response:\n"
      forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Response ()
reResponse

rateLimited
  :: MonadIO m => (Request -> m (Response body)) -> Request -> m (Response body)
rateLimited :: forall (m :: * -> *) body.
MonadIO m =>
(Request -> m (Response body)) -> Request -> m (Response body)
rateLimited = forall (m :: * -> *) body.
MonadIO m =>
Int
-> (Request -> m (Response body)) -> Request -> m (Response body)
rateLimited' Int
10

-- | 'rateLimited' but with configurable retry limit
rateLimited'
  :: MonadIO m
  => Int
  -> (Request -> m (Response body))
  -> Request
  -> m (Response body)
rateLimited' :: forall (m :: * -> *) body.
MonadIO m =>
Int
-> (Request -> m (Response body)) -> Request -> m (Response body)
rateLimited' Int
retryLimit Request -> m (Response body)
f Request
req = do
  Response body
resp <- forall (m :: * -> *) b.
MonadIO m =>
RetryPolicyM m
-> (RetryStatus -> b -> m RetryAction)
-> (RetryStatus -> m b)
-> m b
retryingDynamic
    (Int -> RetryPolicy
limitRetries Int
retryLimit)
    (\RetryStatus
_ ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. b -> (a -> b) -> Maybe a -> b
maybe RetryAction
DontRetry (Int -> RetryAction
ConsultPolicyOverrideDelay forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
microseconds)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. Response body -> Maybe Int
getRetryAfter
    )
    (\RetryStatus
_ -> Request -> m (Response body)
f forall a b. (a -> b) -> a -> b
$ Request -> Request
suppressRetryStatusError Request
req)

  forall (m :: * -> *) body.
MonadIO m =>
Int -> Response body -> m (Response body)
checkRetriesExhausted Int
retryLimit Response body
resp

suppressRetryStatusError :: Request -> Request
suppressRetryStatusError :: Request -> Request
suppressRetryStatusError Request
req = Request
req
  { checkResponse :: Request -> Response BodyReader -> IO ()
checkResponse = \Request
req' Response BodyReader
resp ->
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall a. Response a -> Status
getResponseStatus Response BodyReader
resp forall a. Eq a => a -> a -> Bool
== Status
status429)
      forall a b. (a -> b) -> a -> b
$ Request -> Response BodyReader -> IO ()
originalCheckResponse Request
req' Response BodyReader
resp
  }
  where originalCheckResponse :: Request -> Response BodyReader -> IO ()
originalCheckResponse = Request -> Request -> Response BodyReader -> IO ()
checkResponse Request
req

checkRetriesExhausted :: MonadIO m => Int -> Response body -> m (Response body)
checkRetriesExhausted :: forall (m :: * -> *) body.
MonadIO m =>
Int -> Response body -> m (Response body)
checkRetriesExhausted Int
retryLimit Response body
resp
  | forall a. Response a -> Status
getResponseStatus Response body
resp forall a. Eq a => a -> a -> Bool
== Status
status429 = forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO
  forall a b. (a -> b) -> a -> b
$ RetriesExhausted { reLimit :: Int
reLimit = Int
retryLimit, reResponse :: Response ()
reResponse = forall (f :: * -> *) a. Functor f => f a -> f ()
void Response body
resp }
  | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure Response body
resp

getRetryAfter :: Response body -> Maybe Int
getRetryAfter :: forall body. Response body -> Maybe Int
getRetryAfter Response body
resp = do
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall a. Response a -> Status
getResponseStatus Response body
resp forall a. Eq a => a -> a -> Bool
== Status
status429
  ByteString
header <- forall a. [a] -> Maybe a
listToMaybe forall a b. (a -> b) -> a -> b
$ forall a. HeaderName -> Response a -> [ByteString]
getResponseHeader HeaderName
"Retry-After" Response body
resp
  forall a. Read a => String -> Maybe a
readMaybe forall a b. (a -> b) -> a -> b
$ ByteString -> String
BS8.unpack ByteString
header

microseconds :: Int -> Int
microseconds :: Int -> Int
microseconds = (forall a. Num a => a -> a -> a
* Int
1000000)