{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

{- |
This module provides support for the [wrk](https://github.com/wg/wrk) benchmarking tool.

Given a Servant API and a list of `Endpoint`s, the `export` function can produce a requests file containing
a JSON representation of the provided `Endpoint`s.

In order to provide wrk the request data, you can use a simple lua script as described in [this](http://czerasz.com/2015/07/19/wrk-http-benchmarking-tool-example/) tutorial.

An adapted version of the original script by Michael Czeraszkiewicz can be found in the project's [ repository ](https://github.com/3kyro/servant-benchmark/tree/main/scripts)
-}
module Servant.Benchmark.Tools.Wrk (export) where

import Data.Aeson
import Data.Aeson.Types (Pair)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Data.CaseInsensitive (original)
import Network.HTTP.Types (Header)
import Servant.Benchmark (Endpoint (..))
import Servant.Benchmark.Endpoint (pack)
import Servant.Benchmark.ToText

newtype Output = MkOutput Endpoint

instance ToJSON Output where
    toJSON :: Output -> Value
toJSON (MkOutput Endpoint
endpoint) =
        [Pair] -> Value
object
            [ Text
"path" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Endpoint -> Text
path Endpoint
endpoint
            , Text
"body" Text -> Maybe Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= (ByteString -> Text) -> Maybe ByteString -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Text
forall a. ToText a => a -> Text
toText (Endpoint -> Maybe ByteString
body Endpoint
endpoint)
            , Text
"method" Text -> Maybe Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= (ByteString -> Text) -> Maybe ByteString -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Text
forall a. ToText a => a -> Text
toText (Endpoint -> Maybe ByteString
method Endpoint
endpoint)
            , Text
"headers"
                Text -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= [Pair] -> Value
object
                    (Header -> Pair
headerToValue (Header -> Pair) -> [Header] -> [Pair]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endpoint -> [Header]
headers Endpoint
endpoint)
            ]

    toEncoding :: Output -> Encoding
toEncoding (MkOutput Endpoint
endpoint) =
        Series -> Encoding
pairs (Series -> Encoding) -> Series -> Encoding
forall a b. (a -> b) -> a -> b
$
            Text
"path" Text -> Text -> Series
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Endpoint -> Text
path Endpoint
endpoint
                Series -> Series -> Series
forall a. Semigroup a => a -> a -> a
<> Text
"body" Text -> Maybe Text -> Series
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= (ByteString -> Text) -> Maybe ByteString -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Text
forall a. ToText a => a -> Text
toText (Endpoint -> Maybe ByteString
body Endpoint
endpoint)
                Series -> Series -> Series
forall a. Semigroup a => a -> a -> a
<> Text
"method" Text -> Maybe Text -> Series
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= (ByteString -> Text) -> Maybe ByteString -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Text
forall a. ToText a => a -> Text
toText (Endpoint -> Maybe ByteString
method Endpoint
endpoint)
                Series -> Series -> Series
forall a. Semigroup a => a -> a -> a
<> Text
"headers"
                    Text -> Value -> Series
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= [Pair] -> Value
object
                        (Header -> Pair
headerToValue (Header -> Pair) -> [Header] -> [Pair]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endpoint -> [Header]
headers Endpoint
endpoint)

headerToValue :: Header -> Pair
headerToValue :: Header -> Pair
headerToValue (HeaderName
headerName, ByteString
value) =
    ByteString -> Text
forall a. ToText a => a -> Text
toText (HeaderName -> ByteString
forall s. CI s -> s
original HeaderName
headerName) Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= ByteString -> Text
forall a. ToText a => a -> Text
toText ByteString
value

-- | Export a requests file given a list of `Endpoint`s
export :: FilePath -> [Endpoint] -> IO ()
export :: FilePath -> [Endpoint] -> IO ()
export FilePath
filepath [Endpoint]
endpoints = do
    let encoding :: ByteString
encoding = [Output] -> ByteString
forall a. ToJSON a => a -> ByteString
encode ([Output] -> ByteString) -> [Output] -> ByteString
forall a b. (a -> b) -> a -> b
$ Endpoint -> Output
MkOutput (Endpoint -> Output)
-> (Endpoint -> Endpoint) -> Endpoint -> Output
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Endpoint -> Endpoint
pack (Endpoint -> Output) -> [Endpoint] -> [Output]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Endpoint]
endpoints
    FilePath -> ByteString -> IO ()
BS.writeFile FilePath
filepath (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.toStrict ByteString
encoding