{-# LANGUAGE TemplateHaskell #-}

module Wire.Rpc
  ( Rpc,
    rpc,
    rpcWithRetries,
    runRpcWithHttp,
    x3,
    zUser,
    expect,
  )
where

import Bilge
import Bilge.RPC hiding (rpc)
import Bilge.Retry
import Control.Monad.Catch
import Control.Retry
import Data.ByteString.Conversion
import Data.Id
import Data.Text.Encoding (encodeUtf8)
import Imports
import Network.HTTP.Client qualified as HTTP
import Network.HTTP.Types
import Polysemy
import Util.Options

-- * Effect

type ServiceName = LText

data Rpc m a where
  Rpc :: ServiceName -> Endpoint -> (Request -> Request) -> Rpc m (Response (Maybe LByteString))
  RpcWithRetries :: ServiceName -> Endpoint -> (Request -> Request) -> Rpc m (Response (Maybe LByteString))

makeSem ''Rpc

runRpcWithHttp :: (Member (Embed IO) r) => Manager -> RequestId -> Sem (Rpc : r) a -> Sem r a
runRpcWithHttp :: forall (r :: EffectRow) a.
Member (Embed IO) r =>
Manager -> RequestId -> Sem (Rpc : r) a -> Sem r a
runRpcWithHttp Manager
mgr RequestId
reqId = (forall (rInitial :: EffectRow) x. Rpc (Sem rInitial) x -> Sem r x)
-> Sem (Rpc : r) a -> Sem r a
forall (e :: Effect) (r :: EffectRow) a.
FirstOrder e "interpret" =>
(forall (rInitial :: EffectRow) x. e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall (rInitial :: EffectRow) x.
  Rpc (Sem rInitial) x -> Sem r x)
 -> Sem (Rpc : r) a -> Sem r a)
-> (forall (rInitial :: EffectRow) x.
    Rpc (Sem rInitial) x -> Sem r x)
-> Sem (Rpc : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
  Rpc ServiceName
serviceName Endpoint
ep Request -> Request
req ->
    IO x -> Sem r x
forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO x -> Sem r x) -> IO x -> Sem r x
forall a b. (a -> b) -> a -> b
$ Manager
-> RequestId
-> HttpRpc (Response (Maybe LByteString))
-> IO (Response (Maybe LByteString))
forall a. Manager -> RequestId -> HttpRpc a -> IO a
runHttpRpc Manager
mgr RequestId
reqId (HttpRpc (Response (Maybe LByteString))
 -> IO (Response (Maybe LByteString)))
-> HttpRpc (Response (Maybe LByteString))
-> IO (Response (Maybe LByteString))
forall a b. (a -> b) -> a -> b
$ ServiceName
-> Endpoint
-> (Request -> Request)
-> HttpRpc (Response (Maybe LByteString))
rpcImpl ServiceName
serviceName Endpoint
ep Request -> Request
req
  RpcWithRetries ServiceName
serviceName Endpoint
ep Request -> Request
req ->
    IO x -> Sem r x
forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO x -> Sem r x) -> IO x -> Sem r x
forall a b. (a -> b) -> a -> b
$ Manager
-> RequestId
-> HttpRpc (Response (Maybe LByteString))
-> IO (Response (Maybe LByteString))
forall a. Manager -> RequestId -> HttpRpc a -> IO a
runHttpRpc Manager
mgr RequestId
reqId (HttpRpc (Response (Maybe LByteString))
 -> IO (Response (Maybe LByteString)))
-> HttpRpc (Response (Maybe LByteString))
-> IO (Response (Maybe LByteString))
forall a b. (a -> b) -> a -> b
$ ServiceName
-> Endpoint
-> (Request -> Request)
-> HttpRpc (Response (Maybe LByteString))
rpcWithRetriesImpl ServiceName
serviceName Endpoint
ep Request -> Request
req

rpcImpl :: ServiceName -> Endpoint -> (Request -> Request) -> HttpRpc (Response (Maybe LByteString))
rpcImpl :: ServiceName
-> Endpoint
-> (Request -> Request)
-> HttpRpc (Response (Maybe LByteString))
rpcImpl ServiceName
serviceName Endpoint
ep Request -> Request
req = do
  ServiceName
-> Request
-> (Request -> Request)
-> HttpRpc (Response (Maybe LByteString))
forall (m :: * -> *).
(MonadUnliftIO m, MonadCatch m, MonadHttp m, HasRequestId m) =>
ServiceName
-> Request
-> (Request -> Request)
-> m (Response (Maybe LByteString))
rpc' ServiceName
serviceName Request
empty ((Request -> Request) -> HttpRpc (Response (Maybe LByteString)))
-> (Request -> Request) -> HttpRpc (Response (Maybe LByteString))
forall a b. (a -> b) -> a -> b
$
    Request -> Request
