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

import Prelude

import Control.Exception.Annotated.UnliftIO (Exception (..), throwWithCallStack)
import Control.Monad (guard, unless)
import Control.Monad.IO.Class (MonadIO)
import Control.Retry
import Data.ByteString.Char8 qualified as BS8
import Data.Functor (void)
import Data.Maybe (listToMaybe)
import GHC.Stack (HasCallStack)
import Network.HTTP.Client (Request (..))
import Network.HTTP.Simple
import Network.HTTP.Types.Status (status429)
import Text.Read (readMaybe)

-- | 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
(Int -> RetriesExhausted -> ShowS)
-> (RetriesExhausted -> String)
-> ([RetriesExhausted] -> ShowS)
-> Show RetriesExhausted
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RetriesExhausted -> ShowS
showsPrec :: Int -> RetriesExhausted -> ShowS
$cshow :: RetriesExhausted -> String
show :: RetriesExhausted -> String
$cshowList :: [RetriesExhausted] -> ShowS
showList :: [RetriesExhausted] -> ShowS
Show)

instance Exception RetriesExhausted where
  displayException :: RetriesExhausted -> String
displayException RetriesExhausted {Int
Response ()
reLimit :: RetriesExhausted -> Int
reResponse :: RetriesExhausted -> Response ()
reLimit :: Int
reResponse :: Response ()
..} =
    String
"Retries exhaused after "
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
reLimit
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" attempts. Final response:\n"
      String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Response () -> String
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 = Int
-> (Request -> m (Response body)) -> Request -> m (Response body)
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 <-
    RetryPolicyM m
-> (RetryStatus -> Response body -> m RetryAction)
-> (RetryStatus -> m (Response body))
-> m (Response body)
forall (m :: * -> *) b.
MonadIO m =>
RetryPolicyM m
-> (RetryStatus -> b -> m RetryAction)
-> (RetryStatus -> m b)
-> m b
retryingDynamic
      (Int -> RetryPolicy
limitRetries Int
retryLimit)
      ( \RetryStatus
_ ->
          RetryAction -> m RetryAction
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            (RetryAction -> m RetryAction)
-> (Response body -> RetryAction) -> Response body -> m RetryAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RetryAction -> (Int -> RetryAction) -> Maybe Int -> RetryAction
forall b a. b -> (a -> b) -> Maybe a -> b
maybe RetryAction
DontRetry (Int -> RetryAction
ConsultPolicyOverrideDelay (Int -> RetryAction) -> (Int -> Int) -> Int -> RetryAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
microseconds)
            (Maybe Int -> RetryAction)
-> (Response body -> Maybe Int) -> Response body -> RetryAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Response body -> Maybe Int
forall body. Response body -> Maybe Int
getRetryAfter
      )
      (\RetryStatus
_ -> Request -> m (Response body)
f (Request -> m (Response body)) -> Request -> m (Response body)
forall a b. (a -> b) -> a -> b
$ Request -> Request
suppressRetryStatusError Request
req)

  Int -> Response body -> m (Response body)
forall (m :: * -> *) body.
(MonadIO m, HasCallStack) =>
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
req' Response BodyReader
resp ->
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Response BodyReader -> Status
forall a. Response a -> Status
getResponseStatus Response BodyReader
resp Status -> Status -> Bool
forall a. Eq a => a -> a -> Bool
== Status
status429) (IO () -> IO ()) -> IO () -> IO ()
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, HasCallStack) => Int -> Response body -> m (Response body)
checkRetriesExhausted :: forall (m :: * -> *) body.
(MonadIO m, HasCallStack) =>
Int -> Response body -> m (Response body)
checkRetriesExhausted Int
retryLimit Response body
resp
  | Response body -> Status
forall a. Response a -> Status
getResponseStatus Response body
resp Status -> Status -> Bool
forall a. Eq a => a -> a -> Bool
== Status
status429 =
      RetriesExhausted -> m (Response body)
forall e (m :: * -> *) a.
(MonadIO m, Exception e, HasCallStack) =>
e -> m a
throwWithCallStack (RetriesExhausted -> m (Response body))
-> RetriesExhausted -> m (Response body)
forall a b. (a -> b) -> a -> b
$
        RetriesExhausted {reLimit :: Int
reLimit = Int
retryLimit, reResponse :: Response ()
reResponse = Response body -> Response ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void Response body
resp}
  | Bool
otherwise = Response body -> m (Response body)
forall a. a -> m a
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
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Response body -> Status
forall a. Response a -> Status
getResponseStatus Response body
resp Status -> Status -> Bool
forall a. Eq a => a -> a -> Bool
== Status
status429
  ByteString
header <- [ByteString] -> Maybe ByteString
forall a. [a] -> Maybe a
listToMaybe ([ByteString] -> Maybe ByteString)
-> [ByteString] -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ HeaderName -> Response body -> [ByteString]
forall a. HeaderName -> Response a -> [ByteString]
getResponseHeader HeaderName
"Retry-After" Response body
resp
  String -> Maybe Int
forall a. Read a => String -> Maybe a
readMaybe (String -> Maybe Int) -> String -> Maybe Int
forall a b. (a -> b) -> a -> b
$ ByteString -> String
BS8.unpack ByteString
header

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