natskell-0.0.0.1: bench/Main.hs
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where
import Control.Exception (evaluate)
import Control.Monad (when)
import Data.Bits (xor)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import Data.Int (Int64)
import Data.List (foldl')
import Data.Maybe (fromMaybe)
import Data.Word (Word32, Word64)
import GHC.Clock (getMonotonicTimeNSec)
import GHC.Stats
import Numeric (showFFloat)
import Parser.API
( ParseStep (DropPrefix, Emit, NeedMore, Reject)
, ParsedMessage (ParsedErr, ParsedInfo, ParsedMsg, ParsedOk, ParsedPing, ParsedPong)
, ParserAPI
, parse
)
import Parser.Attoparsec (parserApi)
import System.Environment (getArgs)
import System.Exit (die)
import System.Mem (performGC)
import Text.Printf (printf)
import Types.Err (Err (..))
import Types.Info
( client_id
, connect_urls
, go
, host
, max_payload
, nonce
, port
, proto
, server_id
, version
)
import Types.Msg (Msg (..))
data Config = Config
{ scenarioName :: String
, repetitions :: Int
, scale :: Int
}
data Scenario = Scenario
{ name :: String
, description :: String
, chunks :: [BS.ByteString]
, frameCount :: Int
, inputBytes :: Int
}
data RunSummary = RunSummary
{ parsedFrames :: !Int
, droppedBytes :: !Int
, checksum :: !Int
}
data RuntimeDelta = RuntimeDelta
{ allocatedBytes :: !Word64
, copiedBytes :: !Word64
, peakLiveBytes :: !Word64
, peakMemBytes :: !Word64
, gcCount :: !Word32
, mutatorCpuNs :: !Int64
, gcCpuNs :: !Int64
}
data Aggregate = Aggregate
{ totalFrames :: !Int
, totalDropped :: !Int
, combinedChecksums :: !Int
}
defaultConfig :: Config
defaultConfig =
Config
{ scenarioName = "mixed-chunked"
, repetitions = 5
, scale = 256
}
main :: IO ()
main = do
args <- getArgs
if "--list-scenarios" `elem` args
then listScenarios
else do
config <- parseArgs defaultConfig args
scenario <- resolveScenario config
benchmarkScenario parserApi config scenario
listScenarios :: IO ()
listScenarios =
mapM_ printScenario availableScenarioNames
where
printScenario (scenarioKey, scenarioDescription) =
putStrLn (scenarioKey ++ " - " ++ scenarioDescription)
availableScenarioNames :: [(String, String)]
availableScenarioNames =
[ ("mixed-aligned", "Mixed INFO/PING/PONG/MSG/HMSG/ERR stream with frame-aligned input.")
, ("mixed-chunked", "The same mixed stream split across irregular chunk boundaries.")
, ("mixed-chunked-noisy", "Chunked mixed stream with occasional invalid prefix bytes.")
]
parseArgs :: Config -> [String] -> IO Config
parseArgs config args =
case args of
[] ->
pure config
"--scenario":value:rest ->
parseArgs config { scenarioName = value } rest
"--repetitions":value:rest ->
parsePositiveInt "--repetitions" value >>= \parsed ->
parseArgs config { repetitions = parsed } rest
"--scale":value:rest ->
parsePositiveInt "--scale" value >>= \parsed ->
parseArgs config { scale = parsed } rest
flag:_ | "--" `B8.isPrefixOf` B8.pack flag ->
die ("unknown benchmark flag: " ++ flag)
unexpected:_ ->
die ("unexpected benchmark argument: " ++ unexpected)
parsePositiveInt :: String -> String -> IO Int
parsePositiveInt flag value =
case reads value of
[(parsed, "")]
| parsed > 0 ->
pure parsed
_ ->
die (flag ++ " expects a positive integer, got: " ++ value)
resolveScenario :: Config -> IO Scenario
resolveScenario config =
case scenarioName config of
"mixed-aligned" ->
pure (alignedScenario (scale config))
"mixed-chunked" ->
pure (chunkedScenario (scale config))
"mixed-chunked-noisy" ->
pure (noisyScenario (scale config))
unknown ->
die ("unknown scenario: " ++ unknown)
benchmarkScenario :: ParserAPI ParsedMessage -> Config -> Scenario -> IO ()
benchmarkScenario selectedParser config scenario = do
putStrLn "parser backend: attoparsec"
putStrLn ("scenario: " ++ name scenario)
putStrLn ("description: " ++ description scenario)
putStrLn ("repetitions: " ++ show (repetitions config))
putStrLn ("frames per repetition: " ++ show (frameCount scenario))
putStrLn ("input bytes per repetition: " ++ show (inputBytes scenario))
warmup <- runScenario selectedParser scenario
validateScenario scenario warmup
performGC
beforeStats <- snapshotStats
startNs <- getMonotonicTimeNSec
aggregate <- runRepetitions (repetitions config) (runScenario selectedParser scenario)
endNs <- getMonotonicTimeNSec
performGC
afterStats <- snapshotStats
let elapsedNs = endNs - startNs
elapsedSeconds = fromIntegral elapsedNs / 1e9 :: Double
averageSeconds = elapsedSeconds / fromIntegral (repetitions config)
totalParsedFrames = totalFrames aggregate
totalInputBytes = repetitions config * inputBytes scenario
totalDroppedBytes = totalDropped aggregate
averageRate = fromIntegral totalParsedFrames / elapsedSeconds
averageMiBPerSecond = fromIntegral totalInputBytes / (1024 * 1024) / elapsedSeconds :: Double
finalChecksum = combinedChecksums aggregate
delta = diffStats beforeStats afterStats
allocatedPerFrame = fromIntegral (allocatedBytes delta) / fromIntegral totalParsedFrames :: Double
copiedPerFrame = fromIntegral (copiedBytes delta) / fromIntegral totalParsedFrames :: Double
printf "average delivery rate: %.2f frames/s\n" averageRate
printf "average parser throughput: %.2f MiB/s\n" averageMiBPerSecond
printf "average time per repetition: %.6fs\n" averageSeconds
putStrLn ("dropped invalid prefix bytes: " ++ show totalDroppedBytes)
putStrLn ("checksum: " ++ show finalChecksum)
putStrLn ("allocated bytes: " ++ show (allocatedBytes delta))
printf "allocated bytes per frame: %.2f\n" allocatedPerFrame
putStrLn ("copied bytes: " ++ show (copiedBytes delta))
printf "copied bytes per frame: %.2f\n" copiedPerFrame
putStrLn ("peak live bytes: " ++ show (peakLiveBytes delta))
putStrLn ("peak memory in use bytes: " ++ show (peakMemBytes delta))
putStrLn ("GC count: " ++ show (gcCount delta))
putStrLn ("mutator CPU ms: " ++ formatMillis (mutatorCpuNs delta))
putStrLn ("GC CPU ms: " ++ formatMillis (gcCpuNs delta))
putStrLn "For the full RTS breakdown, rerun with `+RTS -s`."
validateScenario :: Scenario -> RunSummary -> IO ()
validateScenario scenario summary =
when
(parsedFrames summary /= frameCount scenario)
( die
( "benchmark corpus parse mismatch: expected "
++ show (frameCount scenario)
++ " frames but parsed "
++ show (parsedFrames summary)
)
)
runScenario :: ParserAPI ParsedMessage -> Scenario -> IO RunSummary
runScenario parserForScenario scenario =
pure (consumeChunks parserForScenario (chunks scenario))
runRepetitions :: Int -> IO RunSummary -> IO Aggregate
runRepetitions repetitionCount action =
go 0 (Aggregate 0 0 0)
where
go completed aggregate
| completed >= repetitionCount =
pure aggregate
| otherwise = do
summary <- action
let aggregate' =
Aggregate
{ totalFrames = totalFrames aggregate + parsedFrames summary
, totalDropped = totalDropped aggregate + droppedBytes summary
, combinedChecksums = combinedChecksums aggregate * 16777619 + checksum summary
}
evaluate (combinedChecksums aggregate')
go (completed + 1) aggregate'
consumeChunks :: ParserAPI ParsedMessage -> [BS.ByteString] -> RunSummary
consumeChunks parserForScenario =
go BS.empty (RunSummary 0 0 0)
where
go buffer summary remainingChunks =
case remainingChunks of
[] ->
drain buffer summary []
nextChunk:rest ->
drain (buffer <> nextChunk) summary rest
drain buffer summary remainingChunks
| BS.null buffer =
case remainingChunks of
[] ->
summary
_ ->
go buffer summary remainingChunks
| otherwise =
case parse parserForScenario buffer of
Emit message rest
| BS.length rest == BS.length buffer ->
error "parser emitted without consuming input"
| otherwise ->
drain
rest
summary
{ parsedFrames = parsedFrames summary + 1
, checksum = checksum summary `xor` fingerprint message
}
remainingChunks
NeedMore ->
case remainingChunks of
[] ->
error "parser requested more input after the corpus ended"
_ ->
go buffer summary remainingChunks
DropPrefix bytesToDrop _
| bytesToDrop <= 0 ->
error "parser requested a non-positive prefix drop"
| bytesToDrop > BS.length buffer ->
error "parser requested dropping more bytes than available"
| otherwise ->
drain
(BS.drop bytesToDrop buffer)
summary { droppedBytes = droppedBytes summary + bytesToDrop }
remainingChunks
Reject reason ->
error ("parser rejected benchmark corpus: " ++ reason)
fingerprint :: ParsedMessage -> Int
fingerprint parsedMessage =
case parsedMessage of
ParsedPing _ ->
1
ParsedPong _ ->
2
ParsedOk _ ->
3
ParsedErr err ->
11 + BS.length (errReason err)
ParsedInfo info ->
17
+ BS.length (server_id info)
+ BS.length (version info)
+ BS.length (go info)
+ BS.length (host info)
+ port info
+ max_payload info
+ proto info
+ fromMaybe 0 (client_id info)
+ maybe 0 BS.length (nonce info)
+ maybe 0 length (connect_urls info)
ParsedMsg message ->
23
+ BS.length (subject message)
+ BS.length (sid message)
+ maybe 0 BS.length (replyTo message)
+ maybe 0 BS.length (payload message)
+ headersWeight (headers message)
headersWeight :: Maybe [(BS.ByteString, BS.ByteString)] -> Int
headersWeight Nothing = 0
headersWeight (Just pairs) =
foldl' (\total (headerKey, headerValue) -> total + BS.length headerKey + BS.length headerValue) 0 pairs
errReason :: Err -> BS.ByteString
errReason err =
case err of
ErrUnknownOp reason -> reason
ErrRoutePortConn reason -> reason
ErrAuthViolation reason -> reason
ErrAuthTimeout reason -> reason
ErrInvalidProtocol reason -> reason
ErrMaxControlLineEx reason -> reason
ErrErr reason -> reason
ErrTlsRequired reason -> reason
ErrStaleConn reason -> reason
ErrMaxConnsEx reason -> reason
ErrSlowConsumer reason -> reason
ErrMaxPayload reason -> reason
ErrInvalidSubject reason -> reason
ErrPermViolation reason -> reason
snapshotStats :: IO RTSStats
snapshotStats = do
enabled <- getRTSStatsEnabled
if enabled
then getRTSStats
else die "RTS stats are not enabled; the benchmark must be built with -T"
diffStats :: RTSStats -> RTSStats -> RuntimeDelta
diffStats before after =
RuntimeDelta
{ allocatedBytes = allocated_bytes after - allocated_bytes before
, copiedBytes = copied_bytes after - copied_bytes before
, peakLiveBytes = max_live_bytes after
, peakMemBytes = max_mem_in_use_bytes after
, gcCount = gcs after - gcs before
, mutatorCpuNs = mutator_cpu_ns after - mutator_cpu_ns before
, gcCpuNs = gc_cpu_ns after - gc_cpu_ns before
}
formatMillis :: Int64 -> String
formatMillis nanoseconds =
showFFloat (Just 3) (fromIntegral nanoseconds / 1e6 :: Double) "ms"
alignedScenario :: Int -> Scenario
alignedScenario scenarioScale =
Scenario
{ name = "mixed-aligned"
, description = "Mixed control and delivery traffic with each frame presented whole."
, chunks = frames
, frameCount = length frames
, inputBytes = sum (map BS.length frames)
}
where
frames = buildMixedFrames scenarioScale
chunkedScenario :: Int -> Scenario
chunkedScenario scenarioScale =
Scenario
{ name = "mixed-chunked"
, description = "The same mixed traffic split across irregular byte chunks."
, chunks = chunkStream chunkPlan payload
, frameCount = length frames
, inputBytes = BS.length payload
}
where
frames = buildMixedFrames scenarioScale
payload = BS.concat frames
noisyScenario :: Int -> Scenario
noisyScenario scenarioScale =
Scenario
{ name = "mixed-chunked-noisy"
, description = "Chunked mixed traffic with occasional invalid prefix bytes to exercise resynchronization."
, chunks = chunkStream chunkPlan payload
, frameCount = length baseFrames
, inputBytes = BS.length payload
}
where
baseFrames = buildMixedFrames scenarioScale
frames = injectNoise baseFrames
payload = BS.concat frames
chunkPlan :: [Int]
chunkPlan = [17, 31, 64, 127, 257, 511, 1024, 41, 89, 2048, 13]
chunkStream :: [Int] -> BS.ByteString -> [BS.ByteString]
chunkStream plan = go (cycle plan)
where
go _ remaining
| BS.null remaining =
[]
go (nextChunkSize:rest) remaining =
let (chunkBytes, remainingBytes) = BS.splitAt nextChunkSize remaining
in chunkBytes : go rest remainingBytes
go [] _ =
error "chunk plan unexpectedly ended"
injectNoise :: [BS.ByteString] -> [BS.ByteString]
injectNoise =
snd . foldl' inject (0 :: Int, [])
where
inject (index, acc) frame
| index `mod` 29 == 0 =
(index + 1, acc ++ ["Z", frame])
| otherwise =
(index + 1, acc ++ [frame])
buildMixedFrames :: Int -> [BS.ByteString]
buildMixedFrames scenarioScale =
concatMap buildCycle [0 .. scenarioScale - 1]
buildCycle :: Int -> [BS.ByteString]
buildCycle cycleIndex =
[ buildInfoFrame cycleIndex
, "PING\r\n"
, buildMsgFrame cycleIndex "bench.alpha" Nothing 0 Nothing
, buildMsgFrame cycleIndex "bench.beta" (Just "_INBOX.reply") 32 Nothing
, buildMsgFrame cycleIndex "bench.gamma" Nothing 128 Nothing
, buildHMsgFrame cycleIndex "bench.delta" (Just "_INBOX.headers") 256 defaultHeaders
, "PONG\r\n"
, buildMsgFrame cycleIndex "bench.epsilon" Nothing 512 Nothing
, "+OK\r\n"
, buildMsgFrame cycleIndex "bench.zeta" (Just "_INBOX.zeta") 1024 Nothing
, buildHMsgFrame cycleIndex "bench.eta" Nothing 1536 extendedHeaders
, buildMsgFrame cycleIndex "bench.theta" Nothing 2048 Nothing
, "-ERR 'Permissions Violation For Publish To bench.zeta.'\r\n"
, buildMsgFrame cycleIndex "bench.iota" Nothing 64 Nothing
]
defaultHeaders :: [(BS.ByteString, BS.ByteString)]
defaultHeaders =
[ ("Nats-Msg-Id", "msg-123456")
, ("Trace-Id", "trace-a1b2c3d4")
]
extendedHeaders :: [(BS.ByteString, BS.ByteString)]
extendedHeaders =
[ ("Nats-Msg-Id", "msg-654321")
, ("Trace-Id", "trace-z9y8x7w6")
, ("Content-Type", "application/json")
, ("Client", "bench-runner")
]
buildInfoFrame :: Int -> BS.ByteString
buildInfoFrame cycleIndex =
BS.concat
[ "INFO {"
, field "server_id" (quote "bench-server")
, ","
, field "version" (quote "2.10.22")
, ","
, field "go" (quote "1.23.0")
, ","
, field "host" (quote "127.0.0.1")
, ","
, field "port" (decimalBytes 4222)
, ","
, field "max_payload" (decimalBytes 1048576)
, ","
, field "proto" (decimalBytes 1)
, ","
, field "client_id" (decimalBytes (1000 + cycleIndex))
, ","
, field "nonce" (quote ("nonce-" <> decimalBytes cycleIndex))
, ","
, field "headers" "true"
, "}\r\n"
]
field :: BS.ByteString -> BS.ByteString -> BS.ByteString
field key value = BS.concat [quote key, ":", value]
quote :: BS.ByteString -> BS.ByteString
quote bytes = BS.concat ["\"", bytes, "\""]
buildMsgFrame :: Int -> BS.ByteString -> Maybe BS.ByteString -> Int -> Maybe [(BS.ByteString, BS.ByteString)] -> BS.ByteString
buildMsgFrame cycleIndex subjectName replySubject payloadBytes maybeHeaders =
case maybeHeaders of
Nothing ->
BS.concat
[ "MSG "
, subjectName
, " "
, sidBytes cycleIndex
, " "
, maybeReply replySubject
, decimalBytes payloadBytes
, "\r\n"
, payloadOf cycleIndex payloadBytes
, "\r\n"
]
Just headersList ->
buildHMsgFrame cycleIndex subjectName replySubject payloadBytes headersList
buildHMsgFrame :: Int -> BS.ByteString -> Maybe BS.ByteString -> Int -> [(BS.ByteString, BS.ByteString)] -> BS.ByteString
buildHMsgFrame cycleIndex subjectName replySubject payloadBytes headersList =
BS.concat
[ "HMSG "
, subjectName
, " "
, sidBytes cycleIndex
, " "
, maybeReply replySubject
, decimalBytes headerBytes
, " "
, decimalBytes (headerBytes + payloadBytes)
, "\r\n"
, headerBlock
, payloadOf cycleIndex payloadBytes
, "\r\n"
]
where
headerBlock = BS.concat ["NATS/1.0\r\n", renderHeaders headersList, "\r\n"]
headerBytes = BS.length headerBlock
renderHeaders :: [(BS.ByteString, BS.ByteString)] -> BS.ByteString
renderHeaders =
foldMap (\(headerKey, headerValue) -> BS.concat [headerKey, ": ", headerValue, "\r\n"])
maybeReply :: Maybe BS.ByteString -> BS.ByteString
maybeReply Nothing = ""
maybeReply (Just replySubject) = BS.concat [replySubject, " "]
sidBytes :: Int -> BS.ByteString
sidBytes cycleIndex = decimalBytes (100000 + cycleIndex)
payloadOf :: Int -> Int -> BS.ByteString
payloadOf cycleIndex payloadBytes =
BS.take payloadBytes (BS.concat (replicate repeats unit))
where
unit =
BS.concat
[ "payload-"
, decimalBytes cycleIndex
, "-abcdefghijklmnopqrstuvwxyz0123456789"
]
repeats =
max 1 ((payloadBytes `div` BS.length unit) + 1)
decimalBytes :: Int -> BS.ByteString
decimalBytes = B8.pack . show