Haxl/tests/MonadBench.hs

-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file. An additional grant of patent rights can
-- be found in the PATENTS file.

-- | Benchmarking tool for core performance characteristics of the Haxl monad.

module MonadBench (main) where

import Control.Monad
import Data.List as List
import Data.Time.Clock
import System.Environment
import System.Exit
import System.IO
import Text.Printf

import Haxl.Prelude as Haxl
import Prelude()

import Haxl.Core

import ExampleDataSource

testEnv :: IO (Env ())
testEnv = do
  exstate <- ExampleDataSource.initGlobalState
  let st = stateSet exstate stateEmpty
  initEnv st ()

main = do
  [test,n_] <- getArgs
  let n = read n_
  env <- testEnv
  t0 <- getCurrentTime
  case test of
    -- parallel, identical queries
    "par1" -> runHaxl env $
       Haxl.sequence_ (replicate n (listWombats 3))
    -- parallel, distinct queries
    "par2" -> runHaxl env $
       Haxl.sequence_ (map listWombats [1..fromIntegral n])
    -- sequential, identical queries
    "seqr" -> runHaxl env $
       foldr andThen (return ()) (replicate n (listWombats 3))
    -- sequential, left-associated, distinct queries
    "seql" -> runHaxl env $ do
       foldl andThen (return []) (map listWombats [1.. fromIntegral n])
       return ()
    "tree" -> runHaxl env $ void $ tree n
    -- No memoization
    "memo0" -> runHaxl env $
      Haxl.sequence_ [unionWombats | _ <- [1..n]]
    -- One put, N gets.
    "memo1" -> runHaxl env $
      Haxl.sequence_ [memo (42 :: Int) unionWombats | _ <- [1..n]]
    -- N puts, N gets.
    "memo2" -> runHaxl env $
      Haxl.sequence_ [memo (i :: Int) unionWombats | i <- [1..n]]
    _ -> do
      hPutStrLn stderr "syntax: monadbench par1|par2|seqr|seql|memo0|memo1|memo2 NUM"
      exitWith (ExitFailure 1)
  t1 <- getCurrentTime
  printf "%d reqs: %.2fs\n" n (realToFrac (t1 `diffUTCTime` t0) :: Double)
 where
  -- can't use >>, it is aliased to *> and we want the real bind here
  andThen x y = x >>= const y

tree :: Int -> GenHaxl () [Id]
tree 0 = listWombats 0
tree n = concat <$> Haxl.sequence
  [ tree (n-1)
  , listWombats (fromIntegral n), tree (n-1)
  ]

unionWombats :: GenHaxl () [Id]
unionWombats = foldl List.union [] <$> Haxl.mapM listWombats [1..100]