content: {-# LANGUAGE BlockArguments    #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  {-| To benchmark `nix-serve-ng`, run the following commands:
  
      > $ nix build
      > $ PATH="./result/bin:${PATH}" cabal v1-bench --benchmark-option=--time-mode=wall
  
      You can compare against the old `nix-serve` by changing the first command
      to:
  
      > $ nix build --file '<nixpkgs>' nix-serve
  -}
  module Main where
  
  import Control.Applicative (empty)
  import Data.Foldable (traverse_)
  import Data.Vector (Vector)
  import Numeric.Natural (Natural)
  import Turtle (d)
  
  import qualified Control.Concurrent.Async as Async
  import qualified Data.ByteString.Lazy as ByteString
  import qualified Data.ByteString.Lazy.Char8 as Char8
  import qualified Data.Text as Text
  import qualified Data.Vector as Vector
  import qualified Network.HTTP.Client as HTTP
  import qualified System.IO as IO
  import qualified System.IO.Temp as Temp
  import qualified Test.Tasty.Bench as Bench
  import qualified Turtle
  
  prepareMixed :: Natural -> IO (Vector FilePath)
  prepareMixed maxSize =
      Temp.withSystemTempFile "zeros" \tempFile handle -> do
          IO.hClose handle
  
          let generate i = do
                  let size =
                          truncate
                              (   fromIntegral maxSize
                              ** (fromIntegral i / fromIntegral numFiles)
                              :: Double
                              )
  
                  ByteString.writeFile tempFile (ByteString.replicate size 0)
  
                  text <- Turtle.single do
                      Turtle.inproc "nix-store"
                          [ "--add", Text.pack tempFile ]
                          empty
  
                  return (Text.unpack (Turtle.lineToText text))
  
          Vector.generateM (fromIntegral numFiles) generate
  
  prepareConstant :: Natural -> IO (Vector FilePath)
  prepareConstant size =
      Temp.withSystemTempFile "zeros" \tempFile handle -> do
          IO.hClose handle
  
          ByteString.writeFile tempFile (ByteString.replicate (fromIntegral size) 0)
  
          text <- Turtle.single do
              Turtle.inproc "nix-store"
                  [ "--add", Text.pack tempFile ]
                  empty
  
          let path = Text.unpack (Turtle.lineToText text)
  
          return (Vector.replicate (fromIntegral numFiles) path)
  
  port :: Int
  port = 8000
  
  numFiles :: Natural
  numFiles = 10
  
  runNixServe :: IO ()
  runNixServe =
     Turtle.procs "nix-serve"
         [ "--quiet", "--port", Turtle.format d port ]
         empty
  
  host :: String
  host = "http://localhost:" <> show port
  
  main :: IO ()
  main = do
      manager <- HTTP.newManager HTTP.defaultManagerSettings
  
      Async.withAsync runNixServe \_ -> do
          let fetchNarInfo file = do
                  let hash = take 32 (drop 11 file)
  
                  request <- HTTP.parseRequest (host <> "/" <> hash <> ".narinfo")
  
                  response <- HTTP.httpLbs request manager
  
                  return (HTTP.responseBody response)
  
          let fetchNar file = do
                  bytes <- fetchNarInfo file
  
                  let relativePath =
                          ( Char8.unpack
                          . ByteString.drop 5
                          . (!! 1)
                          . Char8.lines
                          ) bytes
  
                  request <- HTTP.parseRequest (host <> "/" <> relativePath)
  
                  response <- HTTP.httpLbs request manager
  
                  return (HTTP.responseBody response)
  
          Bench.defaultMain
              [ Bench.env (prepareMixed 10000000) \files -> do
                  Bench.bench ("fetch present NAR info ×" <> show numFiles)
                      (Bench.nfAppIO (traverse_ fetchNarInfo) files)
              , Bench.bench "fetch absent NAR info ×1"
                  (Bench.nfAppIO fetchNarInfo "/nix/store/00000000000000000000000000000000")
              , Bench.env (prepareConstant 0) \files -> do
                  Bench.bench ("fetch empty NAR ×" <> show numFiles)
                      (Bench.nfAppIO (traverse_ fetchNar) files)
              , Bench.env (prepareConstant 10000000) \files -> do
                  Bench.bench ("fetch 10 MB NAR ×" <> show numFiles)
                      (Bench.nfAppIO (traverse_ fetchNar) files)
              , Bench.env (prepareMixed 10000000) \files -> do
                  Bench.bench ("fetch NARs of mixed sizes up to 10 MB ×" <> show numFiles)
                      (Bench.nfAppIO (traverse_ fetchNar) files)
              ]
  