content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  --
  -- | Support for building a Haxl datasource to connect to a Thrift
  -- service.
  --
  module Haxl.DataSource.Thrift
    ( dispatchSingle
    , dispatchList, dispatchList'
    , dispatchMap, dispatchMap'
    , dispatchVoid, dispatchVoid'
    , dispatchOneway, dispatchOneway'
    , dispatchCommon
    , sendCobSingle, sendCobOneway, recvCobSingle
    , sendCobBatch, recvCobList, recvCobMap, recvCobVoid
    , putFetchErrors
    , ThriftTransientError(..), ThriftLogicError(..)
    , wrapError
    ) where
  
  import Control.Exception
  import Control.Monad
  import Control.Monad.Trans.Class
  import Control.Monad.Trans.Reader
  import Data.Hashable
  import Data.HashMap.Strict (HashMap)
  import Data.Maybe
  import Data.Proxy
  import Data.Text (Text)
  import Data.Typeable
  import TextShow
  import qualified Data.HashMap.Strict as HashMap
  
  import Haxl.Core.Exception
  import Haxl.Core
  
  import Thrift.Channel
  import Thrift.Monad
  import Thrift.Protocol
  import Thrift.Protocol.ApplicationException.Types
  
  -- Exceptions ------------------------------------------------------------------
  
  newtype ThriftTransientError a =
    ThriftTransientError { unThriftTransientError :: a }
    deriving (Show, Eq)
  instance (Show a, Typeable a) => Exception (ThriftTransientError a) where
    toException = transientErrorToException
    fromException = transientErrorFromException
  
  thriftTransientError :: (Show a, Typeable a) => a -> SomeException
  thriftTransientError = toException . ThriftTransientError
  
  newtype ThriftLogicError a =
    ThriftLogicError { unThriftLogicError :: a }
    deriving (Show, Eq)
  instance (Show a, Typeable a) => Exception (ThriftLogicError a) where
    toException = logicErrorToException
    fromException = logicErrorFromException
  
  thriftLogicError :: (Show a, Typeable a) => a -> SomeException
  thriftLogicError = toException . ThriftLogicError
  
  -- Dispatch Functions ----------------------------------------------------------
  
  -- | Type of a generated thrift send_ function
  type Send p s a =
    forall c . (Protocol p, ClientChannel c)
    => Proxy p
    -> c s
    -> Counter
    -> SendCallback
    -> RecvCallback
    -> RpcOptions
    -> a
    -> IO ()
  
  -- | Type of a generated thrift send_ function
  type SendOneway p s a =
    forall c . (Protocol p, ClientChannel c)
    => Proxy p
    -> c s
    -> Counter
    -> SendCallback
    -> RpcOptions
    -> a
    -> IO ()
  
  -- | Type of a generated thrift recv_ function
  type Recv p a = Proxy p -> Response -> Either SomeException a
  
  -- | Dispatch a list of single (unbatched) requests
  dispatchSingle
    :: (Protocol p, ClientChannel c)
    => [(ResultVar a, req)]
    -> Send p s req
    -> Recv p a
    -> ThriftM p c s ()
  dispatchSingle requests send recv = do
    forM_ requests $ \(rvar, req) ->
      dispatchCommon send
        (sendCobSingle rvar)
        (recvCobSingle rvar . recv)
        req
  
  -- | Dispatch a batched request where the response compes back as a list of
  -- results
  dispatchList
    :: (Protocol p, ClientChannel c)
    => Text -- ^ request name
    -> req
    -> [ResultVar a]
    -> Send p s req
    -> Recv p [a]
    -> ThriftM p c s ()
  dispatchList name request rvars send recv = unless (null rvars) $
    dispatchCommon send
      (sendCobBatch rvars)
      (recvCobList name rvars . recv)
      request
  
  dispatchList'
    :: (Protocol p, ClientChannel c)
    => Text -- ^ request name
    -> HashMap (Maybe RpcOptions) [(a, ResultVar b)]
    -> Send p s [a]
    -> Recv p [b]
    -> ThriftM p c s ()
  dispatchList' name reqs send recv =
    dispatchWithOptions send (\_ rvars -> recvCobList name rvars . recv) reqs
  
  -- | Dispatch a batched request where the response comes back in a hashmap from
  -- key to result
  dispatchMap
    :: (Eq k, Hashable k, TextShow k, Protocol p, ClientChannel c)
    => Text         -- ^ request name
    -> (req -> [k]) -- ^ extract keys from request
    -> req
    -> [ResultVar a]
    -> Send p s req
    -> Recv p (HashMap k a)
    -> ThriftM p c s ()
  dispatchMap name getKeys request rvars send recv = unless (null rvars) $
    dispatchCommon send
      (sendCobBatch rvars)
      (recvCobMap name (getKeys request) rvars . recv)
      request
  
  dispatchMap'
    :: (Eq a, Hashable a, TextShow a, Protocol p, ClientChannel c)
    => Text         -- ^ request name
    -> HashMap (Maybe RpcOptions) [(a, ResultVar b)]
    -> Send p s [a]
    -> Recv p (HashMap a b)
    -> ThriftM p c s ()
  dispatchMap' name reqs send recv =
    dispatchWithOptions send (\req rvars -> recvCobMap name req rvars . recv) reqs
  
  -- | Dispatch a batched request where the response is void
  dispatchVoid
    :: (Protocol p, ClientChannel c)
    => [(a, ResultVar ())]
    -> Send p s [a]
    -> Recv p ()
    -> ThriftM p c s ()
  dispatchVoid reqs send recv = unless (null reqs) $
    dispatchCommon send
      (sendCobBatch rvars)
      (recvCobVoid rvars . recv)
      req
    where
      (req, rvars) = unzip reqs
  
  -- | Dispatch a batched request where the response is void
  dispatchVoid'
    :: (Protocol p, ClientChannel c)
    => HashMap (Maybe RpcOptions) [(a, ResultVar ())]
    -> Send p s [a]
    -> Recv p ()
    -> ThriftM p c s ()
  dispatchVoid' reqs send recv =
    dispatchWithOptions send (\_ rvars -> recvCobVoid rvars . recv) reqs
  
  dispatchOneway
    :: (Protocol p, ClientChannel c)
    => [(a, ResultVar ())]
    -> SendOneway p s [a]
    -> ThriftM p c s ()
  dispatchOneway reqs send = unless (null reqs) $ do
    ThriftEnv{..} <- ask
    lift $ send thriftProxy thriftChannel thriftCounter
      (sendCobBatchOneway rvars)
      thriftRpcOpts
      req
    where
      (req, rvars) = unzip reqs
  
  dispatchOneway'
    :: (Protocol p, ClientChannel c)
    => HashMap (Maybe RpcOptions) [(a, ResultVar ())]
    -> SendOneway p s [a]
    -> ThriftM p c s ()
  dispatchOneway' allReqs send =
    forM_ (HashMap.toList allReqs) $ \(opts, reqs) -> do
      let
        (req, rvars) = unzip reqs
      unless (null reqs) $ do
        ThriftEnv{..} <- ask
        lift $ send thriftProxy thriftChannel thriftCounter
          (sendCobBatchOneway rvars)
          (fromMaybe thriftRpcOpts opts)
          req
  
  dispatchWithOptions
    :: forall p c s a b. (Protocol p, ClientChannel c)
    => Send p s [a]
    -> ([a] -> [ResultVar b] -> Proxy p -> RecvCallback)
    -> HashMap (Maybe RpcOptions) [(a, ResultVar b)]
    -> ThriftM p c s ()
  dispatchWithOptions send recvCob allReqs =
    forM_ (HashMap.toList allReqs) $ \(opts, reqs) -> do
      let
        (req, rvars) = unzip reqs
        send' :: Send p s [a]
        send' = case opts of
          Nothing -> send
          Just o -> \p s c sc rc _ -> send p s c sc rc o
      unless (null req) $
        dispatchCommon send'
          (sendCobBatch rvars)
          (recvCob req rvars)
          req
  
  dispatchCommon
    :: (Protocol p, ClientChannel c)
    => Send p s req
    -> SendCallback
    -> (Proxy p -> RecvCallback)
    -> req
    -> ThriftM p c s ()
  dispatchCommon send sendCob recvCob request = do
    ThriftEnv{..} <- ask
    lift $ send thriftProxy thriftChannel thriftCounter
      sendCob
      (recvCob thriftProxy)
      thriftRpcOpts
      request
  
  -- Generic Callbacks for non-Batched Calls -------------------------------------
  
  sendCobSingle :: ResultVar a -> SendCallback
  sendCobSingle _ Nothing = return ()
  sendCobSingle result (Just err) =
    putResultFromChildThread result $ Left $ thriftTransientError err
  
  sendCobOneway :: ResultVar () -> SendCallback
  sendCobOneway result Nothing = putResultFromChildThread result $ Right ()
  sendCobOneway result (Just err) =
    putResultFromChildThread result $ Left $ thriftTransientError err
  
  recvCobSingle
    :: ResultVar a -> (Response -> Either SomeException a) -> RecvCallback
  recvCobSingle result _ (Left err) =
    putResultFromChildThread result $ Left $ thriftTransientError err
  recvCobSingle result deserialize (Right response) = do
    val <- case deserialize response of
      Left err -> Left <$> wrapError err
      res@Right{} -> return res
    putResultFromChildThread result val
  
  -- Generic Callbacks for Batched Calls -----------------------------------------
  
  sendCobBatch :: [ResultVar a] -> Maybe ChannelException -> IO ()
  sendCobBatch _ Nothing = return ()
  sendCobBatch rvars (Just e) = putChannelExceptions rvars e
  
  sendCobBatchOneway :: [ResultVar ()] -> SendCallback
  sendCobBatchOneway rvars Nothing = mapM_ (`putResult` Right ()) rvars
  sendCobBatchOneway rvars (Just e) = putChannelExceptions rvars e
  
  recvCobList
    :: Text
    -> [ResultVar a]
    -> (Response -> Either SomeException [a])
    -> Either ChannelException Response
    -> IO ()
  recvCobList _ rvars _ (Left e) = putChannelExceptions rvars e
  recvCobList name rvars recv (Right r) = case recv r of
    Left e -> putWrappedErrors rvars e
    Right rs
      | length rs == length rvars ->
          zipWithLastM_ putResult putResultFromChildThread rvars (map Right rs)
      | otherwise ->
          putChannelExceptions rvars (if null rs then err_empty else err_size)
    where
    err_empty = ChannelException $
      name <> ": response was empty"
    err_size = ChannelException $
      name <> ": response length does not match request length"
  
  recvCobMap
    :: (Eq k, Hashable k, TextShow k)
    => Text
    -> [k]
    -> [ResultVar a]
    -> (Response -> Either SomeException (HashMap k a))
    -> Either ChannelException Response
    -> IO ()
  recvCobMap _ _ rvars _ (Left e) = putChannelExceptions rvars e
  recvCobMap name keys rvars recv (Right r) =
    case recv r of
      Left e -> putWrappedErrors rvars e
      Right m -> zipWithLastM_ (f putResult m) (f putResultFromChildThread m) keys rvars
    where
      f put m key rvar =
        put rvar $
          case HashMap.lookup key m of
            Nothing  -> Left $ thriftTransientError $
                        ChannelException $
                        name <> ": no response for key " <> showt key
            Just val -> Right val
  
  recvCobVoid
    :: [ResultVar ()]
    -> (Response -> Either SomeException ())
    -> Either ChannelException Response
    -> IO ()
  recvCobVoid rvars _ (Left e) = putChannelExceptions rvars e
  recvCobVoid rvars recv (Right r) =
    zipWithLastM_ putResult putResultFromChildThread rvars $ repeat $ recv r
  
  putChannelExceptions :: [ResultVar a] -> ChannelException -> IO ()
  putChannelExceptions rvars e =
    mapLastM_ (f putResult) (f putResultFromChildThread) rvars
   where
    f put a = put a (Left (thriftTransientError e))
  
  putWrappedErrors :: [ResultVar a] -> SomeException -> IO ()
  putWrappedErrors rvars e =
    mapLastM_ (f putResult) (f putResultFromChildThread) rvars
   where
    f put a = put a . Left =<< wrapError e
  
  putFetchErrors :: [BlockedFetch req] -> SomeException -> IO ()
  putFetchErrors reqs e =
    mapLastM_ (f putResult) (f putResultFromChildThread) reqs
   where
    f :: (forall a . ResultVar a -> Either SomeException a -> IO ())
      -> BlockedFetch u -> IO ()
    f put (BlockedFetch _ rvar) = put rvar . Left =<< wrapError e
  
  wrapError :: SomeException -> IO SomeException
  wrapError err@(SomeException e)
    | Just ae@SomeAsyncException{} <- fromException err = throwIO ae
    | Just ce@ChannelException{} <- fromException err =
        return $ thriftTransientError ce
    | Just ce@ApplicationException{ applicationException_type = t }
        <- fromException err
    , isTransient t = return $ thriftTransientError ce
    -- This contains ApplicationException and all other exception types defined in
    -- the `throws` clause of the Thrift function
    | otherwise = return $ thriftLogicError e
    where
    isTransient ApplicationExceptionType_Timeout = True
    isTransient ApplicationExceptionType_Loadshedding = True
    isTransient _ = False
  
  
  -- zipWithM_ with a different function to call for the last element
  zipWithLastM_ :: (a -> b -> IO ()) -> (a -> b -> IO ()) -> [a] -> [b] -> IO ()
  zipWithLastM_ f flast (a0:as0) (b0:bs0) = go a0 b0 as0 bs0
   where
    go a b (a':as) (b':bs) = f a b >> go a' b' as bs
    go a b _ _ = flast a b
  zipWithLastM_ _ _ _ _ = return ()
  
  -- mapM_  with a different function to call for the last element
  mapLastM_ :: (a -> IO ()) -> (a -> IO ()) -> [a] -> IO ()
  mapLastM_ _ _ [] = return ()
  mapLastM_ f flast (a0:as0) = go a0 as0
   where
    go a [] = flast a
    go a (a':as) = f a >> go a' as
  