ppad-lmdb-0.2.0: bench/Main.hs

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

module Main where

import Control.DeepSeq (force)
import Control.Exception (evaluate)
import Control.Monad (forM_)
import Criterion.Main hiding (env)
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
                }
  -- build keys and values outside the timed region
  kv1  <- evaluate (force (kvs 1000))
  kv10 <- evaluate (force (kvs 10000))
  ks1  <- evaluate (force (scattered 1000))
  ks10 <- evaluate (force (scattered 10000))
  L.withEnv dir flags $ \env -> do
    dbi <- L.openDbi env Nothing True
    -- pre-populate; the put benchmarks then overwrite these keys
    insert env dbi kv10
    defaultMain
      [ bgroup "put"
          [ bench "1k overwrites (one txn)"  $ whnfIO (insert env dbi kv1)
          , bench "10k overwrites (one txn)" $ whnfIO (insert env dbi kv10)
          ]
      , bgroup "get"
          [ bench "1k random hits"           $ whnfIO (readN env dbi ks1)
          , bench "10k random hits"          $ whnfIO (readN env dbi ks10)
          ]
      , bgroup "cursor"
          [ bench "scan 10k entries"         $ whnfIO (scanAll env dbi)
          ]
      ]
  removeDirectoryRecursive dir

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

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

-- every key in [0, n), in a scattered order: 7919 is prime, so
-- coprime with n, and multiplying by it mod n permutes [0, n)
scattered :: Int -> [BS.ByteString]
scattered n = [ mkKey ((i * 7919) `mod` n) | 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 -> L.Dbi -> [(BS.ByteString, BS.ByteString)] -> IO ()
insert env dbi kv = L.withWriteTxn env $ \txn ->
  forM_ kv $ \(k, v) -> L.put txn dbi k v

readN :: L.Env -> L.Dbi -> [BS.ByteString] -> IO ()
readN env dbi ks = L.withReadTxn env $ \txn ->
  forM_ ks $ \k -> do
    _ <- L.get txn dbi k
    pure ()

scanAll :: L.Env -> L.Dbi -> IO Int
scanAll env dbi = L.withReadTxn env $ \txn ->
  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