req
      (Request -> Request) -> (Request -> Request) -> Request -> Request
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Request -> Request
Bilge.host (Text -> ByteString
encodeUtf8 Endpoint
ep.host)
      (Request -> Request) -> (Request -> Request) -> Request -> Request
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Request -> Request
Bilge.port Endpoint
ep.port

rpcWithRetriesImpl :: ServiceName -> Endpoint -> (Request -> Request) -> HttpRpc (Response (Maybe LByteString))
rpcWithRetriesImpl :: ServiceName
-> Endpoint
-> (Request -> Request)
-> HttpRpc (Response (Maybe LByteString))
rpcWithRetriesImpl ServiceName
serviceName Endpoint
ep Request -> Request
req =
  RetryPolicyM HttpRpc
-> [RetryStatus -> Handler HttpRpc Bool]
-> (RetryStatus -> HttpRpc (Response (Maybe LByteString)))
-> HttpRpc (Response (Maybe LByteString))
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
RetryPolicyM m
-> [RetryStatus -> Handler m Bool] -> (RetryStatus -> m a) -> m a
recovering RetryPolicyM HttpRpc
RetryPolicy
x3 [RetryStatus -> Handler HttpRpc Bool]
forall (m :: * -> *) a. Monad m => [a -> Handler m Bool]
rpcHandlers ((RetryStatus -> HttpRpc (Response (Maybe LByteString)))
 -> HttpRpc (Response (Maybe LByteString)))
-> (RetryStatus -> HttpRpc (Response (Maybe LByteString)))
-> HttpRpc (Response (Maybe LByteString))
forall a b. (a -> b) -> a -> b
$
    HttpRpc (Response (Maybe LByteString))
-> RetryStatus -> HttpRpc (Response (Maybe LByteString))
forall a b. a -> b -> a
const (HttpRpc (Response (Maybe LByteString))
 -> RetryStatus -> HttpRpc (Response (Maybe LByteString)))
-> HttpRpc (Response (Maybe LByteString))
-> RetryStatus
-> HttpRpc (Response (Maybe LByteString))
forall a b. (a -> b) -> a -> b
$
      ServiceName
-> Endpoint
-> (Request -> Request)
-> HttpRpc (Response (Maybe LByteString))
rpcImpl ServiceName
serviceName Endpoint
ep Request -> Request
req

-- * Helpers

x3 :: RetryPolicy
x3 :: RetryPolicy
x3 = Int -> RetryPolicy
limitRetries Int
3 RetryPolicyM m -> RetryPolicyM m -> RetryPolicyM m
forall a. Semigroup a => a -> a -> a
<> Int -> RetryPolicyM m
forall (m :: * -> *). Monad m => Int -> RetryPolicyM m
exponentialBackoff Int
100000

zUser :: UserId -> Request -> Request
zUser :: UserId -> Request -> Request
zUser UserId
uid = HeaderName -> ByteString -> Request -> Request
header HeaderName
"Z-User" (UserId -> ByteString
forall a. ToByteString a => a -> ByteString
toByteString' UserId
uid)

expect :: [Status] -> Request -> Request
expect :: [Status] -> Request -> Request
expect [Status]
ss Request
rq = Request
rq {HTTP.checkResponse = check}
  where
    check :: Request -> Response BodyReader -> IO ()
check Request
rq' Response BodyReader
rs = do
      let s :: Status
s = Response BodyReader -> Status
forall body. Response body -> Status
responseStatus Response BodyReader
rs
          rs' :: Response ()
rs' = Response BodyReader
rs {responseBody = ()}
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Status -> Bool
statusIsServerError Status
s Bool -> Bool -> Bool
|| Status
s Status -> [Status] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Status]
ss) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        HttpException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (HttpException -> IO ()) -> HttpException -> IO ()
forall a b. (a -> b) -> a -> b
$
          Request -> HttpExceptionContent -> HttpException
HttpExceptionRequest Request
rq' (Response () -> ByteString -> HttpExceptionContent
HTTP.StatusCodeException Response ()
rs' ByteString
forall a. Monoid a => a
mempty)

-- * Internals

