ppad-lmdb-0.1.0: bench/Main.hs

{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Monad (forM_)
import Criterion.Main
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BSL
import qualified Database.LMDB as L
import System.Directory
  ( createDirectory
  , doesDirectoryExist
  , getTemporaryDirectory
  , removeDirectoryRecursive
  )
import System.FilePath ((</>))

main :: IO ()
main = do
  tmpRoot <- getTemporaryDirectory
  let dir = tmpRoot </> "ppad-lmdb-bench"
  exists <- doesDirectoryExist dir
  if exists then removeDirectoryRecursive dir else pure ()
  createDirectory dir
  let flags = L.defaultEnvFlags
                { L.envMapSize = 128 * 1024 * 1024
                , L.envMaxDbs  = 4
                }
  L.withEnv dir flags $ \env -> do
    -- pre-populate
    L.withWriteTxn env $ \txn -> do
      dbi <- L.openDbi txn Nothing True
      forM_ (kvs 10000) $ \(k, v) -> L.put txn dbi k v
    defaultMain
      [ bgroup "put"
          [ bench "1k inserts (one txn)"  $ whnfIO (insert env 1000)
          , bench "10k inserts (one txn)" $ whnfIO (insert env 10000)
          ]
      , bgroup "get"
          [ bench "1k random hits"        $ whnfIO (readN env 1000)
          , bench "10k random hits"       $ whnfIO (readN env 10000)
          ]
      , bgroup "cursor"
          [ bench "scan 10k entries"      $ whnfIO (scanAll env)
          ]
      ]
  removeDirectoryRecursive dir

-- key/value helpers ----------------------------------------------------------

kvs :: Int -> [(BS.ByteString, BS.ByteString)]
kvs n = [ (mkKey i, mkVal i) | i <- [0 .. n - 1] ]

mkKey :: Int -> BS.ByteString
mkKey i = BSL.toStrict (BB.toLazyByteString (BB.string7 "k" <> BB.intDec i))

mkVal :: Int -> BS.ByteString
mkVal i = BSL.toStrict (BB.toLazyByteString (BB.string7 "v" <> BB.intDec i))

-- workloads ------------------------------------------------------------------

insert :: L.Env -> Int -> IO ()
insert env n = L.withWriteTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing False
  forM_ (kvs n) $ \(k, v) -> L.put txn dbi k v

readN :: L.Env -> Int -> IO ()
readN env n = L.withReadTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing False
  forM_ [0 .. n - 1] $ \i -> do
    _ <- L.get txn dbi (mkKey i)
    pure ()

scanAll :: L.Env -> IO Int
scanAll env = L.withReadTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing False
  L.withCursor txn dbi $ \cur -> do
    let go !acc Nothing  = pure acc
        go !acc (Just _) = L.cursorNext cur >>= go (acc + 1)
    mfirst <- L.cursorFirst cur
    go (0 :: Int) mfirst