{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main (main) where

import           Bench.Consensus.Mempool
import           Bench.Consensus.Mempool.TestBlock (TestBlock)
import qualified Bench.Consensus.Mempool.TestBlock as TestBlock
import           Control.Arrow (first)
import           Control.DeepSeq
import           Control.Monad (unless)
import qualified Control.Tracer as Tracer
import           Data.Aeson
import qualified Data.ByteString.Lazy as BL
import qualified Data.Csv as Csv
import           Data.Maybe (fromMaybe)
import qualified Data.Text as Text
import qualified Data.Text.Read as Text.Read
import           Main.Utf8 (withStdTerminalHandles)
import           Ouroboros.Consensus.Ledger.SupportsMempool (ByteSize32)
import qualified Ouroboros.Consensus.Mempool.Capacity as Mempool
import           System.Exit (die, exitFailure)
import qualified Test.Consensus.Mempool.Mocked as Mocked
import           Test.Consensus.Mempool.Mocked (MockedMempool)
import           Test.Tasty (withResource)
import           Test.Tasty.Bench (CsvPath (CsvPath), bench, benchIngredients,
                     bgroup, whnfIO)
import           Test.Tasty.HUnit (testCase, (@?=))
import           Test.Tasty.Options (changeOption)
import           Test.Tasty.Runners (parseOptions, tryIngredients)

main :: IO ()
IO ()
main = IO () -> IO ()
forall (m :: * -> *) r. (MonadIO m, MonadMask m) => m r -> m r
withStdTerminalHandles (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    let csvFilePath :: FilePath
csvFilePath = FilePath
"mempool-benchmarks.csv"
    FilePath -> IO ()
runBenchmarks FilePath
csvFilePath
    rawValues <- FilePath -> IO (Vector [Text])
forall {a}. FromRecord a => FilePath -> IO (Vector a)
parseBenchmarkResults FilePath
csvFilePath
    convertCsvRowsToJsonObjects rawValues "mempool-benchmarks.json"
  where
    runBenchmarks :: FilePath -> IO ()
runBenchmarks FilePath
csvFilePath = do
        opts <- [Ingredient] -> TestTree -> IO OptionSet
parseOptions [Ingredient]
benchIngredients TestTree
benchmarkJustAddingTransactions
        let opts' = (Maybe CsvPath -> Maybe CsvPath) -> OptionSet -> OptionSet
forall v. IsOption v => (v -> v) -> OptionSet -> OptionSet
changeOption (CsvPath -> Maybe CsvPath
forall a. a -> Maybe a
Just (CsvPath -> Maybe CsvPath)
-> (Maybe CsvPath -> CsvPath) -> Maybe CsvPath -> Maybe CsvPath
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CsvPath -> Maybe CsvPath -> CsvPath
forall a. a -> Maybe a -> a
fromMaybe (FilePath -> CsvPath
CsvPath FilePath
csvFilePath)) OptionSet
opts
        case tryIngredients benchIngredients opts' benchmarkJustAddingTransactions of
          Maybe (IO Bool)
Nothing               -> IO ()
forall a. IO a
exitFailure
          Just    IO Bool
runIngredient -> do
            success <- IO Bool
runIngredient
            unless success exitFailure
      where
        benchmarkJustAddingTransactions :: TestTree
benchmarkJustAddingTransactions =
            FilePath -> [TestTree] -> TestTree
bgroup FilePath
"Just adding" ([TestTree] -> TestTree) -> [TestTree] -> TestTree
forall a b. (a -> b) -> a -> b
$
                (Int -> TestTree) -> [Int] -> [TestTree]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> TestTree
benchAddNTxs [Int
Item [Int]
10_000, Int
Item [Int]
20_000]
          where
            benchAddNTxs :: Int -> TestTree
benchAddNTxs Int
n =
                IO ([MempoolCmd TestBlock], ByteSize32)
-> (([MempoolCmd TestBlock], ByteSize32) -> IO ())
-> (IO ([MempoolCmd TestBlock], ByteSize32) -> TestTree)
-> TestTree
forall a. IO a -> (a -> IO ()) -> (IO a -> TestTree) -> TestTree
withResource
                  (([MempoolCmd TestBlock], ByteSize32)
-> IO ([MempoolCmd TestBlock], ByteSize32)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([MempoolCmd TestBlock], ByteSize32)
 -> IO ([MempoolCmd TestBlock], ByteSize32))
-> ([MempoolCmd TestBlock], ByteSize32)
-> IO ([MempoolCmd TestBlock], ByteSize32)
forall a b. NFData a => (a -> b) -> a -> b
$!!
                     let cmds :: [MempoolCmd TestBlock]
cmds = Int -> [MempoolCmd TestBlock]
mkNTryAddTxs Int
n
                         sz :: ByteSize32
sz   = (GenTx TestBlock -> ByteSize32) -> [GenTx TestBlock] -> ByteSize32
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap GenTx TestBlock -> ByteSize32
TestBlock.txSize ([GenTx TestBlock] -> ByteSize32)
-> [GenTx TestBlock] -> ByteSize32
forall a b. (a -> b) -> a -> b
$ [MempoolCmd TestBlock] -> [GenTx TestBlock]
forall blk. [MempoolCmd blk] -> [GenTx blk]
getCmdsTxs [MempoolCmd TestBlock]
cmds
                     in ([MempoolCmd TestBlock]
cmds, ByteSize32
sz)
                  )
                  (\([MempoolCmd TestBlock], ByteSize32)
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
                  (\IO ([MempoolCmd TestBlock], ByteSize32)
getCmds -> do
                      FilePath -> [TestTree] -> TestTree
bgroup (Int -> FilePath
forall a. Show a => a -> FilePath
show Int
n FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
" transactions") [
                          FilePath -> Benchmarkable -> TestTree
bench FilePath
"setup mempool" (Benchmarkable -> TestTree) -> Benchmarkable -> TestTree
forall a b. (a -> b) -> a -> b
$ IO (MockedMempool IO TestBlock) -> Benchmarkable
forall a. IO a -> Benchmarkable
whnfIO (IO (MockedMempool IO TestBlock) -> Benchmarkable)
-> IO (MockedMempool IO TestBlock) -> Benchmarkable
forall a b. (a -> b) -> a -> b
$ do
                            (_cmds, capacity) <- IO ([MempoolCmd TestBlock], ByteSize32)
getCmds
                            openMempoolWithCapacity capacity
                        , FilePath -> Benchmarkable -> TestTree
bench FilePath
"setup mempool + benchmark" (Benchmarkable -> TestTree) -> Benchmarkable -> TestTree
forall a b. (a -> b) -> a -> b
$ IO () -> Benchmarkable
forall a. IO a -> Benchmarkable
whnfIO (IO () -> Benchmarkable) -> IO () -> Benchmarkable
forall a b. (a -> b) -> a -> b
$ do
                            (cmds, capacity) <- IO ([MempoolCmd TestBlock], ByteSize32)
getCmds
                            mempool <- openMempoolWithCapacity capacity
                            run mempool cmds
                        , FilePath -> IO () -> TestTree
testCase FilePath
"test" (IO () -> TestTree) -> IO () -> TestTree
forall a b. (a -> b) -> a -> b
$ do
                            (cmds, capacity) <- IO ([MempoolCmd TestBlock], ByteSize32)
getCmds
                            mempool <- openMempoolWithCapacity capacity
                            testAddCmds mempool cmds
                        , FilePath -> IO () -> TestTree
testCase FilePath
"cmds length" (IO () -> TestTree) -> IO () -> TestTree
forall a b. (a -> b) -> a -> b
$ do
                            (cmds, _capacity) <- IO ([MempoolCmd TestBlock], ByteSize32)
getCmds
                            length cmds @?= n
                        ]
                  )
              where
                testAddCmds :: MockedMempool IO blk -> [MempoolCmd blk] -> IO ()
testAddCmds MockedMempool IO blk
mempool [MempoolCmd blk]
cmds = do
                    MockedMempool IO blk -> [MempoolCmd blk] -> IO ()
forall (m :: * -> *) blk.
Monad m =>
MockedMempool m blk -> [MempoolCmd blk] -> m ()
run MockedMempool IO blk
mempool [MempoolCmd blk]
cmds
                    mempoolTxs <- MockedMempool IO blk -> IO [GenTx blk]
forall blk.
LedgerSupportsMempool blk =>
MockedMempool IO blk -> IO [GenTx blk]
Mocked.getTxs MockedMempool IO blk
mempool
                    mempoolTxs @?= getCmdsTxs cmds

    parseBenchmarkResults :: FilePath -> IO (Vector a)
parseBenchmarkResults FilePath
csvFilePath = do
        csvData <- FilePath -> IO ByteString
BL.readFile FilePath
csvFilePath
        case Csv.decode Csv.HasHeader csvData of
          Left FilePath
err   -> FilePath -> IO (Vector a)
forall a. FilePath -> IO a
die FilePath
err
          Right Vector a
rows -> Vector a -> IO (Vector a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector a
rows

    -- Output the mempool benchmark results as a JSON file, which conforms to
    -- the input expected by
    -- https://github.com/benchmark-action/github-action-benchmark
    convertCsvRowsToJsonObjects :: f [Text] -> FilePath -> IO ()
convertCsvRowsToJsonObjects f [Text]
rows FilePath
outFilePath =
      FilePath -> f Value -> IO ()
forall a. ToJSON a => FilePath -> a -> IO ()
encodeFile FilePath
outFilePath (f Value -> IO ()) -> f Value -> IO ()
forall a b. (a -> b) -> a -> b
$ ([Text] -> Value) -> f [Text] -> f Value
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Text] -> Value
convertRowToJsonObject f [Text]
rows
      where
        convertRowToJsonObject :: [Text] -> Value
convertRowToJsonObject (Text
name:Text
mean:[Text]
_) =
          [Pair] -> Value
object [ Key
"name"  Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text -> Text
adjustName Text
name
                 , Key
"value" Key -> Integer -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Integer
adjustedMean
                 , Key
"unit"  Key -> FilePath -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= FilePath
unit
                 ]
          where
            adjustName :: Text -> Text
adjustName = HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
Text.replace Text
"."          Text
" "
                       (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
Text.replace Text
".benchmark" Text
""

            adjustedMean :: Integer
            (Integer
adjustedMean, FilePath
unit) = (Double -> Integer) -> (Double, FilePath) -> (Integer, FilePath)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: MapKind) b c d. Arrow a => a b c -> a (b, d) (c, d)
first Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round
                                 ((Double, FilePath) -> (Integer, FilePath))
-> (Double, FilePath) -> (Integer, FilePath)
forall a b. (a -> b) -> a -> b
$ Double -> (Double, FilePath)
convertPicosecondsWithUnit
                                 (Double -> (Double, FilePath)) -> Double -> (Double, FilePath)
forall a b. (a -> b) -> a -> b
$ Integer -> Double
forall a. Num a => Integer -> a
fromInteger
                                 (Integer -> Double) -> Integer -> Double
forall a b. (a -> b) -> a -> b
$ Text -> Integer
textToInt Text
mean
              where
                textToInt :: Text -> Integer
textToInt = (FilePath -> Integer)
-> ((Integer, Text) -> Integer)
-> Either FilePath (Integer, Text)
-> Integer
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either FilePath -> Integer
forall a. HasCallStack => FilePath -> a
error (Integer, Text) -> Integer
forall a b. (a, b) -> a
fst (Either FilePath (Integer, Text) -> Integer)
-> (Text -> Either FilePath (Integer, Text)) -> Text -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Either FilePath (Integer, Text)
forall a. Integral a => Reader a
Text.Read.decimal

            -- Convert a number of picoseconds to the largest time unit that
            -- makes the conversion greater or equal than one.
            convertPicosecondsWithUnit :: Double -> (Double, String)
            convertPicosecondsWithUnit :: Double -> (Double, FilePath)
convertPicosecondsWithUnit Double
n
                |                        Int
numberOfDigits  Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4  = (Double
n       , FilePath
"picoseconds" )
                | Int
4 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
numberOfDigits  Bool -> Bool -> Bool
&& Int
numberOfDigits Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
7  = (Double
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1e3 , FilePath
"nanoseconds" )
                | Int
7 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
numberOfDigits  Bool -> Bool -> Bool
&& Int
numberOfDigits Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
10 = (Double
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1e6 , FilePath
"microseconds")
                | Int
10 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
numberOfDigits Bool -> Bool -> Bool
&& Int
numberOfDigits Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
13 = (Double
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1e9 , FilePath
"milliseconds")
                | Int
13 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
numberOfDigits                         = (Double
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1e12, FilePath
"seconds"     )
              where
                numberOfDigits :: Int
                numberOfDigits :: Int
numberOfDigits = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
10 Double
n) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            convertPicosecondsWithUnit Double
_ = FilePath -> (Double, FilePath)
forall a. HasCallStack => FilePath -> a
error FilePath
"All the cases should be covered by the conditions above"

        convertRowToJsonObject [Text]
_             = FilePath -> Value
forall a. HasCallStack => FilePath -> a
error FilePath
"Wrong format"

{-------------------------------------------------------------------------------
  Adding TestBlock transactions to a mempool
-------------------------------------------------------------------------------}

openMempoolWithCapacity :: ByteSize32 -> IO (MockedMempool IO TestBlock)
openMempoolWithCapacity :: ByteSize32 -> IO (MockedMempool IO TestBlock)
openMempoolWithCapacity ByteSize32
capacity =
    MempoolCapacityBytesOverride
-> Tracer IO (TraceEventMempool TestBlock)
-> InitialMempoolAndModelParams TestBlock
-> IO (MockedMempool IO TestBlock)
forall blk.
(LedgerSupportsMempool blk, HasTxId (GenTx blk),
 ValidateEnvelope blk) =>
MempoolCapacityBytesOverride
-> Tracer IO (TraceEventMempool blk)
-> InitialMempoolAndModelParams blk
-> IO (MockedMempool IO blk)
Mocked.openMockedMempool (ByteSize32 -> MempoolCapacityBytesOverride
Mempool.mkCapacityBytesOverride ByteSize32
capacity)
                             Tracer IO (TraceEventMempool TestBlock)
forall (m :: * -> *) a. Applicative m => Tracer m a
Tracer.nullTracer
                             Mocked.MempoolAndModelParams {
                                 immpInitialState :: LedgerState TestBlock ValuesMK
Mocked.immpInitialState = LedgerState TestBlock ValuesMK
TestBlock.initialLedgerState
                               , immpLedgerConfig :: LedgerConfig TestBlock
Mocked.immpLedgerConfig = LedgerConfig TestBlock
TestBlock.sampleLedgerConfig
                             }

mkNTryAddTxs :: Int -> [MempoolCmd TestBlock.TestBlock]
mkNTryAddTxs :: Int -> [MempoolCmd TestBlock]
mkNTryAddTxs Int
0 = []
mkNTryAddTxs Int
n =        [GenTx TestBlock -> MempoolCmd TestBlock
forall blk. GenTx blk -> MempoolCmd blk
AddTx ([Token] -> [Token] -> GenTx TestBlock
TestBlock.mkTx [] [Int -> Token
TestBlock.Token Int
0])]
                [MempoolCmd TestBlock]
-> [MempoolCmd TestBlock] -> [MempoolCmd TestBlock]
forall a. Semigroup a => a -> a -> a
<> ((Int, Int) -> MempoolCmd TestBlock)
-> [(Int, Int)] -> [MempoolCmd TestBlock]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (GenTx TestBlock -> MempoolCmd TestBlock
forall blk. GenTx blk -> MempoolCmd blk
AddTx (GenTx TestBlock -> MempoolCmd TestBlock)
-> ((Int, Int) -> GenTx TestBlock)
-> (Int, Int)
-> MempoolCmd TestBlock
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> GenTx TestBlock
forall {a} {a}.
(Integral a, Integral a) =>
(a, a) -> GenTx TestBlock
mkSimpleTx) ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
Item [Int]
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2] [Int
Item [Int]
1 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
  where
    mkSimpleTx :: (a, a) -> GenTx TestBlock
mkSimpleTx (a
x, a
y) = [Token] -> [Token] -> GenTx TestBlock
TestBlock.mkTx [Int -> Token
TestBlock.Token (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x)]
                                       [Int -> Token
TestBlock.Token (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
y)]