A Retry combinator for servant

Retry.hs

{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}

{-| Provides `Retry` combinator for Servant clients

This module exposes mechanics to automatically retry Servant client queries until
some condition on `Response`'s '`statusCode` is met, up to some timeout limit. Its
use is similar to standard `Get`, `Post` and other methods combinators:


To declare a custom policy, combining `DefaultPolicy` (wait for status code @200@ for 180
seconds) and a custom policy to retry request until a @404@ response is retrieved:

@
type MyPolicy = 'Also DefaultPolicy ('ExpectStatus 404 180)
@

To declare an API type using that policy, along with corresponding `Proxy` value and
client-side functions:

@
type MyAPI = "node" :>                          GetRetry DefaultPolicy '[JSON] Value
        :<|> "node" :> ReqBody '[JSON] Value :> PostRetry MyPolicy '[JSON] Value

myAPI :: Proxy MyAPI
myAPI = Proxy

getValue :: RetryM Value
postValue :: Value -> RetryM Value
getValue :<|> postValue = client myAPI
@

To run a client-side handler using the given retry policy:

@
do
  env <- makeEnv
  retryHTTP getValue env
@


-}
module Servant.Retry where

import           GHC.TypeLits
import           Network.HTTP.Types (Status (..))
import           Protolude          hiding (catch)
import           Servant
import           Servant.Client
import           Servant.Common.Req
import           System.IO          (hPutStr)

-- * Retry Policy

-- | Possible retry policies
-- This is used as a kind to define which policy applies to which endpoint
data RetryPolicy = ExpectStatus Nat Nat
                 -- ^Expect some status code before some number of seconds.
                 | Also RetryPolicy RetryPolicy
                 -- ^Combines 2 policies
                 | NoRetry
                 -- ^No retryHttp, run query directly

-- | Default policy
-- Expects a 200 status code within 180 seconds
type DefaultPolicy = 'ExpectStatus 200 180

-- | Possible outcome when checking whether or not a given `RetryPolicy`
-- should continue while running request.
data Retries (p :: k) =
  Cont | Stop | Fail

-- | Combines 2 outcomes
combine :: Retries p
        -> Retries p'
        -> Retries ('Also p p')

combine _    Fail = Fail
combine Fail _    = Fail
combine _    Stop = Stop
combine Stop _    = Stop
combine Cont Cont = Cont

-- | Type-class representing application of a `RetryPolicy` at runtime
class Policy p where

  -- | Outcome of this policy from response of a request
  policy :: Proxy p
         -> ServantError
         -- ^Error generated by request
         -> Int
         -- ^Number of seconds already spent retrying query
         -> Retries p


-- | Retry policy implementation for `ExpectStatus`
--
--   * If query fails with a `FailureResponse` whose status code is
--     not the one expected, it will `Retry`
--
--   * If query fails for any other reason (e.g. connectivity) it will
--     `Retry`
instance (KnownNat st, KnownNat countdown) => Policy ('ExpectStatus st countdown) where

  policy _ FailureResponse{responseStatus} timeout
    | fromIntegral (statusCode responseStatus) == natVal (Proxy :: Proxy st) =
        Stop
    | natVal (Proxy :: Proxy countdown) > fromIntegral timeout =
        Cont
    | otherwise =
        Fail

  policy _ err timeout
    | natVal (Proxy :: Proxy countdown) > fromIntegral timeout =
      Cont
    | otherwise =
      Fail


instance (Policy p1, Policy p2) => Policy ('Also p1 p2) where

  policy _ err timeout =
    let p1' = policy (Proxy :: Proxy p1) err timeout
        p2' = policy (Proxy :: Proxy p2) err timeout
    in combine p1' p2'


-- * Retry Combinator

-- | An abstract type used to declare API endpoints that should be retried
--
-- This type provides full control over the endpoints behaviour, one can use
-- type synonyms `GetRetry` and friends for most uses:
--
--   * `meth`: The HTTP method for the endpoint,
--   * `status`: The status code which represents a successful request
--   * `policy`: The `RetryPolicy` applied on that endpoint
--   * `cts`: A type-level list of content-types to represent the returned value
data Retry
     (meth :: StdMethod)
     (status :: Nat)
     (policy :: k)
     (cts :: [*])
     a

type GetRetry    = Retry 'GET 200
type PostRetry   = Retry 'POST 200
type PutRetry    = Retry 'PUT 200
type DeleteRetry = Retry 'DELETE 200


-- | How to handle `Retry` on the client side
instance (Policy p, ReflectMethod meth, MimeUnrender ct' a) =>
  HasClient (Retry meth status p(ct' : cts') a) where

  type Client (Retry meth status p (ct' : cts') a)
    = RetryM p a

  clientWithRoute _ req =
    CWR (Proxy :: Proxy p) req (clientWithRoute (Proxy :: Proxy (Verb meth status (ct':cts') a)) req)


-- | Wrapper over a standard `ClientM` providing type-level information on which `RetryPolicy`
-- to apply
data RetryM (p :: k) a =
  CWR { proxy :: Proxy p, req :: Req, innerClient :: ClientM a }

-- |Default step (in microseconds) between each iteration of retrying request.
step :: Int

step = 1000 * 1000


-- | The outcome of a retriable query
data RetryOutcome a = Succeeded a
                    -- ^The query was a success and returned a value of type `a`. The interpretation
                    -- of what a success is depends on the underlying type of query
                    | Failed ServantError
                    -- ^The query failed in an expected way, e.g. it was expecting a response's status
                    -- code which is a failure code in HTTP protocole (>= 400).
                    | TimedOut ServantError Int
                    -- ^The query timed-out after given number of seconds retrying. It also provides the
                    -- last failure observed.


-- | Runs a retryable query
--
-- This function runs given `Retryable` query applying some retry policy p`. This function
-- returns a `RetryOutcome a` representing the result of repeated retries, possibly after some
-- delay which is controle by the definition of `p`.
retryHttp :: (Policy p, MonadIO m)
          => RetryM p a
          -> ClientEnv
          -> m (RetryOutcome a)

retryHttp  (CWR prox r req) clientenv =
  go 0
  where

    wait err = liftIO $ do
      hPutStr stderr "."
      threadDelay step

    go exhausted = do
      res <- liftIO $ runClientM req clientenv
      case res of
        Left err ->
          case policy prox err exhausted of
            Stop -> pure $ Failed err
            Fail -> pure $ TimedOut err exhausted
            Cont -> wait err >>
                    go (exhausted + 1)
        Right v -> pure $ Succeeded v

-- | Runs a retryable query, throwing `ServantErr` upon `Failed` or `TimedOut`.
--
-- Works as `retryHttp` but also throws `ServantError` when query "fails" before timeout.
retryHttpThrow ::  (Policy p, MonadIO m)
               => RetryM p a
               -> ClientEnv
               -> m a

retryHttpThrow  cwr clientenv =
  retryHttp cwr clientenv >>= maybeThrow
  where
    maybeThrow (Failed e)     = throwIO e
    maybeThrow (TimedOut e _) = throwIO e
    maybeThrow (Succeeded a)  = pure a