{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
import Control.Concurrent
import Control.Concurrent.Async
import Control.Exception (bracket)
import Control.Monad
import Criterion.Main
import Criterion.Types (Config (..))
import qualified Data.ByteString.Lazy as BL
import Data.Word
import GHC.Generics (Generic)
import Network.GRPC.HighLevel.Server hiding (serverLoop)
import Network.GRPC.HighLevel.Server.Unregistered (serverLoop)
import Network.GRPC.LowLevel
import Network.GRPC.LowLevel.GRPC (threadDelaySecs)
import Proto3.Suite.Class
import Proto3.Suite.Types
import System.Random (randomRIO)
data AddRequest = AddRequest {addX :: Fixed Word32
, addY :: Fixed Word32}
deriving (Show, Eq, Ord, Generic)
instance Message AddRequest
data AddResponse = AddResponse (Fixed Word32)
deriving (Show, Eq, Ord, Generic)
instance Message AddResponse
addMethod :: MethodName
addMethod = MethodName "/unary"
addClientStreamMethod :: MethodName
addClientStreamMethod = MethodName "/clientstream"
addServerStreamMethod :: MethodName
addServerStreamMethod = MethodName "/serverstream"
addBiDiMethod :: MethodName
addBiDiMethod = MethodName "/bidistream"
addHandler :: Handler 'Normal
addHandler =
UnaryHandler addMethod $
\c -> do
let b = payload c
return ( AddResponse $ addX b + addY b
, metadata c
, StatusOk
, StatusDetails ""
)
addClientStreamHandler :: Handler 'ClientStreaming
addClientStreamHandler =
ClientStreamHandler addClientStreamMethod $
\_ recv -> do
answer <- go recv 0
return (Just answer, mempty, StatusOk, "")
where go recv !i = do
req <- recv
case req of
Left _ -> return $ AddResponse i
Right Nothing -> return $ AddResponse i
Right (Just (AddRequest x y)) -> go recv (i+x+y)
addServerStreamHandler :: Handler 'ServerStreaming
addServerStreamHandler =
ServerStreamHandler addServerStreamMethod $
\c send -> do
let AddRequest (Fixed x) y = payload c
replicateM_ (fromIntegral x) $ send $ AddResponse y
return (mempty, StatusOk, "")
addBiDiHandler :: Handler 'BiDiStreaming
addBiDiHandler = BiDiStreamHandler addBiDiMethod (go 0)
where go :: Fixed Word32 -> ServerRWHandler AddRequest AddResponse
go !i c recv send = do
req <- recv
case req of
Left _ -> return (mempty, StatusOk, "")
Right Nothing -> return (mempty, StatusOk, "")
Right (Just (AddRequest x y)) -> do
let curr = i + x + y
void $ send $ AddResponse curr
go curr c recv send
serverOpts :: ServerOptions
serverOpts =
defaultOptions{optNormalHandlers = [addHandler]
, optClientStreamHandlers = [addClientStreamHandler]
, optServerStreamHandlers = [addServerStreamHandler]
, optBiDiStreamHandlers = [addBiDiHandler]}
main :: IO ()
main = bracket startServer stopServer $ const $ withGRPC $ \grpc ->
withClient grpc (ClientConfig "localhost" 50051 [] Nothing Nothing) $ \c -> do
rmAdd <- clientRegisterMethodNormal c addMethod
rmClientStream <- clientRegisterMethodClientStreaming c addClientStreamMethod
rmServerStream <- clientRegisterMethodServerStreaming c addServerStreamMethod
rmBiDiStream <- clientRegisterMethodBiDiStreaming c addBiDiMethod
defaultMainWith
defaultConfig{reportFile = Just "benchmarks.html"}
[ bench "unary request" $ nfIO (addRequest c rmAdd)
, bench "client stream: 100 messages" $ nfIO (addClientStream c rmClientStream 100)
, bench "client stream: 1k messages" $ nfIO (addClientStream c rmClientStream 1000)
, bench "client stream: 10k messages" $ nfIO (addClientStream c rmClientStream 10000)
, bench "server stream: 100 messages" $ nfIO (addServerStream c rmServerStream 100)
, bench "server stream: 1k messages" $ nfIO (addServerStream c rmServerStream 1000)
, bench "server stream: 10k messages" $ nfIO (addServerStream c rmServerStream 10000)
, bench "bidi stream: 50 messages up, 50 down" $ nfIO (bidiStream c rmBiDiStream 50)
, bench "bidi stream: 500 message up, 500 down" $ nfIO (bidiStream c rmBiDiStream 500)
, bench "bidi stream: 5000 messages up, 5000 down" $ nfIO (bidiStream c rmBiDiStream 5000)]
where startServer = do
sThrd <- async $ serverLoop serverOpts
threadDelaySecs 1
return sThrd
stopServer sThrd = cancel sThrd >> void (waitCatch sThrd)
encode = BL.toStrict . toLazyByteString
addRequest c rmAdd = do
x <- liftM Fixed $ randomRIO (0,1000)
y <- liftM Fixed $ randomRIO (0,1000)
let addEnc = BL.toStrict . toLazyByteString $ AddRequest x y
clientRequest c rmAdd 5 addEnc mempty >>= \case
Left e -> fail $ "Got client error on add request: " ++ show e
Right r -> case fromByteString (rspBody r) of
Left e -> fail $ "failed to decode add response: " ++ show e
Right dec
| dec == AddResponse (x + y) -> return ()
| otherwise -> fail $ "Got wrong add answer: " ++ show dec ++ "; expected: " ++ show x ++ " + " ++ show y ++ " = " ++ show (x+y)
addClientStream c rm i = do
let msg = encode $ AddRequest 1 0
Right (Just r,_,_,_,_) <- clientWriter c rm 5 mempty $ \send -> do
replicateM_ i $ send msg
let decoded = fromByteString r
when (decoded /= Right (AddResponse (fromIntegral i))) $
fail $ "clientStream: bad answer: " ++ show decoded ++ "; expected: " ++ show i
addServerStream c rm i = do
let msg = encode $ AddRequest (fromIntegral i) 2
Right (_, _, sd) <- clientReader c rm 5 msg mempty $ \_ recv ->
replicateM_ i $ do
Right (Just bs) <- recv
let Right decoded = fromByteString bs
when (decoded /= AddResponse 2) $
fail $ "serverStream: bad response of " ++ show decoded ++ "; expected 2."
when (sd /= mempty) $ fail $ "bad status details: " ++ show sd
bidiStream c rm i = do
Right (_, _, sd) <- clientRW c rm 5 mempty $ \_ recv send done -> do
forM_ (take i [2,4..]) $ \n -> do
void $ send $ encode $ AddRequest 1 1
Right (Just bs) <- recv
let Right decoded = fromByteString bs
when (decoded /= AddResponse n) $
fail $ "bidiStream: got: " ++ show decoded ++ "expected: " ++ show n
void done
when (sd /= mempty) $ fail $ "bad StatusDetails: " ++ show sd