{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}

{- |
Support for the [Drill](https://github.com/fcsonline/drill) load testing application
-}
module Servant.Benchmark.Tools.Drill (Settings (..), export) where

import Data.Aeson (ToJSON (..), object, (.=))
import Data.Aeson.Types (Pair, Value)
import qualified Data.ByteString as BS
import Data.CaseInsensitive (original)
import Data.Ord (comparing)
import qualified Data.Text as T
import qualified Data.Yaml.Pretty as Y
import Network.HTTP.Types (Header)
import Servant.Benchmark.Endpoint
import Servant.Benchmark.ToText

-- | Drill specific settings. See the project's [ documentation ](https://github.com/fcsonline/drill) for more details
data Settings = MkSettings
    { Settings -> Word
concurrency :: Word
    , Settings -> Text
base :: T.Text
    , Settings -> Word
iterations :: Word
    , Settings -> Word
rampup :: Word
    }

data Output = MkOutput Settings [Endpoint]

instance ToJSON Output where
    toJSON :: Output -> Value
toJSON (MkOutput Settings
settings [Endpoint]
plan) =
        [Pair] -> Value
object
            [ Text
"concurrency" Text -> Word -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Settings -> Word
concurrency Settings
settings
            , Text
"base" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Settings -> Text
base Settings
settings
            , Text
"iterations" Text -> Word -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Settings -> Word
iterations Settings
settings
            , Text
"rampup" Text -> Word -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Settings -> Word
rampup Settings
settings
            , Text
"plan"
                Text -> [Value] -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= [Endpoint] -> [Value]
requests [Endpoint]
plan
            ]

requests :: [Endpoint] -> [Value]
requests :: [Endpoint] -> [Value]
requests [Endpoint]
endpoints = Endpoint -> Value
endpointToJSON (Endpoint -> Value) -> [Endpoint] -> [Value]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Endpoint]
endpoints

endpointToJSON :: Endpoint -> Value
endpointToJSON :: Endpoint -> Value
endpointToJSON Endpoint
endpoint =
    [Pair] -> Value
object
        [ Text
"name" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Endpoint -> Text
name Endpoint
endpoint
        , Text
"request"
            Text -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= [Pair] -> Value
object
                [ Text
"url" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= Endpoint -> Text
path Endpoint
endpoint
                , Text
"method" Text -> Maybe Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= (Method -> Text) -> Maybe Method -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Method -> Text
forall a. ToText a => a -> Text
toText (Endpoint -> Maybe Method
method Endpoint
endpoint)
                , Text
"body" Text -> Maybe Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
.= (Method -> Text) -> Maybe Method -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Method -> Text
forall a. ToText a => a -> Text
toText (Endpoint -> Maybe Method
body 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)
                ]
        ]

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

-- | Export a benchmark file given a list of `Endpoint`s
export :: FilePath -> Settings -> [Endpoint] -> IO ()
export :: FilePath -> Settings -> [Endpoint] -> IO ()
export FilePath
filepath Settings
settings [Endpoint]
endpoints = do
    let output :: Output
output = Settings -> [Endpoint] -> Output
MkOutput Settings
settings ([Endpoint] -> Output) -> [Endpoint] -> Output
forall a b. (a -> b) -> a -> b
$ Endpoint -> Endpoint
pack (Endpoint -> Endpoint) -> [Endpoint] -> [Endpoint]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Endpoint]
endpoints
    let encoding :: Method
encoding = Config -> Output -> Method
forall a. ToJSON a => Config -> a -> Method
Y.encodePretty Config
config Output
output
    FilePath -> Method -> IO ()
BS.writeFile FilePath
filepath Method
encoding

config :: Y.Config
config :: Config
config =
    (Text -> Text -> Ordering) -> Config -> Config
Y.setConfCompare Text -> Text -> Ordering
ordering Config
Y.defConfig

-- Explicit ordering for root Yaml fields
ordering :: T.Text -> T.Text -> Ordering
ordering :: Text -> Text -> Ordering
ordering Text
"plan" Text
_ = Ordering
GT
ordering Text
_ Text
"plan" = Ordering
LT
ordering Text
t1 Text
t2 = (Text -> Int) -> Text -> Text -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing Text -> Int
T.length Text
t1 Text
t2