newtype HttpRpc a = HttpRpc {forall a. HttpRpc a -> ReaderT (Manager, RequestId) IO a
unHttpRpc :: ReaderT (Manager, RequestId) IO a}
  deriving newtype
    ( (forall a b. (a -> b) -> HttpRpc a -> HttpRpc b)
-> (forall a b. a -> HttpRpc b -> HttpRpc a) -> Functor HttpRpc
forall a b. a -> HttpRpc b -> HttpRpc a
forall a b. (a -> b) -> HttpRpc a -> HttpRpc b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> HttpRpc a -> HttpRpc b
fmap :: forall a b. (a -> b) -> HttpRpc a -> HttpRpc b
$c<$ :: forall a b. a -> HttpRpc b -> HttpRpc a
<$ :: forall a b. a -> HttpRpc b -> HttpRpc a
Functor,
      Functor HttpRpc
Functor HttpRpc =>
(forall a. a -> HttpRpc a)
-> (forall a b. HttpRpc (a -> b) -> HttpRpc a -> HttpRpc b)
-> (forall a b c.
    (a -> b -> c) -> HttpRpc a -> HttpRpc b -> HttpRpc c)
-> (forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b)
-> (forall a b. HttpRpc a -> HttpRpc b -> HttpRpc a)
-> Applicative HttpRpc
forall a. a -> HttpRpc a
forall a b. HttpRpc a -> HttpRpc b -> HttpRpc a
forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b
forall a b. HttpRpc (a -> b) -> HttpRpc a -> HttpRpc b
forall a b c. (a -> b -> c) -> HttpRpc a -> HttpRpc b -> HttpRpc c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> HttpRpc a
pure :: forall a. a -> HttpRpc a
$c<*> :: forall a b. HttpRpc (a -> b) -> HttpRpc a -> HttpRpc b
<*> :: forall a b. HttpRpc (a -> b) -> HttpRpc a -> HttpRpc b
$cliftA2 :: forall a b c. (a -> b -> c) -> HttpRpc a -> HttpRpc b -> HttpRpc c
liftA2 :: forall a b c. (a -> b -> c) -> HttpRpc a -> HttpRpc b -> HttpRpc c
$c*> :: forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b
*> :: forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b
$c<* :: forall a b. HttpRpc a -> HttpRpc b -> HttpRpc a
<* :: forall a b. HttpRpc a -> HttpRpc b -> HttpRpc a
Applicative,
      Applicative HttpRpc
Applicative HttpRpc =>
(forall a b. HttpRpc a -> (a -> HttpRpc b) -> HttpRpc b)
-> (forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b)
-> (forall a. a -> HttpRpc a)
-> Monad HttpRpc
forall a. a -> HttpRpc a
forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b
forall a b. HttpRpc a -> (a -> HttpRpc b) -> HttpRpc b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. HttpRpc a -> (a -> HttpRpc b) -> HttpRpc b
>>= :: forall a b. HttpRpc a -> (a -> HttpRpc b) -> HttpRpc b
$c>> :: forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b
>> :: forall a b. HttpRpc a -> HttpRpc b -> HttpRpc b
$creturn :: forall a. a -> HttpRpc a
return :: forall a. a -> HttpRpc a
Monad,
      Monad HttpRpc
Monad HttpRpc => (forall a. IO a -> HttpRpc a) -> MonadIO HttpRpc
forall a. IO a -> HttpRpc a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
$cliftIO :: forall a. IO a -> HttpRpc a
liftIO :: forall a. IO a -> HttpRpc a
MonadIO,
      MonadIO HttpRpc
MonadIO HttpRpc =>
(forall b. ((forall a. HttpRpc a -> IO a) -> IO b) -> HttpRpc b)
-> MonadUnliftIO HttpRpc
forall b. ((forall a. HttpRpc a -> IO a) -> IO b) -> HttpRpc b
forall (m :: * -> *).
MonadIO m =>
(forall b. ((forall a. m a -> IO a) -> IO b) -> m b)
-> MonadUnliftIO m
$cwithRunInIO :: forall b. ((forall a. HttpRpc a -> IO a) -> IO b) -> HttpRpc b
withRunInIO :: forall b. ((forall a. HttpRpc a -> IO a) -> IO b) -> HttpRpc b
MonadUnliftIO,
      Monad HttpRpc
Monad HttpRpc =>
(forall e a. (HasCallStack, Exception e) => e -> HttpRpc a)
-> MonadThrow HttpRpc
forall e a. (HasCallStack, Exception e) => e -> HttpRpc a
forall (m :: * -> *).
Monad m =>
(forall e a. (HasCallStack, Exception e) => e -> m a)
-> MonadThrow m
$cthrowM :: forall e a. (HasCallStack, Exception e) => e -> HttpRpc a
throwM :: forall e a. (HasCallStack, Exception e) => e -> HttpRpc a
MonadThrow,
      MonadThrow HttpRpc
MonadThrow HttpRpc =>
(forall e a.
 (HasCallStack, Exception e) =>
 HttpRpc a -> (e -> HttpRpc a) -> HttpRpc a)
-> MonadCatch HttpRpc
forall e a.
(HasCallStack, Exception e) =>
HttpRpc a -> (e -> HttpRpc a) -> HttpRpc a
forall (m :: * -> *).
MonadThrow m =>
(forall e a.
 (HasCallStack, Exception e) =>
 m a -> (e -> m a) -> m a)
-> MonadCatch m
$ccatch :: forall e a.
(HasCallStack, Exception e) =>
HttpRpc a -> (e -> HttpRpc a) -> HttpRpc a
catch :: forall e a.
(HasCallStack, Exception e) =>
HttpRpc a -> (e -> HttpRpc a) -> HttpRpc a
MonadCatch,
      MonadCatch HttpRpc
MonadCatch HttpRpc =>
(forall b.
 HasCallStack =>
 ((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b)
-> (forall b.
    HasCallStack =>
    ((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b)
-> (forall a b c.
    HasCallStack =>
    HttpRpc a
    -> (a -> ExitCase b -> HttpRpc c)
    -> (a -> HttpRpc b)
    -> HttpRpc (b, c))
-> MonadMask HttpRpc
forall b.
HasCallStack =>
((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b
forall a b c.
HasCallStack =>
HttpRpc a
-> (a -> ExitCase b -> HttpRpc c)
-> (a -> HttpRpc b)
-> HttpRpc (b, c)
forall (m :: * -> *).
MonadCatch m =>
(forall b. HasCallStack => ((forall a. m a -> m a) -> m b) -> m b)
-> (forall b.
    HasCallStack =>
    ((forall a. m a -> m a) -> m b) -> m b)
-> (forall a b c.
    HasCallStack =>
    m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c))
-> MonadMask m
$cmask :: forall b.
HasCallStack =>
((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b
mask :: forall b.
HasCallStack =>
((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b
$cuninterruptibleMask :: forall b.
HasCallStack =>
((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b
uninterruptibleMask :: forall b.
HasCallStack =>
((forall a. HttpRpc a -> HttpRpc a) -> HttpRpc b) -> HttpRpc b
$cgeneralBracket :: forall a b c.
HasCallStack =>
HttpRpc a
-> (a -> ExitCase b -> HttpRpc c)
-> (a -> HttpRpc b)
-> HttpRpc (b, c)
generalBracket :: forall a b c.
HasCallStack =>
HttpRpc a
-> (a -> ExitCase b -> HttpRpc c)
-> (a -> HttpRpc b)
-> HttpRpc (b, c)
MonadMask,
      MonadReader (Manager, RequestId)
    )

instance MonadHttp HttpRpc where
  handleRequestWithCont :: Request -> (Response BodyReader -> IO a) -> HttpRpc a
  handleRequestWithCont :: forall a. Request -> (Response BodyReader -> IO a) -> HttpRpc a
handleRequestWithCont Request
req Response BodyReader -> IO a
responseConsumer = do
    Manager
mgr <- ((Manager, RequestId) -> Manager) -> HttpRpc Manager
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Manager, RequestId) -> Manager
forall a b. (a, b) -> a
fst
    Manager -> HttpT HttpRpc a -> HttpRpc a
forall (m :: * -> *) a. Manager -> HttpT m a -> m a
runHttpT Manager
mgr (HttpT HttpRpc a -> HttpRpc a) -> HttpT HttpRpc a -> HttpRpc a
forall a b. (a -> b) -> a -> b
$ Request -> (Response BodyReader -> IO a) -> HttpT HttpRpc a
forall a.
Request -> (Response BodyReader -> IO a) -> HttpT HttpRpc a
forall (m :: * -> *) a.
MonadHttp m =>
Request -> (Response BodyReader -> IO a) -> m a
handleRequestWithCont Request
req Response BodyReader -> IO a
responseConsumer

instance HasRequestId HttpRpc where
  getRequestId :: HttpRpc RequestId
getRequestId = ((Manager, RequestId) -> RequestId) -> HttpRpc RequestId
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Manager, RequestId) -> RequestId
forall a b. (a, b) -> b
snd

runHttpRpc :: Manager -> RequestId -> HttpRpc a -> IO a
runHttpRpc :: forall a. Manager -> RequestId -> HttpRpc a -> IO a
runHttpRpc Manager
mgr RequestId
reqId =
  (ReaderT (Manager, RequestId) IO a -> (Manager, RequestId) -> IO a)
-> (Manager, RequestId)
-> ReaderT (Manager, RequestId) IO a
-> IO a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT (Manager, RequestId) IO a -> (Manager, RequestId) -> IO a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (Manager
mgr, RequestId
reqId) (ReaderT (Manager, RequestId) IO a -> IO a)
-> (HttpRpc a -> ReaderT (Manager, RequestId) IO a)
-> HttpRpc a
-> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HttpRpc a -> ReaderT (Manager, RequestId) IO a
forall a. HttpRpc a -> ReaderT (Manager, RequestId) IO a
unHttpRpc